From f4e29f1b9cc0cc7a86022f8c226bf63d7e311c04 Mon Sep 17 00:00:00 2001
From: Siyuan Miao <i@xswan.net>
Date: Thu, 27 Mar 2025 00:49:42 +0100
Subject: [PATCH] feat: make metrics endpoint opt-in

---
 config.go     |  1 +
 prometheus.go | 15 +++++++++++++++
 web.go        |  3 +--
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/config.go b/config.go
index 642f113..a4f2fda 100644
--- a/config.go
+++ b/config.go
@@ -28,6 +28,7 @@ type Config struct {
 	WakeOnLanDevices     []WakeOnLanDevice  `json:"wake_on_lan_devices"`
 	EdidString           string             `json:"hdmi_edid_string"`
 	ActiveExtension      string             `json:"active_extension"`
+	MetricsEnabled       bool               `json:"enable_metrics"`
 	DisplayMaxBrightness int                `json:"display_max_brightness"`
 	DisplayDimAfterSec   int                `json:"display_dim_after_sec"`
 	DisplayOffAfterSec   int                `json:"display_off_after_sec"`
diff --git a/prometheus.go b/prometheus.go
index 8ebf259..16f44fc 100644
--- a/prometheus.go
+++ b/prometheus.go
@@ -3,8 +3,10 @@ package kvm
 import (
 	"net/http"
 
+	"github.com/gin-gonic/gin"
 	"github.com/prometheus/client_golang/prometheus"
 	versioncollector "github.com/prometheus/client_golang/prometheus/collectors/version"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/prometheus/common/version"
 )
 
@@ -14,4 +16,17 @@ func initPrometheus() {
 	// A Prometheus metrics endpoint.
 	version.Version = builtAppVersion
 	prometheus.MustRegister(versioncollector.NewCollector("jetkvm"))
+
+	promHandler = promhttp.Handler()
+}
+
+func prometheusCheckAuthMiddleware() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		if !config.MetricsEnabled {
+			c.JSON(http.StatusNotFound, gin.H{"error": "Metrics endpoint is disabled"})
+			return
+		}
+
+		c.Next()
+	}
 }
diff --git a/web.go b/web.go
index b35a2db..3fe49e2 100644
--- a/web.go
+++ b/web.go
@@ -12,7 +12,6 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
-	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"golang.org/x/crypto/bcrypt"
 )
 
@@ -87,7 +86,7 @@ func setupRouter() *gin.Engine {
 	r.POST("/device/setup", handleSetup)
 
 	// A Prometheus metrics endpoint.
-	r.GET("/metrics", gin.WrapH(promhttp.Handler()))
+	r.GET("/metrics", prometheusCheckAuthMiddleware(), gin.WrapH(promHandler))
 
 	// Protected routes (allows both password and noPassword modes)
 	protected := r.Group("/")