From e71ac1f04221fa8a219d42d2b46780015bf980f4 Mon Sep 17 00:00:00 2001
From: realityone <realityone@me.com>
Date: Mon, 18 Mar 2024 18:59:17 +0800
Subject: [PATCH] protect internal handler

---
 proxy/debug/debug.go |  3 ++-
 router/mux/mux.go    | 12 +++++++++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/proxy/debug/debug.go b/proxy/debug/debug.go
index d675deb..a4dabdb 100644
--- a/proxy/debug/debug.go
+++ b/proxy/debug/debug.go
@@ -6,6 +6,7 @@ import (
 	"path"
 	"strings"
 
+	rmux "github.com/go-kratos/gateway/router/mux"
 	"github.com/go-kratos/kratos/v2/log"
 	"github.com/gorilla/mux"
 )
@@ -39,7 +40,7 @@ func Register(name string, debuggable Debuggable) {
 func MashupWithDebugHandler(origin http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		if strings.HasPrefix(req.URL.Path, _debugPrefix) {
-			globalService.ServeHTTP(w, req)
+			rmux.ProtectedHandler(globalService).ServeHTTP(w, req)
 			return
 		}
 		origin.ServeHTTP(w, req)
diff --git a/router/mux/mux.go b/router/mux/mux.go
index 6417f49..716a0a0 100644
--- a/router/mux/mux.go
+++ b/router/mux/mux.go
@@ -37,13 +37,23 @@ type muxRouter struct {
 	allCloser []io.Closer
 }
 
+func ProtectedHandler(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("X-Forwarded-For") != "" {
+			http.Error(w, http.StatusText(http.StatusForbidden), http.StatusForbidden)
+			return
+		}
+		h.ServeHTTP(w, r)
+	})
+}
+
 // NewRouter new a mux router.
 func NewRouter(notFoundHandler, methodNotAllowedHandler http.Handler) router.Router {
 	r := &muxRouter{
 		Router: mux.NewRouter().StrictSlash(EnableStrictSlash),
 		wg:     &sync.WaitGroup{},
 	}
-	r.Router.Handle("/metrics", promhttp.Handler())
+	r.Router.Handle("/metrics", ProtectedHandler(promhttp.Handler()))
 	r.Router.NotFoundHandler = notFoundHandler
 	r.Router.MethodNotAllowedHandler = methodNotAllowedHandler
 	return r