1 files changed, 31 insertions, 6 deletions
diff --git a/middleware/pprof/pprof.go b/middleware/pprof/pprof.go
index 677690ff8..2c812f246 100644
--- a/middleware/pprof/pprof.go
+++ b/middleware/pprof/pprof.go
@@ -2,8 +2,9 @@ package pprof
 
 import (
 	"log"
+	"net"
 	"net/http"
-	_ "net/http/pprof"
+	pp "net/http/pprof"
 
 	"github.com/miekg/coredns/middleware"
 
@@ -11,10 +12,10 @@ import (
 	"golang.org/x/net/context"
 )
 
-const addr = "localhost:8053"
-
 type Handler struct {
 	Next middleware.Handler
+	ln   net.Listener
+	mux  *http.ServeMux
 }
 
 // ServeDNS passes all other requests up the chain.
@@ -23,10 +24,34 @@ func (h *Handler) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg
 }
 
 func (h *Handler) Start() error {
+	if ln, err := net.Listen("tcp", addr); err != nil {
+		log.Printf("[ERROR] Failed to start pprof handler: %s", err)
+		return err
+	} else {
+		h.ln = ln
+	}
+
+	h.mux = http.NewServeMux()
+	h.mux.HandleFunc(path+"/", pp.Index)
+	h.mux.HandleFunc(path+"/cmdline", pp.Cmdline)
+	h.mux.HandleFunc(path+"/profile", pp.Profile)
+	h.mux.HandleFunc(path+"/symbol", pp.Symbol)
+	h.mux.HandleFunc(path+"/trace", pp.Trace)
+
 	go func() {
-		if err := http.ListenAndServe(addr, nil); err != nil {
-			log.Printf("[ERROR] Failed to start pprof handler: %s", err)
-		}
+		http.Serve(h.ln, h.mux)
 	}()
 	return nil
 }
+
+func (h *Handler) Shutdown() error {
+	if h.ln != nil {
+		return h.ln.Close()
+	}
+	return nil
+}
+
+const (
+	addr = "localhost:6053"
+	path = "/debug/pprof"
+)