pprof auto

pprof-auth.go

package main

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

func handler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprint(w, "Hello World")
}

func AuthKeyHandler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var token string

		tokens, ok := r.Header["Authorization"]
		if ok && len(tokens) >= 1 {
			token = tokens[0]
			token = strings.TrimPrefix(token, "Bearer ")
		}

		// If the token is empty...
		if token == "" {
			http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
			return
		}

		if token == "abcd" {
			next.ServeHTTP(w, r)
		}
	})
}

func main() {
	pprofHandlers := []struct {
		route   string
		handler func(http.ResponseWriter, *http.Request)
	}{
		{
			route:   "/debug/pprof/",
			handler: pprof.Index,
		},
		{
			route:   "/debug/pprof/cmdline",
			handler: pprof.Cmdline,
		},
		{
			route:   "/debug/pprof/symbol",
			handler: pprof.Symbol,
		},
		{
			route:   "/debug/pprof/trace",
			handler: pprof.Trace,
		},
	}

	pprofMux := http.NewServeMux()
	for _, h := range pprofHandlers {
		pprofMux.Handle(h.route, AuthKeyHandler(http.HandlerFunc(h.handler)))
	}

	http.ListenAndServe(":8080", pprofMux)
}