diff --git a/cmd/echoip/main.go b/cmd/echoip/main.go index ef0965f..3d03634 100644 --- a/cmd/echoip/main.go +++ b/cmd/echoip/main.go @@ -38,6 +38,7 @@ func main() { portLookup := flag.Bool("p", false, "Enable port lookup") template := flag.String("t", "index.html", "Path to template") cacheSize := flag.Int("C", 0, "Size of response cache. Set to 0 to disable") + profile := flag.Bool("P", false, "Enables profiling handlers") var headers multiValueFlag flag.Var(&headers, "H", "Header to trust for remote IP, if present (e.g. X-Real-IP)") flag.Parse() @@ -48,7 +49,7 @@ func main() { log.Fatal(err) } cache := http.NewCache(*cacheSize) - server := http.New(r, cache) + server := http.New(r, cache, *profile) server.IPHeaders = headers if _, err := os.Stat(*template); err == nil { server.Template = *template @@ -69,6 +70,9 @@ func main() { if *cacheSize > 0 { log.Printf("Cache capacity set to %d", *cacheSize) } + if *profile { + log.Printf("Enabling profiling handlers") + } log.Printf("Listening on http://%s", *listen) if err := server.ListenAndServe(*listen); err != nil { log.Fatal(err) diff --git a/http/http.go b/http/http.go index 2b23cf4..0f1f65d 100644 --- a/http/http.go +++ b/http/http.go @@ -10,6 +10,7 @@ import ( "github.com/mpolden/echoip/iputil" "github.com/mpolden/echoip/iputil/geo" "github.com/mpolden/echoip/useragent" + "net/http/pprof" "math/big" "net" @@ -29,6 +30,7 @@ type Server struct { LookupPort func(net.IP, uint64) error cache *Cache gr geo.Reader + profile bool } type Response struct { @@ -57,8 +59,8 @@ type PortResponse struct { Reachable bool `json:"reachable"` } -func New(db geo.Reader, cache *Cache) *Server { - return &Server{cache: cache, gr: db} +func New(db geo.Reader, cache *Cache, profile bool) *Server { + return &Server{cache: cache, gr: db, profile: profile} } func ipFromForwardedForHeader(v string) string { @@ -325,6 +327,13 @@ func cliMatcher(r *http.Request) bool { type appHandler func(http.ResponseWriter, *http.Request) *appError +func wrapHandlerFunc(f http.HandlerFunc) appHandler { + return func(w http.ResponseWriter, r *http.Request) *appError { + f.ServeHTTP(w, r) + return nil + } +} + func (fn appHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { if e := fn(w, r); e != nil { // e is *appError // When Content-Type for error is JSON, we need to marshal the response into JSON @@ -379,6 +388,15 @@ func (s *Server) Handler() http.Handler { r.RoutePrefix("GET", "/port/", s.PortHandler) } + // Profiling + if s.profile { + r.Route("GET", "/debug/pprof/cmdline", wrapHandlerFunc(pprof.Cmdline)) + r.Route("GET", "/debug/pprof/profile", wrapHandlerFunc(pprof.Profile)) + r.Route("GET", "/debug/pprof/symbol", wrapHandlerFunc(pprof.Symbol)) + r.Route("GET", "/debug/pprof/trace", wrapHandlerFunc(pprof.Trace)) + r.RoutePrefix("GET", "/debug/pprof/", wrapHandlerFunc(pprof.Index)) + } + return r.Handler() }