feat: make routes configurable independently of each other

gdegiorgio · Nov 23, 2024 · 5410c7e · 5410c7e
1 parent 5254d9e
commit 5410c7e
Show file tree

Hide file tree

Showing 9 changed files with 109 additions and 64 deletions.
diff --git a/cpu.go b/cpu.go
@@ -2,6 +2,11 @@ package systatus
 
 import "net/http"
 
+type CPUHandlerOpts struct {
+	Middlewares []func(next http.HandlerFunc) http.HandlerFunc
+}
 type CPUResponse struct{}
 
-func HandleCPU(w http.ResponseWriter, r *http.Request) {}
+func HandleCPU(opts CPUHandlerOpts) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {}
+}
diff --git a/disk.go b/disk.go
@@ -2,4 +2,11 @@ package systatus
 
 import "net/http"
 
-func HandleDisk(w http.ResponseWriter, r *http.Request) {}
+type DiskHandlerOpts struct {
+	Middlewares []func(next http.HandlerFunc) http.HandlerFunc
+}
+type DiskResponse struct{}
+
+func HandleDisk(opts DiskHandlerOpts) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {}
+}
diff --git a/env.go b/env.go
@@ -7,22 +7,27 @@ import (
 	"strings"
 )
 
+type EnvHandlerOpts struct {
+	Middlewares []func(next http.HandlerFunc) http.HandlerFunc
+}
+
 type EnvResponse struct {
 	Env map[string]string `json:"env"`
 }
 
