http: Add support for profiling

2 files changed, 25 insertions, 3 deletions

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()
 }