diff --git a/middleware/requestid/README.md b/middleware/requestid/README.md index 07bdde0b..13f5f053 100644 --- a/middleware/requestid/README.md +++ b/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" } - ``` diff --git a/middleware/requestid/requestid.go b/middleware/requestid/requestid.go index ef0de78a..0a5a11a5 100644 --- a/middleware/requestid/requestid.go +++ b/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() } diff --git a/middleware/requestid/requestid_test.go b/middleware/requestid/requestid_test.go index 5903de2a..c2d211c8 100644 --- a/middleware/requestid/requestid_test.go +++ b/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) +} \ No newline at end of file