Merge pull request #948 from codemicro/requestid-locals

🔥 Add request ID to locals

middleware/requestid/README.md

@@ -52,21 +52,25 @@ type Config struct {
 
 	// Generator defines a function to generate the unique identifier.
 	//
-	// Optional. Default: func() string {
-	//   return utils.UUID()
-	// }
+	// Optional. Default: utils.UUID
 	Generator func() string
+
+	// ContextKey defines the key used when storing the request ID in
+	// the locals for a specific request.
+	//
+	// Optional. Default: requestid
+	ContextKey string
 }
 ```
 
 ### Default Config
 ```go
 var ConfigDefault = Config{
-	Next:      nil,
-	Header:    fiber.HeaderXRequestID,
-	Generator: func() string {
+	Next:       nil,
+	Header:     fiber.HeaderXRequestID,
+	Generator:  func() string {
 		return utils.UUID()
 	},
+	ContextKey: "requestid"
 }
-
 ```

middleware/requestid/requestid.go

@@ -21,6 +21,12 @@ type Config struct {
 	//
 	// Optional. Default: utils.UUID
 	Generator func() string
+
+	// ContextKey defines the key used when storing the request ID in
+	// the locals for a specific request.
+	//
+	// Optional. Default: requestid
+	ContextKey string
 }
 
 // ConfigDefault is the default config
@@ -28,6 +34,7 @@ var ConfigDefault = Config{
 	Next:      nil,
 	Header:    fiber.HeaderXRequestID,
 	Generator: utils.UUID,
+	ContextKey: "requestid",
 }
 
 // New creates a new middleware handler
@@ -46,6 +53,9 @@ func New(config ...Config) fiber.Handler {
 		if cfg.Generator == nil {
 			cfg.Generator = ConfigDefault.Generator
 		}
+		if cfg.ContextKey == "" {
+			cfg.ContextKey = ConfigDefault.ContextKey
+		}
 	}
 
 	// Return new handler
@@ -60,6 +70,9 @@ func New(config ...Config) fiber.Handler {
 		// Set new id to response header
 		c.Set(cfg.Header, rid)
 
+		// Add the request ID to locals
+		c.Locals(cfg.ContextKey, rid)
+
 		// Continue stack
 		return c.Next()
 	}

middleware/requestid/requestid_test.go

@@ -48,3 +48,28 @@ func Test_RequestID_Next(t *testing.T) {
 	utils.AssertEqual(t, resp.Header.Get(fiber.HeaderXRequestID), "")
 	utils.AssertEqual(t, fiber.StatusNotFound, resp.StatusCode)
 }
+
+// go test -run Test_RequestID_Locals
+func Test_RequestID_Locals(t *testing.T) {
+	reqId := "ThisIsARequestId"
+	ctxKey := "ThisIsAContextKey"
+
+	app := fiber.New()
+	app.Use(New(Config{
+		Generator: func() string {
+			return reqId
+		},
+		ContextKey: ctxKey,
+	}))
+
+	var ctxVal string
+
+	app.Use(func (c *fiber.Ctx) error {
+		ctxVal = c.Locals(ctxKey).(string)
+		return c.Next()
+	})
+
+	_, err := app.Test(httptest.NewRequest("GET", "/", nil))
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, reqId, ctxVal)
+}