From 58605718d012e0a9a7bf373dd367f60620ab5802 Mon Sep 17 00:00:00 2001
From: Siyuan Miao <i@xswan.net>
Date: Tue, 15 Apr 2025 01:05:30 +0200
Subject: [PATCH] feat(developer): add pprof endpoint

---
 web.go | 66 +++++++++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 63 insertions(+), 3 deletions(-)

diff --git a/web.go b/web.go
index e74d551..31d3133 100644
--- a/web.go
+++ b/web.go
@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"io/fs"
 	"net/http"
+	"net/http/pprof"
 	"path/filepath"
 	"strings"
 	"time"
@@ -103,6 +104,25 @@ func setupRouter() *gin.Engine {
 	// A Prometheus metrics endpoint.
 	r.GET("/metrics", gin.WrapH(promhttp.Handler()))
 
+	// Developer mode protected routes
+	developerModeRouter := r.Group("/developer/")
+	developerModeRouter.Use(basicAuthProtectedMiddleware(true))
+	{
+		// pprof
+		developerModeRouter.GET("/pprof/", gin.WrapF(pprof.Index))
+		developerModeRouter.GET("/pprof/cmdline", gin.WrapF(pprof.Cmdline))
+		developerModeRouter.GET("/pprof/profile", gin.WrapF(pprof.Profile))
+		developerModeRouter.POST("/pprof/symbol", gin.WrapF(pprof.Symbol))
+		developerModeRouter.GET("/pprof/symbol", gin.WrapF(pprof.Symbol))
+		developerModeRouter.GET("/pprof/trace", gin.WrapF(pprof.Trace))
+		developerModeRouter.GET("/pprof/allocs", gin.WrapH(pprof.Handler("allocs")))
+		developerModeRouter.GET("/pprof/block", gin.WrapH(pprof.Handler("block")))
+		developerModeRouter.GET("/pprof/goroutine", gin.WrapH(pprof.Handler("goroutine")))
+		developerModeRouter.GET("/pprof/heap", gin.WrapH(pprof.Handler("heap")))
+		developerModeRouter.GET("/pprof/mutex", gin.WrapH(pprof.Handler("mutex")))
+		developerModeRouter.GET("/pprof/threadcreate", gin.WrapH(pprof.Handler("threadcreate")))
+	}
+
 	// Protected routes (allows both password and noPassword modes)
 	protected := r.Group("/")
 	protected.Use(protectedMiddleware())
@@ -464,11 +484,51 @@ func protectedMiddleware() gin.HandlerFunc {
 	}
 }
 
+func sendErrorJsonThenAbort(c *gin.Context, status int, message string) {
+	c.JSON(status, gin.H{"error": message})
+	c.Abort()
+}
+
+func basicAuthProtectedMiddleware(requireDeveloperMode bool) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		if requireDeveloperMode {
+			devModeState, err := rpcGetDevModeState()
+			if err != nil {
+				sendErrorJsonThenAbort(c, http.StatusInternalServerError, "Failed to get developer mode state")
+				return
+			}
+
+			if !devModeState.Enabled {
+				sendErrorJsonThenAbort(c, http.StatusUnauthorized, "Developer mode is not enabled")
+				return
+			}
+		}
+
+		if config.LocalAuthMode == "noPassword" {
+			sendErrorJsonThenAbort(c, http.StatusForbidden, "The resource is not available in noPassword mode")
+			return
+		}
+
+		// calculate basic auth credentials
+		_, password, ok := c.Request.BasicAuth()
+		if !ok {
+			c.Header("WWW-Authenticate", "Basic realm=\"JetKVM\"")
+			sendErrorJsonThenAbort(c, http.StatusUnauthorized, "Basic auth is required")
+			return
+		}
+
+		err := bcrypt.CompareHashAndPassword([]byte(config.HashedPassword), []byte(password))
+		if err != nil {
+			sendErrorJsonThenAbort(c, http.StatusUnauthorized, "Invalid password")
+			return
+		}
+
+		c.Next()
+	}
+}
+
 func RunWebServer() {
 	r := setupRouter()
-	//if strings.Contains(builtAppVersion, "-dev") {
-	//	pprof.Register(r)
-	//}
 	err := r.Run(":80")
 	if err != nil {
 		panic(err)