diff --git a/cmd/serve/cmd.go b/cmd/serve/cmd.go index cec3f67..73f13e2 100644 --- a/cmd/serve/cmd.go +++ b/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 diff --git a/server/admin/admin.go b/server/admin/admin.go index 7b7a329..5d7035e 100644 --- a/server/admin/admin.go +++ b/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") } diff --git a/server/server.go b/server/server.go index cf4ce67..2338ed7 100644 --- a/server/server.go +++ b/server/server.go @@ -5,6 +5,7 @@ import ( "io" "net" "net/http" + "net/http/pprof" "os" "os/signal" "sync" @@ -27,10 +28,11 @@ 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) { @@ -38,8 +40,9 @@ func NewTohServer(options Options) (*TohServer, error) { 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), @@ -47,16 +50,21 @@ func NewTohServer(options Options) (*TohServer, error) { 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)