-func HandleEnv(w http.ResponseWriter, r *http.Request) {
+func HandleEnv(opts EnvHandlerOpts) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		res := EnvResponse{}
+		env := os.Environ()
 
-	res := EnvResponse{}
-	env := os.Environ()
+		res.Env = make(map[string]string, len(env))
 
-	res.Env = make(map[string]string, len(env))
+		for _, val := range env {
+			split := strings.Split(val, "=")
+			res.Env[split[0]] = split[1]
+		}
 
-	for _, val := range env {
-		split := strings.Split(val, "=")
-		res.Env[split[0]] = split[1]
+		w.Header().Add("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(res)
 	}
-
-	w.Header().Add("Content-Type", "application/json")
-	json.NewEncoder(w).Encode(res)
 }
diff --git a/examples/stdlib/main.go b/examples/stdlib/main.go
@@ -16,7 +16,14 @@ func customHealthCheck(w http.ResponseWriter, r *http.Request) {
 }
 
 func main() {
-	opts := systatus.SystatusOptions{Prefix: "/dev", ExposeEnv: true, Healthcheck: customHealthCheck, Middlewares: []func(next http.HandlerFunc) http.HandlerFunc{middleware.Logger}}
+	opts := systatus.SystatusOptions{
+		Prefix:    "/dev",
+		ExposeEnv: true,
+		HealthHandlerOpts: systatus.HealthHandlerOpts{
+			Middlewares: []func(next http.HandlerFunc) http.HandlerFunc{middleware.Logger},
+			Healthcheck: customHealthCheck,
+		},
+	}
 	systatus.Enable(opts)
 	log.Info().Msg("Starting server on :3333")
 	if err := http.ListenAndServe(":3333", nil); err != nil {

diff --git a/health.go b/health.go
@@ -3,13 +3,25 @@ package systatus
 import (
 	"encoding/json"
 	"net/http"
+
+	"github.com/rs/zerolog/log"
 )
 
+type HealthHandlerOpts struct {
+	Middlewares []func(next http.HandlerFunc) http.HandlerFunc
+	Healthcheck func(w http.ResponseWriter, r *http.Request)
+}
+
 type HealthResponse struct {
 	Status string `json:"status"`
 }
 
-func HandleHealth(w http.ResponseWriter, r *http.Request) {
-
-	json.NewEncoder(w).Encode(HealthResponse{Status: "HEALTHY"})
+func HandleHealth(opts HealthHandlerOpts) func(w http.ResponseWriter, r *http.Request) {
+	if opts.Healthcheck != nil {
+		log.Debug().Msg("Found a custom healthcheck")
+		return opts.Healthcheck
+	}
+	return func(w http.ResponseWriter, r *http.Request) {
+		json.NewEncoder(w).Encode(HealthResponse{Status: "HEALTHY"})
+	}
 }
diff --git a/health_test.go b/health_test.go
@@ -28,7 +28,7 @@ func TestHealthEndpointDefaultHandler(t *testing.T) {
 	var h HealthResponse
 	err := json.NewDecoder(res.Body).Decode(&h)
 	assert.NoError(t, err)
-	assert.Equal(t, h.Status, "HEALTHY")
+	assert.Equal(t, "HEALTHY", h.Status)
 }
 
 func TestHealthEndpointCustomHandler(t *testing.T) {
@@ -37,11 +37,14 @@ func TestHealthEndpointCustomHandler(t *testing.T) {
 	opts := SystatusOptions{
 		Prefix: "",
 		Mux:    mux,
-		Healthcheck: func(w http.ResponseWriter, r *http.Request) {
-			w.WriteHeader(200)
-			json.NewEncoder(w).Encode(HealthResponse{Status: "HEALTHY-CUSTOM"})
+		HealthHandlerOpts: HealthHandlerOpts{
+			Healthcheck: func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(200)
+				json.NewEncoder(w).Encode(HealthResponse{Status: "HEALTHY-CUSTOM"})
+			},
 		},
 	}
+
 	Enable(opts)
 
 	ts := httptest.NewServer(mux)
@@ -53,5 +56,5 @@ func TestHealthEndpointCustomHandler(t *testing.T) {
 	var h HealthResponse
 	err := json.NewDecoder(res.Body).Decode(&h)
 	assert.NoError(t, err)
-	assert.Equal(t, h.Status, "HEALTHY-CUSTOM")
+	assert.Equal(t, "HEALTHY-CUSTOM", h.Status)
 }
diff --git a/mem.go b/mem.go
@@ -6,23 +6,26 @@ import (
 	"runtime"
 )
 
+type MemHandlerOpts struct {
+	Middlewares []func(next http.HandlerFunc) http.HandlerFunc
+}
+
 type MemResponse struct {
 	TotalAlloc uint64 `json:"total_alloc"`
 	Alloc      uint64 `json:"alloc"`
 	Sys        uint64 `json:"sys"`
 }
 
-func HandleMem(w http.ResponseWriter, r *http.Request) {
-	res := &MemResponse{}
-
-	var stats runtime.MemStats
-
-	runtime.ReadMemStats(&stats)
-
-	res.Sys = stats.Sys
-	res.TotalAlloc = stats.TotalAlloc
-	res.Alloc = stats.Alloc
+func HandleMem(opts MemHandlerOpts) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		res := &MemResponse{}
+		var stats runtime.MemStats
+		runtime.ReadMemStats(&stats)
+		res.Sys = stats.Sys
+		res.TotalAlloc = stats.TotalAlloc
+		res.Alloc = stats.Alloc
+		w.Header().Add("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(res)
+	}
 
-	w.Header().Add("Content-Type", "application/json")
-	json.NewEncoder(w).Encode(res)
 }
diff --git a/systatus.go b/systatus.go
@@ -6,11 +6,15 @@ import (
 )
 
 type SystatusOptions struct {
-	Mux         *http.ServeMux
-	Prefix      string
-	ExposeEnv   bool
-	Healthcheck func(w http.ResponseWriter, r *http.Request)
-	Middlewares []func(next http.HandlerFunc) http.HandlerFunc
+	Mux               *http.ServeMux
+	Prefix            string
+	ExposeEnv         bool
+	HealthHandlerOpts HealthHandlerOpts
+	CPUHandlerOpts    CPUHandlerOpts
+	EnvHandlerOpts    EnvHandlerOpts
+	DiskHandlerOpts   DiskHandlerOpts
+	UptimeHandlerOpts UptimeHandlerOpts
+	MemHandlerOpts    MemHandlerOpts
 }
 
 func useMiddlewares(handler func(w http.ResponseWriter, r *http.Request), middlewares []func(next http.HandlerFunc) http.HandlerFunc) func(w http.ResponseWriter, r *http.Request) {
@@ -21,25 +25,19 @@ func useMiddlewares(handler func(w http.ResponseWriter, r *http.Request), middle
 }
 
 func Enable(opts SystatusOptions) {
-	var healthcheck = HandleHealth
 
 	mux := http.DefaultServeMux
 
 	if opts.Mux != nil {
 		mux = opts.Mux
 	}
-
-	if opts.Healthcheck != nil {
-		healthcheck = opts.Healthcheck
-	}
-
-	mux.HandleFunc(fmt.Sprintf("%s/health", opts.Prefix), useMiddlewares(healthcheck, opts.Middlewares))
-	mux.HandleFunc(fmt.Sprintf("%s/uptime", opts.Prefix), useMiddlewares(HandleUptime, opts.Middlewares))
-	mux.HandleFunc(fmt.Sprintf("%s/cpu", opts.Prefix), useMiddlewares(HandleCPU, opts.Middlewares))
-	mux.HandleFunc(fmt.Sprintf("%s/mem", opts.Prefix), useMiddlewares(HandleCPU, opts.Middlewares))
-	mux.HandleFunc(fmt.Sprintf("%s/disk", opts.Prefix), useMiddlewares(HandleDisk, opts.Middlewares))
+	mux.HandleFunc(fmt.Sprintf("%s/health", opts.Prefix), useMiddlewares(HandleHealth(opts.HealthHandlerOpts), opts.CPUHandlerOpts.Middlewares))
+	mux.HandleFunc(fmt.Sprintf("%s/uptime", opts.Prefix), useMiddlewares(HandleUptime(opts.UptimeHandlerOpts), opts.UptimeHandlerOpts.Middlewares))
+	mux.HandleFunc(fmt.Sprintf("%s/cpu", opts.Prefix), useMiddlewares(HandleCPU(opts.CPUHandlerOpts), opts.CPUHandlerOpts.Middlewares))
+	mux.HandleFunc(fmt.Sprintf("%s/mem", opts.Prefix), useMiddlewares(HandleMem(opts.MemHandlerOpts), opts.MemHandlerOpts.Middlewares))
+	mux.HandleFunc(fmt.Sprintf("%s/disk", opts.Prefix), useMiddlewares(HandleDisk(opts.DiskHandlerOpts), opts.DiskHandlerOpts.Middlewares))
 
 	if opts.ExposeEnv {
-		mux.HandleFunc(fmt.Sprintf("%s/env", opts.Prefix), HandleEnv)
+		mux.HandleFunc(fmt.Sprintf("%s/env", opts.Prefix), useMiddlewares(HandleEnv(opts.EnvHandlerOpts), opts.EnvHandlerOpts.Middlewares))
 	}
 }
diff --git a/uptime.go b/uptime.go
@@ -7,25 +7,30 @@ import (
 	"strings"
 )
 
+type UptimeHandlerOpts struct {
+	Middlewares []func(next http.HandlerFunc) http.HandlerFunc
+}
 type UptimeResponse struct {
 	Systime string `json:"systime"`
 	Uptime  string `json:"uptime"`
 }
 
-func HandleUptime(w http.ResponseWriter, r *http.Request) {
-	res := UptimeResponse{}
-	cmdoutput, err := exec.Command("/bin/uptime").Output()
-	if err != nil {
-		http.Error(w, "Could not exec uptime command on this machine", http.StatusInternalServerError)
-		return
-	}
-	split := strings.Split(string(cmdoutput), " ")
+func HandleUptime(opts UptimeHandlerOpts) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		res := UptimeResponse{}
+		cmdoutput, err := exec.Command("/bin/uptime").Output()
+		if err != nil {
+			http.Error(w, "Could not exec uptime command on this machine", http.StatusInternalServerError)
+			return
+		}
+		split := strings.Split(string(cmdoutput), " ")
 
-	res.Systime = split[1]
-	// Remove comma e.g 3:05,
-	res.Uptime = strings.Split(split[4], ",")[0]
+		res.Systime = split[1]
+		// Remove comma e.g 3:05,
+		res.Uptime = strings.Split(split[4], ",")[0]
 
-	w.WriteHeader(200)
-	w.Header().Add("Content-Type", "application/json")
-	json.NewEncoder(w).Encode(res)
+		w.WriteHeader(200)
+		w.Header().Add("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(res)
+	}
 }