go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/pprof/pprof.go (about)

     1  // Copyright 2017 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package pprof is similar to net/http/pprof, except it supports auth.
    16  //
    17  // Use it instead of net/http/pprof in LUCI server environments.
    18  //
    19  // It uses temporary HMAC-based tokens (generated through admin portal) for
    20  // authenticating requests. Requires a secret store to be installed in the
    21  // base middleware.
    22  package pprof
    23  
    24  import (
    25  	"fmt"
    26  	"html"
    27  	"net/http"
    28  
    29  	"go.chromium.org/luci/common/retry/transient"
    30  
    31  	"go.chromium.org/luci/server/pprof/internal"
    32  	"go.chromium.org/luci/server/router"
    33  )
    34  
    35  var pprofRoutes = map[string]http.HandlerFunc{
    36  	"cmdline": internal.Cmdline,
    37  	"profile": internal.Profile,
    38  	"symbol":  internal.Symbol,
    39  	"trace":   internal.Trace,
    40  }
    41  
    42  // InstallHandlers installs HTTP handlers for pprof routes.
    43  func InstallHandlers(r *router.Router, base router.MiddlewareChain) {
    44  	// Pprof native routing structure is not supported by julienschmidt/httprouter
    45  	// since it mixes prefix matches and direct matches. So we'll have to do some
    46  	// manual routing for paths under /debug/pprof :(
    47  	r.GET("/debug/pprof/*path", base, func(c *router.Context) {
    48  		// Validate the token generated through the portal page.
    49  		tok := c.Request.FormValue("tok")
    50  		if tok == "" {
    51  			sendError(c.Writer, "Missing 'tok' query parameter with a pprof token", http.StatusBadRequest)
    52  			return
    53  		}
    54  		switch err := checkToken(c.Request.Context(), tok); {
    55  		case transient.Tag.In(err):
    56  			sendError(c.Writer, fmt.Sprintf("Transient error, please retry: %s", err), http.StatusInternalServerError)
    57  			return
    58  		case err != nil:
    59  			sendError(c.Writer, fmt.Sprintf("Bad pprof token: %s", err), http.StatusBadRequest)
    60  			return
    61  		}
    62  
    63  		// Manually route. See init() in go/src/net/http/pprof/pprof.go.
    64  		h := pprofRoutes[c.Params.ByName("path")]
    65  		if h == nil {
    66  			h = internal.Index
    67  		}
    68  		h(c.Writer, c.Request)
    69  	})
    70  }
    71  
    72  const errorPage = `<html>
    73  <head>
    74  <title>/debug/pprof/</title>
    75  </head>
    76  <body>
    77  <h2>Error</h2>
    78  <p>%s.</p>
    79  <hr>
    80  <p>To generate a new pprof token visit <a href="/admin/portal/pprof">the admin
    81  portal</a> and follow instructions there.</p>
    82  </body>
    83  </html>
    84  `
    85  
    86  func sendError(w http.ResponseWriter, msg string, code int) {
    87  	w.Header().Set("Content-Type", "text/html")
    88  	w.WriteHeader(code)
    89  	w.Write([]byte(fmt.Sprintf(errorPage, html.EscapeString(msg))))
    90  }