🔥 feat: Add Context Support to RequestID Middleware (#3200)

* Rename UserContext() to Context(). Rename Context() to RequestCtx()

* feat: add requestID in UserContext

* Update Ctxt docs and What's new

* Remove extra blank lines

* ♻️ Refactor: merge issue #3186

* 🔥 Feature: improve FromContext func and test

* 📚 Doc: improve requestid middleware

* ♻️ Refactor: Rename interface to any

* fix: Modify structure sorting to reduce memory usage

---------

Co-authored-by: Juan Calderon-Perez <jgcalderonperez@protonmail.com>
Co-authored-by: Juan Calderon-Perez <835733+gaby@users.noreply.github.com>

docs/middleware/requestid.md

@@ -49,6 +49,16 @@ func handler(c fiber.Ctx) error {
 }
 ```
 
+In version v3, Fiber will inject `requestID` into the built-in `Context` of Go.
+
+```go
+func handler(c fiber.Ctx) error {
+    id := requestid.FromContext(c.Context())
+    log.Printf("Request ID: %s", id)
+    return c.SendString("Hello, World!")
+}
+```
+
 ## Config
 
 | Property   | Type                    | Description                                                                                       | Default        |

middleware/requestid/requestid.go

@@ -1,7 +1,10 @@
 package requestid
 
 import (
+	"context"
+
 	"github.com/gofiber/fiber/v3"
+	"github.com/gofiber/fiber/v3/log"
 )
 
 // The contextKey type is unexported to prevent collisions with context keys defined in
@@ -36,6 +39,10 @@ func New(config ...Config) fiber.Handler {
 		// Add the request ID to locals
 		c.Locals(requestIDKey, rid)
 
+		// Add the request ID to UserContext
+		ctx := context.WithValue(c.Context(), requestIDKey, rid)
+		c.SetContext(ctx)
+
 		// Continue stack
 		return c.Next()
 	}
@@ -43,9 +50,21 @@ func New(config ...Config) fiber.Handler {
 
 // FromContext returns the request ID from context.
 // If there is no request ID, an empty string is returned.
-func FromContext(c fiber.Ctx) string {
+// Supported context types:
-	if rid, ok := c.Locals(requestIDKey).(string); ok {
+// - fiber.Ctx: Retrieves request ID from Locals
-		return rid
+// - context.Context: Retrieves request ID from context values
+func FromContext(c any) string {
+	switch ctx := c.(type) {
+	case fiber.Ctx:
+		if rid, ok := ctx.Locals(requestIDKey).(string); ok {
+			return rid
+		}
+	case context.Context:
+		if rid, ok := ctx.Value(requestIDKey).(string); ok {
+			return rid
+		}
+	default:
+		log.Errorf("Unsupported context type: %T. Expected fiber.Ctx or context.Context", c)
 	}
 	return ""
 }

middleware/requestid/requestid_test.go

@@ -51,26 +51,59 @@ func Test_RequestID_Next(t *testing.T) {
 	require.Equal(t, fiber.StatusNotFound, resp.StatusCode)
 }
 
-// go test -run Test_RequestID_Locals
+// go test -run Test_RequestID_FromContext
 func Test_RequestID_FromContext(t *testing.T) {
 	t.Parallel()
+
 	reqID := "ThisIsARequestId"
 
-	app := fiber.New()
+	type args struct {
-	app.Use(New(Config{
+		inputFunc func(c fiber.Ctx) any
-		Generator: func() string {
+	}
-			return reqID
+
+	tests := []struct {
+		args args
+		name string
+	}{
+		{
+			name: "From fiber.Ctx",
+			args: args{
+				inputFunc: func(c fiber.Ctx) any {
+					return c
+				},
+			},
 		},
-	}))
+		{
+			name: "From context.Context",
+			args: args{
+				inputFunc: func(c fiber.Ctx) any {
+					return c.Context()
+				},
+			},
+		},
+	}
 
-	var ctxVal string
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
 
-	app.Use(func(c fiber.Ctx) error {
+			app := fiber.New()
-		ctxVal = FromContext(c)
+			app.Use(New(Config{
-		return c.Next()
+				Generator: func() string {
-	})
+					return reqID
+				},
+			}))
 
-	_, err := app.Test(httptest.NewRequest(fiber.MethodGet, "/", nil))
+			var ctxVal string
-	require.NoError(t, err)
+
-	require.Equal(t, reqID, ctxVal)
+			app.Use(func(c fiber.Ctx) error {
+				ctxVal = FromContext(tt.args.inputFunc(c))
+				return c.Next()
+			})
+
+			_, err := app.Test(httptest.NewRequest(fiber.MethodGet, "/", nil))
+			require.NoError(t, err)
+			require.Equal(t, reqID, ctxVal)
+		})
+	}
 }