Replaced http wrapper field in plugin Options with HTTPServer. (#229)

Replaced http wrapper field in plugin Options with HTTPServer.

HTTPServer allows somebody embedding pprof to control how the
hosting HTTP server is constructed which makes pprof work in
more environments.

internal/driver/webui.go

@@ -23,6 +23,7 @@ import (
 	gourl "net/url"
 	"os"
 	"os/exec"
+	"strconv"
 	"strings"
 	"time"
 
@@ -81,6 +82,18 @@ type webArgs struct {
 }
 
 func serveWebInterface(hostport string, p *profile.Profile, o *plugin.Options) error {
+	host, portStr, err := net.SplitHostPort(hostport)
+	if err != nil {
+		return fmt.Errorf("could not split http address: %v", err)
+	}
+	port, err := strconv.Atoi(portStr)
+	if err != nil {
+		return fmt.Errorf("invalid port number: %v", err)
+	}
+	if host == "" {
+		host = "localhost"
+	}
+
 	interactiveMode = true
 	ui := makeWebInterface(p, o)
 	for n, c := range pprofCommands {
@@ -93,47 +106,51 @@ func serveWebInterface(hostport string, p *profile.Profile, o *plugin.Options) e
 	ui.help["graph"] = "Display profile as a directed graph"
 	ui.help["reset"] = "Show the entire profile"
 
-	ln, url, isLocal, err := newListenerAndURL(hostport)
-	if err != nil {
-		return err
+	server := o.HTTPServer
+	if server == nil {
+		server = defaultWebServer
 	}
-
-	// authorization wrapper
-	wrap := o.HTTPWrapper
-	if wrap == nil {
-		if isLocal {
-			// Only allow requests from local host.
-			wrap = checkLocalHost
-		} else {
-			wrap = func(h http.Handler) http.Handler { return h }
-		}
+	args := &plugin.HTTPServerArgs{
+		Hostport: net.JoinHostPort(host, portStr),
+		Host:     host,
+		Port:     port,
+		Handlers: map[string]http.Handler{
+			"/":       http.HandlerFunc(ui.dot),
+			"/top":    http.HandlerFunc(ui.top),
+			"/disasm": http.HandlerFunc(ui.disasm),
+			"/source": http.HandlerFunc(ui.source),
+			"/peek":   http.HandlerFunc(ui.peek),
+		},
 	}
 
-	mux := http.NewServeMux()
-	mux.Handle("/", wrap(http.HandlerFunc(ui.dot)))
-	mux.Handle("/top", wrap(http.HandlerFunc(ui.top)))
-	mux.Handle("/disasm", wrap(http.HandlerFunc(ui.disasm)))
-	mux.Handle("/source", wrap(http.HandlerFunc(ui.source)))
-	mux.Handle("/peek", wrap(http.HandlerFunc(ui.peek)))
-
-	s := &http.Server{Handler: mux}
-	go openBrowser(url, o)
-	return s.Serve(ln)
+	go openBrowser("http://"+args.Hostport, o)
+	return server(args)
 }
 
-func newListenerAndURL(hostport string) (ln net.Listener, url string, isLocal bool, err error) {
-	host, _, err := net.SplitHostPort(hostport)
+func defaultWebServer(args *plugin.HTTPServerArgs) error {
+	ln, err := net.Listen("tcp", args.Hostport)
 	if err != nil {
-		return nil, "", false, err
-	}
-	if host == "" {
-		host = "localhost"
-	}
-	if ln, err = net.Listen("tcp", hostport); err != nil {
-		return nil, "", false, err
+		return err
 	}
-	url = fmt.Sprint("http://", net.JoinHostPort(host, fmt.Sprint(ln.Addr().(*net.TCPAddr).Port)))
-	return ln, url, isLocalhost(host), nil
+	isLocal := isLocalhost(args.Host)
+	handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		if isLocal {
+			// Only allow local clients
+			host, _, err := net.SplitHostPort(req.RemoteAddr)
+			if err != nil || !isLocalhost(host) {
+				http.Error(w, "permission denied", http.StatusForbidden)
+				return
+			}
+		}
+		h := args.Handlers[req.URL.Path]
+		if h == nil {
+			// Fall back to default behavior
+			h = http.DefaultServeMux
+		}
+		h.ServeHTTP(w, req)
+	})
+	s := &http.Server{Handler: handler}
+	return s.Serve(ln)
 }
 
 func isLocalhost(host string) bool {
@@ -179,17 +196,6 @@ func openBrowser(url string, o *plugin.Options) {
 	o.UI.PrintErr(u.String())
 }
 
-func checkLocalHost(h http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
-		host, _, err := net.SplitHostPort(req.RemoteAddr)
-		if err != nil || !isLocalhost(host) {
-			http.Error(w, "permission denied", http.StatusForbidden)
-			return
-		}
-		h.ServeHTTP(w, req)
-	})
-}
-
 func varsFromURL(u *gourl.URL) variables {
 	vars := pprofVariables.makeCopy()
 	vars["focus"].value = u.Query().Get("f")
@@ -241,12 +247,6 @@ func (ui *webInterface) render(w http.ResponseWriter, baseURL, tmpl string,
 
 // dot generates a web page containing an svg diagram.
 func (ui *webInterface) dot(w http.ResponseWriter, req *http.Request) {
-	// Disable prefix matching behavior of net/http.
-	if req.URL.Path != "/" {
-		http.NotFound(w, req)
-		return
-	}
-
 	rpt, errList := ui.makeReport(w, req, []string{"svg"})
 	if rpt == nil {
 		return // error already reported

internal/driver/webui_test.go

@@ -23,7 +23,6 @@ import (
 	"net/url"
 	"os/exec"
 	"regexp"
-	"runtime"
 	"testing"
 
 	"github.com/google/pprof/internal/plugin"
@@ -32,27 +31,28 @@ import (
 
 func TestWebInterface(t *testing.T) {
 	prof := makeFakeProfile()
-	ui := makeWebInterface(prof, &plugin.Options{
-		Obj: fakeObjTool{},
-		UI:  &stdUI{},
-	})
 
-	// Start test server.
-	server := httptest.NewServer(http.HandlerFunc(
-		func(w http.ResponseWriter, r *http.Request) {
-			switch r.URL.Path {
-			case "/":
-				ui.dot(w, r)
-			case "/top":
-				ui.top(w, r)
-			case "/disasm":
-				ui.disasm(w, r)
-			case "/peek":
-				ui.peek(w, r)
-			case "/source":
-				ui.source(w, r)
-			}
-		}))
+	// Custom http server creator
+	var server *httptest.Server
+	serverCreated := make(chan bool)
+	creator := func(a *plugin.HTTPServerArgs) error {
+		server = httptest.NewServer(http.HandlerFunc(
+			func(w http.ResponseWriter, r *http.Request) {
+				if h := a.Handlers[r.URL.Path]; h != nil {
+					h.ServeHTTP(w, r)
+				}
+			}))
+		serverCreated <- true
+		return nil
+	}
+
+	// Start server and wait for it to be initialized
+	go serveWebInterface("unused:1234", prof, &plugin.Options{
+		Obj:        fakeObjTool{},
+		UI:         &stdUI{},
+		HTTPServer: creator,
+	})
+	<-serverCreated
 	defer server.Close()
 
 	haveDot := false
@@ -195,55 +195,6 @@ func makeFakeProfile() *profile.Profile {
 	}
 }
 
-func TestNewListenerAndURL(t *testing.T) {
-	if runtime.GOOS == "nacl" {
-		t.Skip("test assumes tcp available")
-	}
-
-	tests := []struct {
-		hostport  string
-		wantErr   bool
-		wantURLRe *regexp.Regexp
-		wantLocal bool
-	}{
-		{
-			hostport:  ":",
-			wantURLRe: regexp.MustCompile(`http://localhost:\d+`),
-			wantLocal: true,
-		},
-		{
-			hostport:  "localhost:",
-			wantURLRe: regexp.MustCompile(`http://localhost:\d+`),
-			wantLocal: true,
-		},
-		{
-			hostport: "http://localhost:12345",
-			wantErr:  true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.hostport, func(t *testing.T) {
-			_, url, isLocal, err := newListenerAndURL(tt.hostport)
-			if tt.wantErr {
-				if err == nil {
-					t.Fatalf("newListenerAndURL(%v) want error; got none", tt.hostport)
-				}
-				return
-			}
-			if err != nil {
-				t.Fatalf("newListenerAndURL(%v) = %v", tt.hostport, err)
-			}
-			if !tt.wantURLRe.MatchString(url) {
-				t.Errorf("newListenerAndURL(%v) URL = %v; want URL matching %v", tt.hostport, url, tt.wantURLRe)
-			}
-			if got, want := isLocal, tt.wantLocal; got != want {
-				t.Errorf("newListenerAndURL(%v) isLocal = %v; want %v", tt.hostport, got, want)
-			}
-		})
-	}
-}
-
 func TestIsLocalHost(t *testing.T) {
 	for _, s := range []string{"localhost:10000", "[::1]:10000", "127.0.0.1:10000"} {
 		host, _, err := net.SplitHostPort(s)

internal/plugin/plugin.go

@@ -33,14 +33,15 @@ type Options struct {
 	Obj     ObjTool
 	UI      UI
 
-	// HTTPWrapper takes a pprof http handler as an argument and
-	// returns the actual handler that should be invoked by http.
-	// A typical use is to add authentication before calling the
-	// pprof handler.
+	// HTTPServer is a function that should block serving http requests,
+	// including the handlers specfied in args.  If non-nil, pprof will
+	// invoke this function if necessary to provide a web interface.
 	//
-	// If HTTPWrapper is nil, a default wrapper will be used that
-	// disallows all requests except from the localhost.
-	HTTPWrapper func(http.Handler) http.Handler
+	// If HTTPServer is nil, pprof will use its own internal HTTP server.
+	//
+	// A common use for a custom HTTPServer is to provide custom
+	// authentication checks.
+	HTTPServer func(args *HTTPServerArgs) error
 }
 
 // Writer provides a mechanism to write data under a certain name,
@@ -195,3 +196,17 @@ type UI interface {
 	// the auto-completion of cmd, if the UI supports auto-completion at all.
 	SetAutoComplete(complete func(string) string)
 }
+
+// HTTPServerArgs contains arguments needed by an HTTP server that
+// is exporting a pprof web interface.
+type HTTPServerArgs struct {
+	// Hostport contains the http server address (derived from flags).
+	Hostport string
+
+	Host string // Host portion of Hostport
+	Port int    // Port portion of Hostport
+
+	// Handlers maps from URL paths to the handler to invoke to
+	// serve that path.
+	Handlers map[string]http.Handler
+}