pprof: Only handle if path matches /debug/pprof, add tests

caddy/setup/pprof.go

@@ -20,7 +20,8 @@ func PProf(c *Controller) (middleware.Middleware, error) {
 		}
 		found = true
 	}
+
 	return func(next middleware.Handler) middleware.Handler {
-		return pprof.New(next)
+		return &pprof.Handler{Next: next, Mux: pprof.NewMux()}
 	}, nil
 }

dist/CHANGES.txt

 CHANGES
 
 <master>
+- New pprof directive for exposing process performance profile
+- Toggle case-sensitive path matching with environment variable
 - proxy: New max_conns setting to limit max connections per upstream
 - Internal improvements, restructuring, and bug fixes
 

middleware/pprof/pprof.go

@@ -7,28 +7,35 @@ import (
 	"github.com/mholt/caddy/middleware"
 )
 
-//Handler is a simple struct whose ServeHTTP will delegate relevant pprof endpoints to net/http/pprof
-type handler struct {
-	mux *http.ServeMux
+// BasePath is the base path to match for all pprof requests.
+const BasePath = "/debug/pprof"
+
+// Handler is a simple struct whose ServeHTTP will delegate pprof
+// endpoints to their equivalent net/http/pprof handlers.
+type Handler struct {
+	Next middleware.Handler
+	Mux  *http.ServeMux
 }
 
-//New creates a new pprof middleware
-func New(next middleware.Handler) middleware.Handler {
-	//pretty much copying what pprof does on init: https://golang.org/src/net/http/pprof/pprof.go#L67
-	mux := http.NewServeMux()
-	mux.HandleFunc("/debug/pprof/", pp.Index)
-	mux.HandleFunc("/debug/pprof/cmdline", pp.Cmdline)
-	mux.HandleFunc("/debug/pprof/profile", pp.Profile)
-	mux.HandleFunc("/debug/pprof/symbol", pp.Symbol)
-	mux.HandleFunc("/debug/pprof/trace", pp.Trace)
-	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-		next.ServeHTTP(w, r)
-	})
-	return &handler{mux}
+// ServeHTTP handles requests to BasePath with pprof, or passes
+// all other requests up the chain.
+func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
+	if middleware.Path(r.URL.Path).Matches(BasePath) {
+		h.Mux.ServeHTTP(w, r)
+		return 0, nil
+	}
+	return h.Next.ServeHTTP(w, r)
 }
 
-func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
-	rec := middleware.NewResponseRecorder(w)
-	h.mux.ServeHTTP(rec, r)
-	return rec.Status(), nil
+// NewMux returns a new http.ServeMux that routes pprof requests.
+// It pretty much copies what the std lib pprof does on init:
+// https://golang.org/src/net/http/pprof/pprof.go#L67
+func NewMux() *http.ServeMux {
+	mux := http.NewServeMux()
+	mux.HandleFunc(BasePath+"/", pp.Index)
+	mux.HandleFunc(BasePath+"/cmdline", pp.Cmdline)
+	mux.HandleFunc(BasePath+"/profile", pp.Profile)
+	mux.HandleFunc(BasePath+"/symbol", pp.Symbol)
+	mux.HandleFunc(BasePath+"/trace", pp.Trace)
+	return mux
 }

middleware/pprof/pprof_test.go

+package pprof
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/mholt/caddy/middleware"
+)
+
+func TestServeHTTP(t *testing.T) {
+	h := Handler{
+		Next: middleware.HandlerFunc(nextHandler),
+		Mux:  NewMux(),
+	}
+
+	w := httptest.NewRecorder()
+	r, err := http.NewRequest("GET", "/debug/pprof", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	status, err := h.ServeHTTP(w, r)
+
+	if status != 0 {
+		t.Errorf("Expected status %d but got %d", 0, status)
+	}
+	if err != nil {
+		t.Errorf("Expected nil error, but got: %v", err)
+	}
+	if w.Body.String() == "content" {
+		t.Errorf("Expected pprof to handle request, but it didn't")
+	}
+
+	w = httptest.NewRecorder()
+	r, err = http.NewRequest("GET", "/foo", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	status, err = h.ServeHTTP(w, r)
+	if status != http.StatusNotFound {
+		t.Errorf("Test two: Expected status %d but got %d", http.StatusNotFound, status)
+	}
+	if err != nil {
+		t.Errorf("Test two: Expected nil error, but got: %v", err)
+	}
+	if w.Body.String() != "content" {
+		t.Errorf("Expected pprof to pass the request thru, but it didn't; got: %s", w.Body.String())
+	}
+}
+
+func nextHandler(w http.ResponseWriter, r *http.Request) (int, error) {
+	fmt.Fprintf(w, "content")
+	return http.StatusNotFound, nil
+}