From f359aea2faa5004b385d91a6c8668d6ba6f4cdc2 Mon Sep 17 00:00:00 2001
From: Chris O'Haver <cohaver@infoblox.com>
Date: Mon, 24 Apr 2017 10:27:26 -0400
Subject: [PATCH] Pprof listen (#639)

* add listen addr option

* Add listen address option to pprof

* There is configuration

* code styling
---
 middleware/pprof/README.md     | 18 +++++++++++++++---
 middleware/pprof/pprof.go      |  8 ++++----
 middleware/pprof/setup.go      | 15 +++++++++++++--
 middleware/pprof/setup_test.go |  2 ++
 4 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/middleware/pprof/README.md b/middleware/pprof/README.md
index 86bdc750b..b1a5f12c4 100644
--- a/middleware/pprof/README.md
+++ b/middleware/pprof/README.md
@@ -11,14 +11,14 @@ For more information, please see [Go's pprof
 documentation](https://golang.org/pkg/net/http/pprof/) and read
 [Profiling Go Programs](https://blog.golang.org/profiling-go-programs).
 
-There is not configuration.
-
 ## Syntax
 
 ~~~
-pprof
+pprof [ADDRESS]
 ~~~
 
+If not specified, ADDRESS defaults to localhost:6053.
+
 ## Examples
 
 Enable pprof endpoints:
@@ -26,3 +26,15 @@ Enable pprof endpoints:
 ~~~
 pprof
 ~~~
+
+Listen on an alternate address:
+
+~~~
+pprof 10.9.8.7:6060
+~~~
+
+Listen on an all addresses on port 6060:
+
+~~~
+pprof :6060
+~~~
diff --git a/middleware/pprof/pprof.go b/middleware/pprof/pprof.go
index aef834a68..020776ecf 100644
--- a/middleware/pprof/pprof.go
+++ b/middleware/pprof/pprof.go
@@ -10,12 +10,13 @@ import (
 )
 
 type handler struct {
-	ln  net.Listener
-	mux *http.ServeMux
+	addr string
+	ln   net.Listener
+	mux  *http.ServeMux
 }
 
 func (h *handler) Startup() error {
-	ln, err := net.Listen("tcp", addr)
+	ln, err := net.Listen("tcp", h.addr)
 	if err != nil {
 		log.Printf("[ERROR] Failed to start pprof handler: %s", err)
 		return err
@@ -44,6 +45,5 @@ func (h *handler) Shutdown() error {
 }
 
 const (
-	addr = "localhost:6053"
 	path = "/debug/pprof"
 )
diff --git a/middleware/pprof/setup.go b/middleware/pprof/setup.go
index 725dc0d63..f94eef710 100644
--- a/middleware/pprof/setup.go
+++ b/middleware/pprof/setup.go
@@ -1,6 +1,7 @@
 package pprof
 
 import (
+	"net"
 	"sync"
 
 	"github.com/coredns/coredns/middleware"
@@ -8,6 +9,8 @@ import (
 	"github.com/mholt/caddy"
 )
 
+const defaultAddr = "localhost:6053"
+
 func init() {
 	caddy.RegisterPlugin("pprof", caddy.Plugin{
 		ServerType: "dns",
@@ -17,11 +20,20 @@ func init() {
 
 func setup(c *caddy.Controller) error {
 	found := false
+	h := &handler{addr: defaultAddr}
 	for c.Next() {
 		if found {
 			return middleware.Error("pprof", c.Err("pprof can only be specified once"))
 		}
-		if len(c.RemainingArgs()) != 0 {
+		args := c.RemainingArgs()
+		if len(args) == 1 {
+			h.addr = args[0]
+			_, _, e := net.SplitHostPort(h.addr)
+			if e != nil {
+				return e
+			}
+		}
+		if len(args) > 1 {
 			return middleware.Error("pprof", c.ArgErr())
 		}
 		if c.NextBlock() {
@@ -30,7 +42,6 @@ func setup(c *caddy.Controller) error {
 		found = true
 	}
 
-	h := &handler{}
 	pprofOnce.Do(func() {
 		c.OnStartup(h.Startup)
 		c.OnShutdown(h.Shutdown)
diff --git a/middleware/pprof/setup_test.go b/middleware/pprof/setup_test.go
index af46fd415..eaa4cb37e 100644
--- a/middleware/pprof/setup_test.go
+++ b/middleware/pprof/setup_test.go
@@ -12,6 +12,8 @@ func TestPProf(t *testing.T) {
 		shouldErr bool
 	}{
 		{`pprof`, false},
+		{`pprof 1.2.3.4:1234`, false},
+		{`pprof :1234`, false},
 		{`pprof {}`, true},
 		{`pprof /foo`, true},
 		{`pprof {