From 124fc0ecacae55f68899720e944d5c8b00be0c07 Mon Sep 17 00:00:00 2001
From: mstar <me@mstar.dev>
Date: Tue, 15 Oct 2024 21:26:48 +0200
Subject: [PATCH] Add profiling endpoints

---
 config/config.go           |  6 ++++
 server/constants.go        |  1 +
 server/healthAndMetrics.go | 68 ++++++++++++++++++++++++++++++++++++++
 server/server.go           | 15 +++++++--
 4 files changed, 88 insertions(+), 2 deletions(-)
 create mode 100644 server/healthAndMetrics.go

diff --git a/config/config.go b/config/config.go
index 3addf44..e53b696 100644
--- a/config/config.go
+++ b/config/config.go
@@ -45,6 +45,11 @@ type ConfigAdmin struct {
 	// A one time password used to verify account access to the root admin
 	// after a server has been created and before the account could be linked to a passkey
 	FirstTimeSetupOTP string `toml:"first_time_setup_otp"`
+	// Password for protecting profiling data from unauthorised access
+	// An empty string equals to no password
+	// The password has to be supplied in the `password` GET form value for all requests
+	// to /profiling/*
+	ProfilingPassword string `toml:"profiling_password"`
 }
 
 type ConfigStorage struct {
@@ -106,6 +111,7 @@ var defaultConfig Config = Config{
 	Admin: ConfigAdmin{
 		Username:          "server-admin",
 		FirstTimeSetupOTP: "Example otp password",
+		ProfilingPassword: "Example profiling password",
 	},
 	Webauthn: ConfigWebAuthn{
 		DisplayName: "Linstrom",
diff --git a/server/constants.go b/server/constants.go
index f62c6fa..d9ba58b 100644
--- a/server/constants.go
+++ b/server/constants.go
@@ -14,4 +14,5 @@ const (
 	HttpErrIdMissingContextValue
 	HttpErrIdDbFailure
 	HttpErrIdNotAuthenticated
+	HttpErrIdJsonMarshalFail
 )
diff --git a/server/healthAndMetrics.go b/server/healthAndMetrics.go
new file mode 100644
index 0000000..d9c52e8
--- /dev/null
+++ b/server/healthAndMetrics.go
@@ -0,0 +1,68 @@
+package server
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/pprof"
+	"runtime"
+	"runtime/debug"
+	"time"
+
+	"gitlab.com/mstarongitlab/goutils/other"
+)
+
+func setupProfilingHandler() http.Handler {
+	router := http.NewServeMux()
+	router.HandleFunc("GET /current-goroutines", metricActiveGoroutinesHandler)
+	router.HandleFunc("GET /memory", metricMemoryStatsHandler)
+	router.HandleFunc("GET /pprof/cpu", pprof.Profile)
+	router.Handle("GET /pprof/memory", pprof.Handler("heap"))
+	router.Handle("GET /pprof/goroutines", pprof.Handler("goroutine"))
+	router.Handle("GET /pprof/blockers", pprof.Handler("block"))
+
+	return router
+}
+
+func isAliveHandler(w http.ResponseWriter, r *http.Request) {
+	fmt.Fprint(w, "yup")
+}
+
+func metricActiveGoroutinesHandler(w http.ResponseWriter, r *http.Request) {
+	fmt.Fprintf(w, "{\"goroutines\": %d}", runtime.NumGoroutine())
+}
+
+func metricMemoryStatsHandler(w http.ResponseWriter, r *http.Request) {
+	type OutData struct {
+		CollectedAt          time.Time `json:"collected_at"`
+		HeapUsed             uint64    `json:"heap_used"`
+		HeapIdle             uint64    `json:"heap_idle"`
+		StackUsed            uint64    `json:"stack_used"`
+		GCLastFired          time.Time `json:"gc_last_fired"`
+		GCNextTargetHeapSize uint64    `json:"gc_next_target_heap_size"`
+	}
+	stats := runtime.MemStats{}
+	gcStats := debug.GCStats{}
+	runtime.ReadMemStats(&stats)
+	debug.ReadGCStats(&gcStats)
+	outData := OutData{
+		CollectedAt:          time.Now(),
+		HeapUsed:             stats.HeapInuse,
+		HeapIdle:             stats.HeapIdle,
+		StackUsed:            stats.StackInuse,
+		GCLastFired:          gcStats.LastGC,
+		GCNextTargetHeapSize: stats.NextGC,
+	}
+
+	jsonData, err := json.Marshal(&outData)
+	if err != nil {
+		other.HttpErr(
+			w,
+			HttpErrIdJsonMarshalFail,
+			"Failed to encode return data",
+			http.StatusInternalServerError,
+		)
+		return
+	}
+	fmt.Fprint(w, string(jsonData))
+}
diff --git a/server/server.go b/server/server.go
index 27e415e..47f7882 100644
--- a/server/server.go
+++ b/server/server.go
@@ -1,13 +1,13 @@
 package server
 
 import (
-	"fmt"
 	"io/fs"
 	"net/http"
 
 	"github.com/mstarongithub/passkey"
 	"github.com/rs/zerolog/log"
 	"gitlab.com/mstarongitlab/goutils/other"
+	"gitlab.com/mstarongitlab/linstrom/config"
 	"gitlab.com/mstarongitlab/linstrom/storage"
 )
 
@@ -32,7 +32,18 @@ func buildRootHandler(pkey *passkey.Passkey, reactiveFS, staticFS fs.FS) http.Ha
 	pkey.MountRoutes(mux, "/webauthn/")
 	mux.Handle("/", setupFrontendRouter(reactiveFS, staticFS))
 	mux.Handle("/pk/", http.StripPrefix("/pk", http.FileServer(http.Dir("pk-auth"))))
-	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, true) })
+	mux.HandleFunc("/alive", isAliveHandler)
+
+	profilingHandler := setupProfilingHandler()
+	mux.Handle("/profiling/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Only allow access to profiling if a password is provided and it matches the one in the config
+		if r.FormValue("password") != config.GlobalConfig.Admin.ProfilingPassword {
+			// Specifically reply with a plain 404
+			http.Error(w, "", 404)
+			return
+		}
+		profilingHandler.ServeHTTP(w, r)
+	}))
 	mux.Handle(
 		"/authonly/",
 		pkey.Auth(