summaryrefslogtreecommitdiff
path: root/vendor/codeberg.org/gruf/go-debug/pprof_on.go
blob: 4f91aa0924e6a301917cca7f1f001f781c171f34 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
//go:build debug || debugenv
// +build debug debugenv

package debug

import (
	"net/http"
	"net/http/pprof"
	"strings"
)

// ServePprof will start an HTTP server serving /debug/pprof only if debug enabled.
func ServePprof(addr string) error {
	if !DEBUG {
		// debug disabled in env
		return nil
	}
	handler := WithPprof(nil)
	return http.ListenAndServe(addr, handler)
}

// WithPprof will add /debug/pprof handling (provided by "net/http/pprof") only if debug enabled.
func WithPprof(handler http.Handler) http.Handler {
	if !DEBUG {
		// debug disabled in env
		return handler
	}

	// Default serve mux is setup with pprof
	pprofmux := http.DefaultServeMux

	if pprofmux == nil {
		// Someone nil'ed the default mux
		pprofmux = &http.ServeMux{}
		pprofmux.HandleFunc("/debug/pprof/", pprof.Index)
		pprofmux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
		pprofmux.HandleFunc("/debug/pprof/profile", pprof.Profile)
		pprofmux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
		pprofmux.HandleFunc("/debug/pprof/trace", pprof.Trace)
	}

	if handler == nil {
		// Ensure handler is non-nil
		handler = http.NotFoundHandler()
	}

	// Debug enabled, return wrapped handler func
	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		const prefix = "/debug/pprof"

		// /debug/pprof(/.*)? -> pass to pprofmux
		if strings.HasPrefix(r.URL.Path, prefix) {
			path := r.URL.Path[len(prefix):]
			if path == "" || path[0] == '/' {
				pprofmux.ServeHTTP(rw, r)
				return
			}
		}

		// .* -> pass to handler
		handler.ServeHTTP(rw, r)
	})
}