From 4d9741dda6fa233b535ae8a5fb50e0e833710dc2 Mon Sep 17 00:00:00 2001
From: Matthew Holt <Matthew.Holt+git@gmail.com>
Date: Sat, 19 Mar 2016 19:49:19 -0600
Subject: [PATCH] pprof: Only handle if path matches /debug/pprof, add tests

---
 caddy/setup/pprof.go           |  3 +-
 dist/CHANGES.txt               |  2 ++
 middleware/pprof/pprof.go      | 49 +++++++++++++++++-------------
 middleware/pprof/pprof_test.go | 55 ++++++++++++++++++++++++++++++++++
 4 files changed, 87 insertions(+), 22 deletions(-)
 create mode 100644 middleware/pprof/pprof_test.go

diff --git a/caddy/setup/pprof.go b/caddy/setup/pprof.go
index 26943945..01048502 100644
--- a/caddy/setup/pprof.go
+++ b/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
 }
diff --git a/dist/CHANGES.txt b/dist/CHANGES.txt
index 7fb0f558..299c5b91 100644
--- a/dist/CHANGES.txt
+++ b/dist/CHANGES.txt
@@ -1,6 +1,8 @@
 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
 
diff --git a/middleware/pprof/pprof.go b/middleware/pprof/pprof.go
index d78d34ba..8d8e9c78 100644
--- a/middleware/pprof/pprof.go
+++ b/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
+// 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)
+}
+
+// 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("/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}
-}
-
-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
+	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
 }
diff --git a/middleware/pprof/pprof_test.go b/middleware/pprof/pprof_test.go
new file mode 100644
index 00000000..a9aee20c
--- /dev/null
+++ b/middleware/pprof/pprof_test.go
@@ -0,0 +1,55 @@
+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
+}