From 33becf1aa6a847c4e43ea608780851949f54284b Mon Sep 17 00:00:00 2001
From: Junwei Zhao <wahyd4@gmail.com>
Date: Wed, 13 Oct 2021 16:29:30 +1100
Subject: [PATCH] Add ability to capture some more info about the api value
 request

---
 handlers/badges.go | 14 ++++++++++++--
 utils/request.go   | 29 +++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)
 create mode 100644 utils/request.go

diff --git a/handlers/badges.go b/handlers/badges.go
index 62602ab..48768e5 100644
--- a/handlers/badges.go
+++ b/handlers/badges.go
@@ -22,6 +22,7 @@ import (
 	"github.com/prometheus/common/model"
 	"github.com/sirupsen/logrus"
 	"github.com/wahyd4/badger/models"
+	"github.com/wahyd4/badger/utils"
 )
 
 var (
@@ -32,6 +33,11 @@ var (
 		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", "ip"})
+
 	badgeGenerationTime = promauto.NewHistogram(prometheus.HistogramOpts{
 		Name: "badge_generation_time",
 		Help: "The time consumed for badge generation",
@@ -181,7 +187,6 @@ func (handler *BadgesHandler) saveBadgeValue(request *APIValueRequest, badge *mo
 }
 
 func (handler *BadgesHandler) generateBadgeValue(c *gin.Context) error {
-	startTime := time.Now()
 	_, internalRequest := c.GetQuery("i")
 
 	queryName := c.Param("name")
@@ -214,8 +219,13 @@ func (handler *BadgesHandler) generateBadgeValue(c *gin.Context) error {
 	}
 
 	if !internalRequest {
-		badgeGenerationTime.Observe(float64(time.Since(startTime).Microseconds()))
 		badgeRequested.With(map[string]string{"badge": queryName}).Inc()
+		aria2UIImageCounter.With(map[string]string{
+			"badge":   queryName,
+			"version": c.Query("version"),
+			"arch":    c.Request.UserAgent(),
+			"ip":      utils.GetUserIP(c.Request),
+		}).Inc()
 	}
 	c.String(http.StatusOK, "%s", metric.TextValue)
 	return nil
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)
+}