diff --git a/handlers/badges.go b/handlers/badges.go index b9e141e..1b09c9a 100644 --- a/handlers/badges.go +++ b/handlers/badges.go @@ -7,6 +7,7 @@ import ( "fmt" "net/http" "strconv" + "strings" "time" svg "github.com/ajstarks/svgo" @@ -21,6 +22,7 @@ import ( "github.com/prometheus/common/model" "github.com/sirupsen/logrus" "github.com/wahyd4/badger/models" + "github.com/wahyd4/badger/utils" ) var ( @@ -29,7 +31,12 @@ var ( badgeRequested = promauto.NewCounterVec(prometheus.CounterOpts{ Name: "badge_requested_total", Help: "The total request number of badges", - }, []string{"badge"}) + }, []string{"badge", "useragent", "ip"}) + + 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", "useragent", "ip"}) badgeGenerationTime = promauto.NewHistogram(prometheus.HistogramOpts{ Name: "badge_generation_time", @@ -43,6 +50,7 @@ const ( 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 @@ -54,9 +62,10 @@ const ( ) type Metric struct { - Label string - Value int - Time time.Time + Label string + Value int + TextValue string `json:"text_value"` // for api value + Time time.Time } type BadgesHandler struct { @@ -66,6 +75,10 @@ type BadgesHandler struct { type ColorScheme string +type APIValueRequest struct { + Value string +} + func (handler *BadgesHandler) ValidateBadge(c *gin.Context) { queryName := c.Param("name") @@ -94,6 +107,52 @@ func (handler *BadgesHandler) GetBadge(c *gin.Context) { } } +func (handler *BadgesHandler) GetBadgeValue(c *gin.Context) { + 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 @@ -120,6 +179,62 @@ func (handler *BadgesHandler) GetBadgeWithColorScheme(c *gin.Context) { } } +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, + "useragent": c.Request.UserAgent(), + "ip": utils.GetUserIP(c.Request), + }).Inc() + aria2UIImageCounter.With(map[string]string{ + "badge": queryName, + "version": c.Query("version"), + "useragent": c.Request.UserAgent(), + "ip": utils.GetUserIP(c.Request), + }).Inc() + } + c.String(http.StatusOK, "%s", metric.TextValue) + return nil +} + func (handler *BadgesHandler) generateSVG(c *gin.Context, colorScheme ColorScheme) error { startTime := time.Now() _, internalRequest := c.GetQuery("i") @@ -185,7 +300,11 @@ func (handler *BadgesHandler) generateSVG(c *gin.Context, colorScheme ColorSchem if !internalRequest { badgeGenerationTime.Observe(float64(time.Since(startTime).Microseconds())) - badgeRequested.With(map[string]string{"badge": queryName}).Inc() + badgeRequested.With(map[string]string{ + "badge": queryName, + "useragent": c.Request.UserAgent(), + "ip": utils.GetUserIP(c.Request), + }).Inc() } return nil @@ -272,6 +391,21 @@ func (handler *BadgesHandler) CreateBadge(c *gin.Context) { 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 { @@ -292,6 +426,18 @@ func (handler *BadgesHandler) queryMetric(badge *models.Badge, internalRequest b }, nil } +func (handler *BadgesHandler) fetchBadgeValueFromDB(badge *models.Badge) (*models.BadgeAPIValue, error) { + 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) { var viewCount models.BadgeViewCount err := handler.DB.Find(&viewCount, "ulid = ?", badge.ULID).Error diff --git a/main.go b/main.go index ad8b94d..1d22817 100644 --- a/main.go +++ b/main.go @@ -61,7 +61,10 @@ func main() { &models.User{}, &models.ViewRecord{}, &models.FeatureToggle{}, - &models.BadgeViewCount{}) + &models.BadgeViewCount{}, + &models.BadgeAPIValue{}, + &models.UserAPIToken{}, + ) client, err := api.NewClient(api.Config{ Address: "https://prometheus-api.junv.me", @@ -89,6 +92,8 @@ func main() { }) r.GET("/svg/:name", handler.GetBadge) + r.GET("/val/:name", handler.GetBadgeValue) + r.PUT("/val/:name", handler.UpdateBadgeValue) r.GET("/badges/:name/:color", handler.GetBadgeWithColorScheme) diff --git a/models/api_value.go b/models/api_value.go new file mode 100644 index 0000000..7cfed4f --- /dev/null +++ b/models/api_value.go @@ -0,0 +1,9 @@ +package models + +import "github.com/jinzhu/gorm" + +type BadgeAPIValue struct { + gorm.Model + BadgeULID string `gorm:"column:ulid;type:varchar;not null" json:"ulid"` + Value string `gorm:"not null;" json:"value"` +} diff --git a/models/user_api_token.go b/models/user_api_token.go new file mode 100644 index 0000000..2969f0e --- /dev/null +++ b/models/user_api_token.go @@ -0,0 +1,10 @@ +package models + +import "github.com/jinzhu/gorm" + +type UserAPIToken struct { + gorm.Model + Token string + ExpiresAt string + UserID int +} diff --git a/utils/request.go b/utils/request.go new file mode 100644 index 0000000..66070d3 --- /dev/null +++ b/utils/request.go @@ -0,0 +1,29 @@ +package utils + +import ( + "fmt" + "net/http" + "net/http/httputil" +) + +func GetUserIP(r *http.Request) string { + IPAddress := r.Header.Get("X-Real-Ip") + if IPAddress == "" { + IPAddress = r.Header.Get("X-Forwarded-For") + } + if IPAddress == "" { + IPAddress = r.RemoteAddr + } + return IPAddress +} + +func DumpRequest(r *http.Request) { + if r == nil { + return + } + dump, err := httputil.DumpRequest(r, true) + if err != nil { + fmt.Println(err.Error()) + } + fmt.Println(dump) +}