From 98745b3bb9652c3e7cddee0f929fbac19702bf18 Mon Sep 17 00:00:00 2001
From: James Hillyerd <james@hillyerd.com>
Date: Sat, 20 Oct 2018 16:16:09 -0700
Subject: [PATCH] web: Optionally mount /debug/pprof for #120

- web: eliminate use of http.DefaultServeMux
---
 CHANGELOG.md             |  1 +
 doc/config.md            | 15 +++++++++++++++
 pkg/config/config.go     |  1 +
 pkg/server/web/server.go | 14 ++++++++++++--
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 24b7a16..369f279 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 ### Added
 - Use Go 1.11 modules for reproducible builds.
 - SMTP TLS support (thanks kingforaday.)
+- `INBUCKET_WEB_PPROF` configuration option for performance profiling.
 
 ### Changed
 - Docker build now uses Go 1.11 and Alpine 3.8
diff --git a/doc/config.md b/doc/config.md
index e003d3a..c880d4e 100644
--- a/doc/config.md
+++ b/doc/config.md
@@ -35,6 +35,7 @@ variables it supports:
     INBUCKET_WEB_COOKIEAUTHKEY                              Session cipher key (text)
     INBUCKET_WEB_MONITORVISIBLE         true                Show monitor tab in UI?
     INBUCKET_WEB_MONITORHISTORY         30                  Monitor remembered messages
+    INBUCKET_WEB_PPROF                  false               Expose profiling tools on /debug/pprof
     INBUCKET_STORAGE_TYPE               memory              Storage impl: file or memory
     INBUCKET_STORAGE_PARAMS                                 Storage impl parameters, see docs.
     INBUCKET_STORAGE_RETENTIONPERIOD    24h                 Duration to retain messages
@@ -377,6 +378,20 @@ them.
 - Default: `30`
 - Values: Integer greater than or equal to 0
 
+### Performance Profiling & Debug Tools
+
+`INBUCKET_WEB_PPROF`
+
+If true, Go's pprof package will be installed to the `/debug/pprof` URI.  This
+exposes detailed memory and CPU performance data for debugging Inbucket.  If you
+enable this option, please make sure it is not exposed to the public internet,
+as its use can significantly impact performance.
+
+For example usage, see https://golang.org/pkg/net/http/pprof/
+
+- Default: `false`
+- Values: `true` or `false`
+
 
 ## Storage
 
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 9c1228f..b1bf39a 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -100,6 +100,7 @@ type Web struct {
 	CookieAuthKey  string `desc:"Session cipher key (text)"`
 	MonitorVisible bool   `required:"true" default:"true" desc:"Show monitor tab in UI?"`
 	MonitorHistory int    `required:"true" default:"30" desc:"Monitor remembered messages"`
+	PProf          bool   `required:"true" default:"false" desc:"Expose profiling tools on /debug/pprof"`
 }
 
 // Storage contains the mail store configuration.
diff --git a/pkg/server/web/server.go b/pkg/server/web/server.go
index a60735a..bc7bc5e 100644
--- a/pkg/server/web/server.go
+++ b/pkg/server/web/server.go
@@ -6,6 +6,7 @@ import (
 	"expvar"
 	"net"
 	"net/http"
+	"net/http/pprof"
 	"path/filepath"
 	"time"
 
@@ -70,7 +71,16 @@ func Initialize(
 		Msg("Web UI content mapped")
 	Router.PathPrefix("/public/").Handler(http.StripPrefix("/public/",
 		http.FileServer(http.Dir(staticPath))))
-	http.Handle("/", Router)
+	Router.Handle("/debug/vars", expvar.Handler())
+	if conf.Web.PProf {
+		Router.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+		Router.HandleFunc("/debug/pprof/profile", pprof.Profile)
+		Router.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+		Router.HandleFunc("/debug/pprof/trace", pprof.Trace)
+		Router.PathPrefix("/debug/pprof/").HandlerFunc(pprof.Index)
+		log.Warn().Str("module", "web").Str("phase", "startup").
+			Msg("Go pprof tools installed to /debug/pprof")
+	}
 
 	// Session cookie setup
 	if conf.Web.CookieAuthKey == "" {
@@ -88,7 +98,7 @@ func Initialize(
 func Start(ctx context.Context) {
 	server = &http.Server{
 		Addr:         rootConfig.Web.Addr,
-		Handler:      nil,
+		Handler:      Router,
 		ReadTimeout:  60 * time.Second,
 		WriteTimeout: 60 * time.Second,
 	}