badger/handlers/badges.go

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
}