package handlers import ( "context" "encoding/json" "errors" "fmt" "net/http" "strconv" "strings" "time" svg "github.com/ajstarks/svgo" "github.com/getsentry/sentry-go" "github.com/gin-gonic/gin" "github.com/jinzhu/gorm" "github.com/jinzhu/gorm/dialects/postgres" "github.com/prometheus/client_golang/api" v1 "github.com/prometheus/client_golang/api/prometheus/v1" "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promauto" "github.com/prometheus/common/model" "github.com/sirupsen/logrus" "github.com/wahyd4/badger/models" "github.com/wahyd4/badger/utils" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" ) var ( ErrNoSuchBadgeFound = errors.New("no such badge can be found, please make sure you have provided the correct information") badgeRequested = promauto.NewCounterVec(prometheus.CounterOpts{ Name: "badge_requested_total", Help: "The total request number of badges", }, []string{"badge"}) aria2UIImageCounter = promauto.NewCounterVec(prometheus.CounterOpts{ Name: "aria2_ui_docker_counter", Help: "The counter for junv's aria2-ui docker image updates checker", }, []string{"badge", "version", "arch"}) badgeGenerationTime = promauto.NewHistogram(prometheus.HistogramOpts{ Name: "badge_generation_time", Help: "The time consumed for badge generation", Buckets: []float64{10, 20, 30, 50, 80, 130, 210, 340, 550}, }) getBadgeTime = promauto.NewHistogramVec(prometheus.HistogramOpts{ Name: "get_badge_time", Help: "The time consumed for badge generation", Buckets: []float64{10, 20, 30, 50, 80, 130, 210, 340, 550}, }, []string{"badge"}) ) const ( ColorSchemeGreen = ColorScheme("fill-opacity:1.00; fill:rgb(49, 197, 83);") ColorSchemeOrange = ColorScheme("fill-opacity:1.00; fill:rgb(255, 87, 34);") ColorSchemeBlue = ColorScheme("fill-opacity:1.00; fill:rgb(32, 99, 155);") BadgeTypePrometheus = "PROMETHEUS" BadgeTypePageView = "PAGE_VIEW" BadgeTypeAPIValue = "API_VALUE" pageViewsLabel = "Page Views" textMarginX = 6.0 textMarginY = 14 badgeHeight = 20 textStyle = "text-anchor:start;font-size:12px;fill:white;font-family:-apple-system,BlinkMacSystemFont,Segoe UI,Helvetica,Arial,sans-serif,Apple Color Emoji,Segoe UI Emoji;" badgeExpireSeconds = 10 metricRequestTimeout = 10 * time.Second ErrorMessageInternalError = "Oops, there's something wrong, please try it again later." ) type Metric struct { Label string Value int TextValue string `json:"text_value"` // for api value Time time.Time } type BadgesHandler struct { DB *gorm.DB PrometheusAPIClient api.Client } type ColorScheme string type APIValueRequest struct { Value string } func (handler *BadgesHandler) ValidateBadge(c *gin.Context) { queryName := c.Param("name") badge, err := handler.queryBadgeFromDB(queryName) if err != nil { if err == gorm.ErrRecordNotFound { c.AbortWithStatusJSON(http.StatusNotFound, GeneralError{"resource not found"}) return } sentry.CaptureException(err) logrus.Error(err) } c.Set("badge", badge) } func (handler *BadgesHandler) GetBadge(c *gin.Context) { span := tracer.StartSpan("web.request", tracer.ResourceName("/GET/svg/badge")) defer span.Finish() startTime := time.Now() if err := handler.generateSVG(c, ColorSchemeGreen); err != nil { logrus.Error(err) sentry.CaptureException(err) if err == ErrNoSuchBadgeFound { c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) return } c.JSON(http.StatusInternalServerError, gin.H{"error": ErrorMessageInternalError}) } getBadgeTime.WithLabelValues(c.Param("name")).Observe(float64(time.Since(startTime).Milliseconds())) } func (handler *BadgesHandler) GetBadgeValue(c *gin.Context) { span := tracer.StartSpan("web.request", tracer.ResourceName("/GET/val/badge")) defer span.Finish() if err := handler.generateBadgeValue(c); err != nil { logrus.Error(err) sentry.CaptureException(err) if err == ErrNoSuchBadgeFound { c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) return } c.JSON(http.StatusInternalServerError, gin.H{"error": ErrorMessageInternalError}) } } func (handler *BadgesHandler) UpdateBadgeValue(c *gin.Context) { request := APIValueRequest{} if err := c.BindJSON(&request); err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request"}) return } if request.Value == "" || strings.TrimSpace(request.Value) == "" { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request"}) return } queryName := c.Param("name") badge, err := handler.queryBadgeFromDB(queryName) if err != nil { if err == gorm.ErrRecordNotFound { c.JSON(http.StatusNotFound, gin.H{"error": "Badge name is invalid"}) return } sentry.CaptureException(err) return } if err := handler.saveBadgeValue(&request, badge); err != nil { logrus.Error(err) sentry.CaptureException(err) c.JSON(http.StatusInternalServerError, gin.H{"error": ErrorMessageInternalError}) return } logrus.Infof("API value for Badge %s updated", badge.ULID) } func (handler *BadgesHandler) GetBadgeWithColorScheme(c *gin.Context) { color := c.Param("color") var colorScheme ColorScheme switch color { case "green.svg": colorScheme = ColorSchemeGreen case "orange.svg": colorScheme = ColorSchemeOrange case "blue.svg": colorScheme = ColorSchemeBlue default: colorScheme = ColorSchemeGreen } if err := handler.generateSVG(c, colorScheme); err != nil { logrus.Error(err) if err == ErrNoSuchBadgeFound { c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) return } sentry.CaptureException(err) c.JSON(http.StatusInternalServerError, gin.H{"error": ErrorMessageInternalError}) } } func (handler *BadgesHandler) saveBadgeValue(request *APIValueRequest, badge *models.Badge) error { return handler.DB.Save(&models.BadgeAPIValue{ BadgeULID: badge.ULID, Value: request.Value, }).Error } func (handler *BadgesHandler) generateBadgeValue(c *gin.Context) error { _, internalRequest := c.GetQuery("i") queryName := c.Param("name") badge, err := handler.queryBadgeFromDB(queryName) if err != nil { if err == gorm.ErrRecordNotFound { return ErrNoSuchBadgeFound } sentry.CaptureException(err) return err } var metric *Metric if badge.Type == BadgeTypePrometheus { prometheusOptions := badge.OptionStruct.(*models.PrometheusOptions) metric, err = handler.queryPrometheus(badge, prometheusOptions) } else if badge.Type == BadgeTypeAPIValue { metric, err = handler.queryBadgeValue(badge, internalRequest) } else { metric, err = handler.queryMetric(badge, internalRequest) } if err != nil { logrus.Errorf("failed to query metric for badge %s : %v", badge.ULID, err) metric = &Metric{ Label: badge.Label, TextValue: "No Value posted to Badger yet", } } if !internalRequest { badgeRequested.With(map[string]string{ "badge": queryName, }).Inc() aria2UIImageCounter.With(map[string]string{ "badge": queryName, "version": c.Query("version"), "arch": c.Query("arch"), }).Inc() } logrus.Infof("Request %s from IP %s version and arch %s", badge.ULID, utils.GetUserIP(c.Request), c.Query("version"), c.Query("arch")) c.String(http.StatusOK, "%s", metric.TextValue) return nil } func (handler *BadgesHandler) generateSVG(c *gin.Context, colorScheme ColorScheme) error { span := tracer.StartSpan("svc.generateSVG") defer span.Finish() startTime := time.Now() _, internalRequest := c.GetQuery("i") queryName := c.Param("name") badge, err := handler.queryBadgeFromDB(queryName) if err != nil { if err == gorm.ErrRecordNotFound { return ErrNoSuchBadgeFound } sentry.CaptureException(err) return err } var metric *Metric if badge.Type == BadgeTypePrometheus { prometheusOptions := badge.OptionStruct.(*models.PrometheusOptions) metric, err = handler.queryPrometheus(badge, prometheusOptions) } else { metric, err = handler.queryMetric(badge, internalRequest) } if err != nil { logrus.Errorf("failed to query metric for badge %s : %v", badge.ULID, err) metric = &Metric{ Label: badge.Label, Value: -1, } } metricValue := metric.Value label := metric.Label valueRectWidth := 40 labelContainerWidth := 110 // Calculate labelContainerWidth based on label length labelContainerWidth = len(label)*6 + 20 valueRectWidth = len(strconv.Itoa(metricValue))*6 + 21 badgeWidth := labelContainerWidth + valueRectWidth w := c.Writer w.Header().Set("Content-Type", "image/svg+xml") w.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d, s-maxage=%d", badgeExpireSeconds, badgeExpireSeconds)) canvas := svg.New(w) canvas.Start(badgeWidth, badgeHeight) canvas.Rect(0, 0, badgeWidth, badgeHeight, "fill-opacity:1.00; fill:rgb(90,90,90); padding: 2px 5px;") canvas.Rect(0, 0, int(labelContainerWidth), badgeHeight, string(colorScheme)) canvas.Text(textMarginX, textMarginY, label, textStyle) valueText := strconv.Itoa(metricValue) canvas.Text(labelContainerWidth+textMarginX, textMarginY, valueText, textStyle) canvas.End() if !internalRequest { badgeGenerationTime.Observe(float64(time.Since(startTime).Milliseconds())) badgeRequested.With(map[string]string{ "badge": queryName, }).Inc() } return nil } func (handler *BadgesHandler) queryBadgeFromDB(queryName string) (*models.Badge, error) { span := tracer.StartSpan("db.queryBadge") defer span.Finish() var badge models.Badge err := handler.DB.Find(&badge, "ulid = ?", queryName).Error if err != nil { if err != gorm.ErrRecordNotFound { sentry.CaptureException(err) return nil, fmt.Errorf("failed to query metric %w", err) } return nil, err } if badge.Type == BadgeTypePrometheus { prometheusOptions := models.PrometheusOptions{} if err = json.Unmarshal(badge.Options.RawMessage, &prometheusOptions); err != nil { sentry.CaptureException(err) return nil, fmt.Errorf("failed to unmarshal badge options %w", err) } badge.OptionStruct = &prometheusOptions } return &badge, nil } func (handler *BadgesHandler) MyBadges(c *gin.Context) { userInterface, exists := c.Get("currentUser") if !exists { logrus.Error("user cannot be found from context") c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"msg": ErrorMessageInternalError}) return } currentUser := userInterface.(models.User) badges := make([]*models.Badge, 0) if err := handler.DB.Find(&badges, "user_id = ?", currentUser.ID).Error; err != nil { sentry.CaptureException(err) c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"msg": ErrorMessageInternalError}) } c.JSON(http.StatusOK, badges) } func (handler *BadgesHandler) CreateBadge(c *gin.Context) { newID := NewULID() jsonB := new(postgres.Jsonb) options := models.PrometheusOptions{ ValueField: "", Query: fmt.Sprintf("ceil(sum(increase(badge_requested_total{badge=\"%s\"}[500d])))", newID), } bytes, err := json.Marshal(&options) if err != nil { logrus.Errorf("failed to marshal option to bytes %v", err) sentry.CaptureException(err) c.JSON(http.StatusInternalServerError, gin.H{"msg": ErrorMessageInternalError}) return } if err = jsonB.Scan(bytes); err != nil { logrus.Error(err.Error()) c.JSON(http.StatusInternalServerError, gin.H{"msg": ErrorMessageInternalError}) return } badge := models.Badge{ Label: pageViewsLabel, ULID: newID, Type: BadgeTypePageView, // Options: jsonB, } if c.GetBool("loggedIn") { badge.UserID = c.GetInt("currentUserID") } if err = handler.DB.Save(&badge).Error; err != nil { logrus.Error(err.Error()) sentry.CaptureException(err) c.JSON(http.StatusInternalServerError, gin.H{"msg": ErrorMessageInternalError}) return } c.JSON(http.StatusOK, gin.H{"id": newID}) } func (handler *BadgesHandler) queryBadgeValue(badge *models.Badge, internalRequest bool) (*Metric, error) { value, err := handler.fetchBadgeValueFromDB(badge) if err != nil { sentry.CaptureException(err) logrus.Error(err.Error()) return nil, err } handler.DB.Save(&models.ViewRecord{ULID: badge.ULID}) return &Metric{ Label: badge.Label, TextValue: value.Value, Time: value.UpdatedAt, }, nil } func (handler *BadgesHandler) queryMetric(badge *models.Badge, internalRequest bool) (*Metric, error) { viewCount, err := handler.fetchViewCountFromDB(badge) if err != nil { return nil, err } // only increase count for non internal request if !internalRequest { if err = handler.updateMetric(viewCount); err != nil { return nil, err } } return &Metric{ Label: badge.Label, Value: viewCount.Count, Time: viewCount.UpdatedAt, }, nil } func (handler *BadgesHandler) fetchBadgeValueFromDB(badge *models.Badge) (*models.BadgeAPIValue, error) { span := tracer.StartSpan("db.fetchBadgeValue") defer span.Finish() var badgeValue models.BadgeAPIValue err := handler.DB.Order("id DESC").First(&badgeValue, "ulid = ?", badge.ULID).Error if err != nil { if err != gorm.ErrRecordNotFound { return nil, err } return nil, fmt.Errorf("no Badge API value found for %s", badge.ULID) } return &badgeValue, nil } func (handler *BadgesHandler) fetchViewCountFromDB(badge *models.Badge) (*models.BadgeViewCount, error) { span := tracer.StartSpan("db.fetchViewCount") defer span.Finish() var viewCount models.BadgeViewCount err := handler.DB.Find(&viewCount, "ulid = ?", badge.ULID).Error if err != nil { if err != gorm.ErrRecordNotFound { return nil, err } logrus.Infof("no db viewcount for %s", badge.ULID) return &models.BadgeViewCount{ BadgeULID: badge.ULID, Count: 0, }, nil } return &viewCount, nil } func (handler *BadgesHandler) updateMetric(viewCount *models.BadgeViewCount) error { viewCount.Count++ err := handler.DB.Save(&viewCount).Error if err != nil { sentry.CaptureException(err) logrus.Errorf("failed to update badge %s viewcount %v", viewCount.BadgeULID, err) return err } handler.DB.Save(&models.ViewRecord{ULID: viewCount.BadgeULID}) return nil } func (handler *BadgesHandler) queryPrometheus(badge *models.Badge, options *models.PrometheusOptions) (*Metric, error) { v1api := v1.NewAPI(handler.PrometheusAPIClient) ctx, cancel := context.WithTimeout(context.Background(), metricRequestTimeout) defer cancel() result, warnings, err := v1api.Query(ctx, options.Query, time.Now()) if err != nil { sentry.CaptureException(err) return nil, fmt.Errorf("error querying Prometheus: %w", err) } if len(warnings) > 0 { logrus.Warnf("Warnings: %v\n", warnings) } var metric Metric for _, v := range result.(model.Vector) { metric = Metric{ Label: string(v.Metric[model.LabelName(options.ValueField)]), Value: int(float64(v.Value)), Time: v.Timestamp.Time(), } } // handle label is empty, then use badge label from database if len(metric.Label) == 0 { metric.Label = badge.Label } handler.DB.Save(&models.ViewRecord{ULID: badge.ULID}) return &metric, nil }