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

@@ -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,
 }
-```
+```

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{

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)

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("<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)
+}