From 5ff820cfd03877c35184d19607bb5bec532ac8fc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan-Otto=20Kr=C3=B6pke?= <github@jkroepke.de>
Date: Fri, 3 May 2024 07:30:40 +0200
Subject: [PATCH] Ability to disable debug endpoints (#1448)

---
 exporter.go | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/exporter.go b/exporter.go
index 5ff8c8c2..2a0aecbb 100644
--- a/exporter.go
+++ b/exporter.go
@@ -11,7 +11,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
-	_ "net/http/pprof"
+	"net/http/pprof"
 	"os"
 	"os/user"
 	"runtime"
@@ -80,6 +80,10 @@ func main() {
 			"scrape.timeout-margin",
 			"Seconds to subtract from the timeout allowed by the client. Tune to allow for overhead or high loads.",
 		).Default("0.5").Float64()
+		debugEnabled = app.Flag(
+			"debug.enabled",
+			"If true, windows_exporter will expose debug endpoints under /debug/pprof.",
+		).Default("false").Bool()
 	)
 
 	winlogConfig := &winlog.Config{}
@@ -174,15 +178,16 @@ func main() {
 
 	_ = level.Info(logger).Log("msg", fmt.Sprintf("Enabled collectors: %v", strings.Join(enabledCollectorList, ", ")))
 
-	http.HandleFunc(*metricsPath, withConcurrencyLimit(*maxRequests, collectors.BuildServeHTTP(*disableExporterMetrics, *timeoutMargin)))
-	http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
+	mux := http.NewServeMux()
+	mux.HandleFunc(*metricsPath, withConcurrencyLimit(*maxRequests, collectors.BuildServeHTTP(*disableExporterMetrics, *timeoutMargin)))
+	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "application/json")
 		_, err := fmt.Fprintln(w, `{"status":"ok"}`)
 		if err != nil {
 			_ = level.Debug(logger).Log("Failed to write to stream", "err", err)
 		}
 	})
-	http.HandleFunc("/version", func(w http.ResponseWriter, r *http.Request) {
+	mux.HandleFunc("/version", func(w http.ResponseWriter, r *http.Request) {
 		// we can't use "version" directly as it is a package, and not an object that
 		// can be serialized.
 		err := json.NewEncoder(w).Encode(prometheusVersion{
@@ -197,6 +202,15 @@ func main() {
 			http.Error(w, fmt.Sprintf("error encoding JSON: %s", err), http.StatusInternalServerError)
 		}
 	})
+
+	if *debugEnabled {
+		mux.HandleFunc("/debug/pprof/", pprof.Index)
+		mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+		mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+		mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+		mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+	}
+
 	if *metricsPath != "/" && *metricsPath != "" {
 		landingConfig := web.LandingConfig{
 			Name:        "Windows Exporter",
@@ -222,7 +236,7 @@ func main() {
 			_ = level.Error(logger).Log("msg", "failed to generate landing page", "err", err)
 			os.Exit(1)
 		}
-		http.Handle("/", landingPage)
+		mux.Handle("/", landingPage)
 	}
 
 	_ = level.Info(logger).Log("msg", "Starting windows_exporter", "version", version.Info())
@@ -230,7 +244,7 @@ func main() {
 	_ = level.Debug(logger).Log("msg", "Go MAXPROCS", "procs", runtime.GOMAXPROCS(0))
 
 	go func() {
-		server := &http.Server{}
+		server := &http.Server{Handler: mux}
 		if err := web.ListenAndServe(server, webConfig, logger); err != nil {
 			_ = level.Error(logger).Log("msg", "cannot start windows_exporter", "err", err)
 			os.Exit(1)