Include handler and template

kolide · Oct 23, 2017 · 53cc8b2 · 53cc8b2
1 parent f3aa699
commit 53cc8b2
Show file tree

Hide file tree

Showing 3 changed files with 93 additions and 17 deletions.
diff --git a/Gopkg.lock b/Gopkg.lock
diff --git a/Gopkg.toml b/Gopkg.toml
@@ -45,7 +45,3 @@
 
 [[constraint]]
   name = "google.golang.org/grpc"
-
-[[constraint]]
-  branch = "master"
-  name = "github.com/e-dard/netbug"
diff --git a/debug/debug.go b/debug/debug.go
@@ -3,21 +3,26 @@ package debug
 import (
 	"context"
 	"fmt"
+	"html/template"
 	"io/ioutil"
 	"net"
 	"net/http"
+	nhpprof "net/http/pprof"
+	"net/url"
 	"os"
 	"os/signal"
+	"runtime/pprof"
+	"strings"
 	"syscall"
 
-	"github.com/e-dard/netbug"
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
 	"github.com/google/uuid"
 	"github.com/pkg/errors"
 )
 
 const debugSignal = syscall.SIGUSR1
+const debugPrefix = "/debug/"
 
 // AttachDebugHandler will attach a signal handler that will toggle the debug
 // server state when SIGUSR1 is sent to the process.
@@ -44,11 +49,12 @@ func AttachDebugHandler(addrPath string, logger log.Logger) {
 					"msg", "error shutting down debug server",
 					"err", err,
 				)
-			} else {
-				level.Info(logger).Log(
-					"msg", "shutdown debug server",
-				)
+				continue
 			}
+
+			level.Info(logger).Log(
+				"msg", "shutdown debug server",
+			)
 		}
 	}()
 }
@@ -62,7 +68,7 @@ func startDebugServer(addrPath string, logger log.Logger) (*http.Server, error)
 
 	// Start the debug server
 	r := http.NewServeMux()
-	netbug.RegisterAuthHandler(token.String(), "/debug/", r)
+	registerAuthHandler(token.String(), r, logger)
 	serv := http.Server{
 		Handler: r,
 	}
@@ -90,3 +96,83 @@ func startDebugServer(addrPath string, logger log.Logger) (*http.Server, error)
 
 	return &serv, nil
 }
+
+// The below handler code is adapted from MIT licensed github.com/e-dard/netbug
+func handler(token string, logger log.Logger) http.Handler {
+	info := struct {
+		Profiles []*pprof.Profile
+		Token    string
+	}{
+		Profiles: pprof.Profiles(),
+		Token:    url.QueryEscape(token),
+	}
+
+	h := func(w http.ResponseWriter, r *http.Request) {
+		name := strings.TrimPrefix(r.URL.Path, "/")
+		switch name {
+		case "":
+			// Index page.
+			if err := indexTmpl.Execute(w, info); err != nil {
+				level.Info(logger).Log(
+					"msg", "error rendering debug template",
+					"err", err,
+				)
+				return
+			}
+		case "cmdline":
+			nhpprof.Cmdline(w, r)
+		case "profile":
+			nhpprof.Profile(w, r)
+		case "trace":
+			nhpprof.Trace(w, r)
+		case "symbol":
+			nhpprof.Symbol(w, r)
+		default:
+			// Provides access to all profiles under runtime/pprof
+			nhpprof.Handler(name).ServeHTTP(w, r)
+		}
+	}
+	return http.HandlerFunc(h)
+}
+
+func authHandler(token string, logger log.Logger) http.Handler {
+	h := handler(token, logger)
+	ah := func(w http.ResponseWriter, r *http.Request) {
+		if r.FormValue("token") == token {
+			h.ServeHTTP(w, r)
+		} else {
+			w.WriteHeader(http.StatusUnauthorized)
+			fmt.Fprintln(w, "Unauthorized.")
+		}
+	}
+	return http.HandlerFunc(ah)
+}
+
+func registerAuthHandler(token string, mux *http.ServeMux, logger log.Logger) {
+	mux.Handle(debugPrefix, http.StripPrefix(debugPrefix, authHandler(token, logger)))
+}
+
+var indexTmpl = template.Must(template.New("index").Parse(`<html>
+  <head>
+    <title>Debug Information</title>
+  </head>
+  <br>
+  <body>
+    profiles:<br>
+    <table>
+    {{range .Profiles}}
+      <tr><td align=right>{{.Count}}<td><a href="{{.Name}}?debug=1&token={{$.Token}}">{{.Name}}</a>
+    {{end}}
+    <tr><td align=right><td><a href="profile?token={{.Token}}">CPU</a>
+    <tr><td align=right><td><a href="trace?seconds=5&token={{.Token}}">5-second trace</a>
+    <tr><td align=right><td><a href="trace?seconds=30&token={{.Token}}">30-second trace</a>
+    </table>
+    <br>
+    debug information:<br>
+    <table>
+      <tr><td align=right><td><a href="cmdline?token={{.Token}}">cmdline</a>
+      <tr><td align=right><td><a href="symbol?token={{.Token}}">symbol</a>
+    <tr><td align=right><td><a href="goroutine?debug=2&token={{.Token}}">full goroutine stack dump</a><br>
+    <table>
+  </body>
+</html>`))