Skip to content

Commit

Permalink
Add support for configurable target header for the request_id middleware
Browse files Browse the repository at this point in the history
  • Loading branch information
Guilherme Cardoso authored and aldas committed Dec 7, 2021
1 parent b437ee3 commit c32fafa
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 4 deletions.
1 change: 1 addition & 0 deletions echo.go
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,7 @@ const (
HeaderXHTTPMethodOverride = "X-HTTP-Method-Override"
HeaderXRealIP = "X-Real-IP"
HeaderXRequestID = "X-Request-ID"
HeaderXCorrelationID = "X-Correlation-ID"
HeaderXRequestedWith = "X-Requested-With"
HeaderServer = "Server"
HeaderOrigin = "Origin"
Expand Down
15 changes: 11 additions & 4 deletions middleware/request_id.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,18 @@ type (

// RequestIDHandler defines a function which is executed for a request id.
RequestIDHandler func(echo.Context, string)

// TargetHeader defines what header to look for to populate the id
TargetHeader string
}
)

var (
// DefaultRequestIDConfig is the default RequestID middleware config.
DefaultRequestIDConfig = RequestIDConfig{
Skipper: DefaultSkipper,
Generator: generator,
Skipper: DefaultSkipper,
Generator: generator,
TargetHeader: echo.HeaderXRequestID,
}
)

Expand All @@ -42,6 +46,9 @@ func RequestIDWithConfig(config RequestIDConfig) echo.MiddlewareFunc {
if config.Generator == nil {
config.Generator = generator
}
if config.TargetHeader == "" {
config.TargetHeader = echo.HeaderXRequestID
}

return func(next echo.HandlerFunc) echo.HandlerFunc {
return func(c echo.Context) error {
Expand All @@ -51,11 +58,11 @@ func RequestIDWithConfig(config RequestIDConfig) echo.MiddlewareFunc {

req := c.Request()
res := c.Response()
rid := req.Header.Get(echo.HeaderXRequestID)
rid := req.Header.Get(config.TargetHeader)
if rid == "" {
rid = config.Generator()
}
res.Header().Set(echo.HeaderXRequestID, rid)
res.Header().Set(config.TargetHeader, rid)
if config.RequestIDHandler != nil {
config.RequestIDHandler(c, rid)
}
Expand Down
31 changes: 31 additions & 0 deletions middleware/request_id_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -55,3 +55,34 @@ func TestRequestID_IDNotAltered(t *testing.T) {
_ = h(c)
assert.Equal(t, rec.Header().Get(echo.HeaderXRequestID), "<sample-request-id>")
}

func TestRequestIDConfigDifferentHeader(t *testing.T) {
e := echo.New()
req := httptest.NewRequest(http.MethodGet, "/", nil)
rec := httptest.NewRecorder()
c := e.NewContext(req, rec)
handler := func(c echo.Context) error {
return c.String(http.StatusOK, "test")
}

rid := RequestIDWithConfig(RequestIDConfig{TargetHeader: echo.HeaderXCorrelationID})
h := rid(handler)
h(c)
assert.Len(t, rec.Header().Get(echo.HeaderXCorrelationID), 32)

// Custom generator and handler
customID := "customGenerator"
calledHandler := false
rid = RequestIDWithConfig(RequestIDConfig{
Generator: func() string { return customID },
TargetHeader: echo.HeaderXCorrelationID,
RequestIDHandler: func(_ echo.Context, id string) {
calledHandler = true
assert.Equal(t, customID, id)
},
})
h = rid(handler)
h(c)
assert.Equal(t, rec.Header().Get(echo.HeaderXCorrelationID), "customGenerator")
assert.True(t, calledHandler)
}

0 comments on commit c32fafa

Please sign in to comment.