From ee4dfca47691f740eb50e066c719612e7b18b122 Mon Sep 17 00:00:00 2001
From: Martin Angers <martin.n.angers@gmail.com>
Date: Mon, 2 May 2022 08:58:03 -0400
Subject: [PATCH] Use a debug mux to support both fleet-authenticated and
 token-auth debug paths (#5424)

---
 changes/issue-5268-fix-fleet-serve-panic |  1 +
 cmd/fleet/serve.go                       | 27 ++++++++++++-
 cmd/fleet/serve_test.go                  | 50 ++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 changes/issue-5268-fix-fleet-serve-panic

diff --git a/changes/issue-5268-fix-fleet-serve-panic b/changes/issue-5268-fix-fleet-serve-panic
new file mode 100644
index 0000000000..d487f2992f
--- /dev/null
+++ b/changes/issue-5268-fix-fleet-serve-panic
@@ -0,0 +1 @@
+* Fix a panic when running `fleet serve --debug` due to duplicate registration of `/debug/` paths.
diff --git a/cmd/fleet/serve.go b/cmd/fleet/serve.go
index c4436f696b..a5bab6f58f 100644
--- a/cmd/fleet/serve.go
+++ b/cmd/fleet/serve.go
@@ -451,7 +451,11 @@ the way that the Fleet server works.
 
 			rootMux.Handle("/api/", apiHandler)
 			rootMux.Handle("/", frontendHandler)
-			rootMux.Handle("/debug/", service.MakeDebugHandler(svc, config, logger, eh, ds))
+
+			debugHandler := &debugMux{
+				fleetAuthenticatedHandler: service.MakeDebugHandler(svc, config, logger, eh, ds),
+			}
+			rootMux.Handle("/debug/", debugHandler)
 
 			if path, ok := os.LookupEnv("FLEET_TEST_PAGE_PATH"); ok {
 				// test that we can load this
@@ -477,7 +481,7 @@ the way that the Fleet server works.
 				if err != nil {
 					initFatal(err, "generating debug token")
 				}
-				rootMux.Handle("/debug/", http.StripPrefix("/debug/", netbug.AuthHandler(debugToken)))
+				debugHandler.tokenAuthenticatedHandler = http.StripPrefix("/debug/", netbug.AuthHandler(debugToken))
 				fmt.Printf("*** Debug mode enabled ***\nAccess the debug endpoints at /debug/?token=%s\n", url.QueryEscape(debugToken))
 			}
 
@@ -760,3 +764,22 @@ func argsToString(args []driver.NamedValue) string {
 	allArgs.WriteString("}")
 	return allArgs.String()
 }
+
+// The debugMux directs the request to either the fleet-authenticated handler,
+// which is the standard handler for debug endpoints (using a Fleet
+// authorization bearer token), or to the token-authenticated handler if a
+// query-string token is provided and such a handler is set. The only wayt to
+// set this handler is if the --debug flag was provided to the fleet serve
+// command.
+type debugMux struct {
+	fleetAuthenticatedHandler http.Handler
+	tokenAuthenticatedHandler http.Handler
+}
+
+func (m *debugMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.URL.Query().Has("token") && m.tokenAuthenticatedHandler != nil {
+		m.tokenAuthenticatedHandler.ServeHTTP(w, r)
+		return
+	}
+	m.fleetAuthenticatedHandler.ServeHTTP(w, r)
+}
diff --git a/cmd/fleet/serve_test.go b/cmd/fleet/serve_test.go
index a5f161ea5f..0e9efbc499 100644
--- a/cmd/fleet/serve_test.go
+++ b/cmd/fleet/serve_test.go
@@ -505,3 +505,53 @@ func TestBasicAuthHandler(t *testing.T) {
 		})
 	}
 }
+
+func TestDebugMux(t *testing.T) {
+	h1 := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(200) })
+	h2 := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(400) })
+
+	cases := []struct {
+		desc string
+		mux  debugMux
+		tok  string
+		want int
+	}{
+		{
+			"only fleet auth handler, no token",
+			debugMux{fleetAuthenticatedHandler: h1},
+			"",
+			200,
+		},
+		{
+			"only fleet auth handler, with token",
+			debugMux{fleetAuthenticatedHandler: h1},
+			"token",
+			200,
+		},
+		{
+			"both handlers, no token",
+			debugMux{fleetAuthenticatedHandler: h1, tokenAuthenticatedHandler: h2},
+			"",
+			200,
+		},
+		{
+			"both handlers, with token",
+			debugMux{fleetAuthenticatedHandler: h1, tokenAuthenticatedHandler: h2},
+			"token",
+			400,
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.desc, func(t *testing.T) {
+			path := "/debug/pprof"
+			if c.tok != "" {
+				path += "?token=" + c.tok
+			}
+			req := httptest.NewRequest("GET", path, nil)
+			res := httptest.NewRecorder()
+			c.mux.ServeHTTP(res, req)
+			require.Equal(t, c.want, res.Code)
+		})
+	}
+}