Add URL prefix to pprof middleware (#2194)

* Add URL prefix to pprof middleware

Signed-off-by: Glenn Lewis <6598971+gmlewis@users.noreply.github.com>

* Minor tweak

Signed-off-by: Glenn Lewis <6598971+gmlewis@users.noreply.github.com>

Signed-off-by: Glenn Lewis <6598971+gmlewis@users.noreply.github.com>
pull/2107/head
Glenn Lewis 2022-11-08 22:59:33 -08:00 committed by GitHub
parent 55fcddda6f
commit 13247206ab
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 134 additions and 14 deletions

View File

@ -24,6 +24,16 @@ After you initiate your Fiber app, you can use the following possibilities:
app.Use(pprof.New())
```
In systems where you have multiple ingress endpoints, it is common to add a URL prefix, like so:
```go
// Default middleware
app.Use(pprof.New(pprof.Config{Prefix: "/endpoint-prefix"}))
```
This prefix will be added to the default path of "/debug/pprof/", for a resulting URL of:
"/endpoint-prefix/debug/pprof/".
## Config
```go
@ -33,6 +43,13 @@ type Config struct {
//
// Optional. Default: nil
Next func(c *fiber.Ctx) bool
// Prefix defines a URL prefix added before "/debug/pprof".
// Note that it should start with (but not end with) a slash.
// Example: "/federated-fiber"
//
// Optional. Default: ""
Prefix string
}
```
@ -42,4 +59,4 @@ type Config struct {
var ConfigDefault = Config{
Next: nil,
}
```
```

View File

@ -8,6 +8,13 @@ type Config struct {
//
// Optional. Default: nil
Next func(c *fiber.Ctx) bool
// Prefix defines a URL prefix added before "/debug/pprof".
// Note that it should start with (but not end with) a slash.
// Example: "/federated-fiber"
//
// Optional. Default: ""
Prefix string
}
var ConfigDefault = Config{

View File

@ -37,39 +37,39 @@ func New(config ...Config) fiber.Handler {
path := c.Path()
// We are only interested in /debug/pprof routes
if len(path) < 12 || !strings.HasPrefix(path, "/debug/pprof") {
if len(path) < 12 || !strings.HasPrefix(path, cfg.Prefix+"/debug/pprof") {
return c.Next()
}
// Switch to original path without stripped slashes
switch path {
case "/debug/pprof/":
case cfg.Prefix + "/debug/pprof/":
pprofIndex(c.Context())
case "/debug/pprof/cmdline":
case cfg.Prefix + "/debug/pprof/cmdline":
pprofCmdline(c.Context())
case "/debug/pprof/profile":
case cfg.Prefix + "/debug/pprof/profile":
pprofProfile(c.Context())
case "/debug/pprof/symbol":
case cfg.Prefix + "/debug/pprof/symbol":
pprofSymbol(c.Context())
case "/debug/pprof/trace":
case cfg.Prefix + "/debug/pprof/trace":
pprofTrace(c.Context())
case "/debug/pprof/allocs":
case cfg.Prefix + "/debug/pprof/allocs":
pprofAllocs(c.Context())
case "/debug/pprof/block":
case cfg.Prefix + "/debug/pprof/block":
pprofBlock(c.Context())
case "/debug/pprof/goroutine":
case cfg.Prefix + "/debug/pprof/goroutine":
pprofGoroutine(c.Context())
case "/debug/pprof/heap":
case cfg.Prefix + "/debug/pprof/heap":
pprofHeap(c.Context())
case "/debug/pprof/mutex":
case cfg.Prefix + "/debug/pprof/mutex":
pprofMutex(c.Context())
case "/debug/pprof/threadcreate":
case cfg.Prefix + "/debug/pprof/threadcreate":
pprofThreadcreate(c.Context())
default:
// pprof index only works with trailing slash
if strings.HasSuffix(path, "/") {
path = strings.TrimRight(path, "/")
} else {
path = "/debug/pprof/"
path = cfg.Prefix + "/debug/pprof/"
}
return c.Redirect(path, fiber.StatusFound)

View File

@ -28,6 +28,24 @@ func Test_Non_Pprof_Path(t *testing.T) {
utils.AssertEqual(t, "escaped", string(b))
}
func Test_Non_Pprof_Path_WithPrefix(t *testing.T) {
app := fiber.New(fiber.Config{DisableStartupMessage: true})
app.Use(New(Config{Prefix: "/federated-fiber"}))
app.Get("/", func(c *fiber.Ctx) error {
return c.SendString("escaped")
})
resp, err := app.Test(httptest.NewRequest(fiber.MethodGet, "/", nil))
utils.AssertEqual(t, nil, err)
utils.AssertEqual(t, 200, resp.StatusCode)
b, err := io.ReadAll(resp.Body)
utils.AssertEqual(t, nil, err)
utils.AssertEqual(t, "escaped", string(b))
}
func Test_Pprof_Index(t *testing.T) {
app := fiber.New(fiber.Config{DisableStartupMessage: true})
@ -47,6 +65,25 @@ func Test_Pprof_Index(t *testing.T) {
utils.AssertEqual(t, true, bytes.Contains(b, []byte("<title>/debug/pprof/</title>")))
}
func Test_Pprof_Index_WithPrefix(t *testing.T) {
app := fiber.New(fiber.Config{DisableStartupMessage: true})
app.Use(New(Config{Prefix: "/federated-fiber"}))
app.Get("/", func(c *fiber.Ctx) error {
return c.SendString("escaped")
})
resp, err := app.Test(httptest.NewRequest(fiber.MethodGet, "/federated-fiber/debug/pprof/", nil))
utils.AssertEqual(t, nil, err)
utils.AssertEqual(t, 200, resp.StatusCode)
utils.AssertEqual(t, fiber.MIMETextHTMLCharsetUTF8, resp.Header.Get(fiber.HeaderContentType))
b, err := io.ReadAll(resp.Body)
utils.AssertEqual(t, nil, err)
utils.AssertEqual(t, true, bytes.Contains(b, []byte("<title>/debug/pprof/</title>")))
}
func Test_Pprof_Subs(t *testing.T) {
app := fiber.New(fiber.Config{DisableStartupMessage: true})
@ -74,6 +111,33 @@ func Test_Pprof_Subs(t *testing.T) {
}
}
func Test_Pprof_Subs_WithPrefix(t *testing.T) {
app := fiber.New(fiber.Config{DisableStartupMessage: true})
app.Use(New(Config{Prefix: "/federated-fiber"}))
app.Get("/", func(c *fiber.Ctx) error {
return c.SendString("escaped")
})
subs := []string{
"cmdline", "profile", "symbol", "trace", "allocs", "block",
"goroutine", "heap", "mutex", "threadcreate",
}
for _, sub := range subs {
t.Run(sub, func(t *testing.T) {
target := "/federated-fiber/debug/pprof/" + sub
if sub == "profile" {
target += "?seconds=1"
}
resp, err := app.Test(httptest.NewRequest(fiber.MethodGet, target, nil), 5000)
utils.AssertEqual(t, nil, err)
utils.AssertEqual(t, 200, resp.StatusCode)
})
}
}
func Test_Pprof_Other(t *testing.T) {
app := fiber.New(fiber.Config{DisableStartupMessage: true})
@ -88,6 +152,20 @@ func Test_Pprof_Other(t *testing.T) {
utils.AssertEqual(t, 302, resp.StatusCode)
}
func Test_Pprof_Other_WithPrefix(t *testing.T) {
app := fiber.New(fiber.Config{DisableStartupMessage: true})
app.Use(New(Config{Prefix: "/federated-fiber"}))
app.Get("/", func(c *fiber.Ctx) error {
return c.SendString("escaped")
})
resp, err := app.Test(httptest.NewRequest(fiber.MethodGet, "/federated-fiber/debug/pprof/302", nil))
utils.AssertEqual(t, nil, err)
utils.AssertEqual(t, 302, resp.StatusCode)
}
// go test -run Test_Pprof_Next
func Test_Pprof_Next(t *testing.T) {
t.Parallel()
@ -104,3 +182,21 @@ func Test_Pprof_Next(t *testing.T) {
utils.AssertEqual(t, nil, err)
utils.AssertEqual(t, 404, resp.StatusCode)
}
// go test -run Test_Pprof_Next_WithPrefix
func Test_Pprof_Next_WithPrefix(t *testing.T) {
t.Parallel()
app := fiber.New()
app.Use(New(Config{
Next: func(_ *fiber.Ctx) bool {
return true
},
Prefix: "/federated-fiber",
}))
resp, err := app.Test(httptest.NewRequest(fiber.MethodGet, "/federated-fiber/debug/pprof/", nil))
utils.AssertEqual(t, nil, err)
utils.AssertEqual(t, 404, resp.StatusCode)
}