pkg server: add an optional debugMode

cmd/serve/cmd.go

@@ -18,6 +18,7 @@ func init() {
 	Cmd.Flags().String("acl", "acl.json", "file containing access control rules")
 	Cmd.Flags().String("admin-key", "", "key to access the admin api (leave blank to disable admin api)")
 	Cmd.Flags().String("copy-buf", "32ki", "buffer size for copying network data")
+	Cmd.Flags().Bool("debug-mode", false, "run in debug mode (for developers)")
 	Cmd.Flags().StringP("listen", "l", "127.0.0.1:9986", "http server listen address")
 }
 
@@ -47,6 +48,10 @@ func processServerOptions(cmd *cobra.Command) (options server.Options, err error
 	if err != nil {
 		return
 	}
+	options.DebugMode, err = cmd.Flags().GetBool("debug-mode")
+	if err != nil {
+		return
+	}
 	copyBuf, err := cmd.Flags().GetString("copy-buf")
 	if err != nil {
 		return

server/admin/admin.go

@@ -14,14 +14,14 @@ type AdminAPI struct {
 	ACL *acl.ACL
 }
 
-func (s *AdminAPI) Register() {
+func (s *AdminAPI) Register(mux *http.ServeMux) {
 	if !s.ACL.AdminEnabled() {
 		return
 	}
-	http.HandleFunc("/admin/acl/key", s.withAuth(s.HandleKey))
-	http.HandleFunc("/admin/acl/limit", s.withAuth(s.HandleLimit))
-	http.HandleFunc("/admin/acl/usage", s.withAuth(s.HandleUsage))
-	http.HandleFunc("/admin/acl", s.withAuth(s.HandleShowACL))
+	mux.HandleFunc("/admin/acl/key", s.withAuth(s.HandleKey))
+	mux.HandleFunc("/admin/acl/limit", s.withAuth(s.HandleLimit))
+	mux.HandleFunc("/admin/acl/usage", s.withAuth(s.HandleUsage))
+	mux.HandleFunc("/admin/acl", s.withAuth(s.HandleShowACL))
 	logrus.Info("admin api(/admin/**) is enabled")
 }
 

server/server.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"net/http/pprof"
 	"os"
 	"os/signal"
 	"sync"
@@ -27,36 +28,43 @@ type TohServer struct {
 }
 
 type Options struct {
-	Listen   string // http server listen address. required
-	ACL      string // acl json file path. required
-	Buf      uint64 // pipe buffer size, default is 1472. optional
-	AdminKey string // admin api authenticate key. optional
+	Listen    string // http server listen address. required
+	ACL       string // acl json file path. required
+	Buf       uint64 // pipe buffer size, default is 1472. optional
+	AdminKey  string // admin api authenticate key. optional
+	DebugMode bool   // optional
 }
 
 func NewTohServer(options Options) (*TohServer, error) {
 	acl, err := acl.NewACL(options.ACL, options.AdminKey)
 	if err != nil {
 		return nil, err
 	}
-	return &TohServer{
-		httpServer:       &http.Server{Addr: options.Listen},
+	mux := http.NewServeMux()
+	srv := TohServer{
+		httpServer:       &http.Server{Addr: options.Listen, Handler: mux},
 		acl:              acl,
 		adminAPI:         &admin.AdminAPI{ACL: acl},
 		trafficEventChan: make(chan *TrafficEvent, 2048),
 		bufPool: &sync.Pool{New: func() any {
 			buf := make([]byte, max(1472, options.Buf))
 			return &buf
 		}},
-	}, nil
+	}
+
+	srv.adminAPI.Register(mux)
+	mux.HandleFunc("/stats", srv.handleShowStats)
+	mux.HandleFunc("/", srv.handleUpgradeWebSocket)
+	if options.DebugMode {
+		srv.debugPprofRegister(mux)
+	}
+
+	return &srv, nil
 }
 
 func (s *TohServer) Run() {
 	go s.runTrafficEventConsumeLoop()
 	go s.runShutdownListener()
-	s.adminAPI.Register()
-
-	http.HandleFunc("/stats", s.handleShowStats)
-	http.HandleFunc("/", s.handleUpgradeWebSocket)
 
 	logrus.Infof("server listen on %s now", s.httpServer.Addr)
 	err := s.httpServer.ListenAndServe()
@@ -69,6 +77,21 @@ func (s *TohServer) Shutdown(ctx context.Context) error {
 	return s.httpServer.Shutdown(ctx)
 }
 
+func (s *TohServer) debugPprofRegister(mux *http.ServeMux) {
+	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)
+	mux.HandleFunc("/debug/pprof/allocs", pprof.Handler("allocs").ServeHTTP)
+	mux.HandleFunc("/debug/pprof/block", pprof.Handler("block").ServeHTTP)
+	mux.HandleFunc("/debug/pprof/goroutine", pprof.Handler("goroutine").ServeHTTP)
+	mux.HandleFunc("/debug/pprof/heap", pprof.Handler("heap").ServeHTTP)
+	mux.HandleFunc("/debug/pprof/mutex", pprof.Handler("mutex").ServeHTTP)
+	mux.HandleFunc("/debug/pprof/threadcreate", pprof.Handler("threadcreate").ServeHTTP)
+	logrus.Info("debug api(/debug/pprof/**) is enabled")
+}
+
 func (s *TohServer) handleUpgradeWebSocket(w http.ResponseWriter, r *http.Request) {
 	key := r.Header.Get(spec.HeaderHandshakeKey)
 	network := r.Header.Get(spec.HeaderHandshakeNet)