From 13247206ab63de182748c23d277675b574ff7c95 Mon Sep 17 00:00:00 2001 From: Glenn Lewis <6598971+gmlewis@users.noreply.github.com> Date: Tue, 8 Nov 2022 22:59:33 -0800 Subject: [PATCH] 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> --- middleware/pprof/README.md | 19 ++++++- middleware/pprof/config.go | 7 +++ middleware/pprof/pprof.go | 26 ++++----- middleware/pprof/pprof_test.go | 96 ++++++++++++++++++++++++++++++++++ 4 files changed, 134 insertions(+), 14 deletions(-) diff --git a/middleware/pprof/README.md b/middleware/pprof/README.md index 5f266438..2a5bbc2f 100644 --- a/middleware/pprof/README.md +++ b/middleware/pprof/README.md @@ -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, } -``` \ No newline at end of file +``` diff --git a/middleware/pprof/config.go b/middleware/pprof/config.go index 93f58fc9..e69ffd70 100644 --- a/middleware/pprof/config.go +++ b/middleware/pprof/config.go @@ -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{ diff --git a/middleware/pprof/pprof.go b/middleware/pprof/pprof.go index b4c7ba33..8ad85c9a 100644 --- a/middleware/pprof/pprof.go +++ b/middleware/pprof/pprof.go @@ -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) diff --git a/middleware/pprof/pprof_test.go b/middleware/pprof/pprof_test.go index ed0a41b5..09d7206d 100644 --- a/middleware/pprof/pprof_test.go +++ b/middleware/pprof/pprof_test.go @@ -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("/debug/pprof/"))) } +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("/debug/pprof/"))) +} + 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) +}