2017-02-10 16:21:32 +09:00
|
|
|
package middleware
|
|
|
|
|
|
|
|
import (
|
2021-07-15 23:34:01 +03:00
|
|
|
"github.com/labstack/echo/v5"
|
2017-02-10 16:21:32 +09:00
|
|
|
)
|
|
|
|
|
2021-07-15 23:34:01 +03:00
|
|
|
// RequestIDConfig defines the config for RequestID middleware.
|
|
|
|
type RequestIDConfig struct {
|
|
|
|
// Skipper defines a function to skip middleware.
|
|
|
|
Skipper Skipper
|
2017-03-06 13:16:05 -08:00
|
|
|
|
2021-07-15 23:34:01 +03:00
|
|
|
// Generator defines a function to generate an ID.
|
|
|
|
// Optional. Default value random.String(32).
|
|
|
|
Generator func() string
|
2021-06-23 01:04:23 +09:00
|
|
|
|
2021-07-15 23:34:01 +03:00
|
|
|
// RequestIDHandler defines a function which is executed for a request id.
|
|
|
|
RequestIDHandler func(c echo.Context, requestID string)
|
2021-12-07 10:56:32 +00:00
|
|
|
|
2021-07-15 23:34:01 +03:00
|
|
|
// TargetHeader defines what header to look for to populate the id
|
|
|
|
TargetHeader string
|
|
|
|
}
|
2017-02-10 16:21:32 +09:00
|
|
|
|
|
|
|
// RequestID returns a X-Request-ID middleware.
|
|
|
|
func RequestID() echo.MiddlewareFunc {
|
2021-07-15 23:34:01 +03:00
|
|
|
return RequestIDWithConfig(RequestIDConfig{})
|
2017-02-10 16:21:32 +09:00
|
|
|
}
|
|
|
|
|
2021-07-15 23:34:01 +03:00
|
|
|
// RequestIDWithConfig returns a X-Request-ID middleware with config or panics on invalid configuration.
|
2017-02-10 16:21:32 +09:00
|
|
|
func RequestIDWithConfig(config RequestIDConfig) echo.MiddlewareFunc {
|
2021-07-15 23:34:01 +03:00
|
|
|
return toMiddlewareOrPanic(config)
|
|
|
|
}
|
|
|
|
|
|
|
|
// ToMiddleware converts RequestIDConfig to middleware or returns an error for invalid configuration
|
|
|
|
func (config RequestIDConfig) ToMiddleware() (echo.MiddlewareFunc, error) {
|
2017-02-10 16:21:32 +09:00
|
|
|
if config.Skipper == nil {
|
2021-07-15 23:34:01 +03:00
|
|
|
config.Skipper = DefaultSkipper
|
2017-02-10 16:21:32 +09:00
|
|
|
}
|
2017-03-06 13:16:05 -08:00
|
|
|
if config.Generator == nil {
|
2021-07-15 23:34:01 +03:00
|
|
|
config.Generator = createRandomStringGenerator(32)
|
2017-03-06 13:16:05 -08:00
|
|
|
}
|
2021-12-07 10:56:32 +00:00
|
|
|
if config.TargetHeader == "" {
|
|
|
|
config.TargetHeader = echo.HeaderXRequestID
|
|
|
|
}
|
2017-02-10 16:21:32 +09:00
|
|
|
|
|
|
|
return func(next echo.HandlerFunc) echo.HandlerFunc {
|
|
|
|
return func(c echo.Context) error {
|
|
|
|
if config.Skipper(c) {
|
|
|
|
return next(c)
|
|
|
|
}
|
|
|
|
|
|
|
|
req := c.Request()
|
|
|
|
res := c.Response()
|
2021-12-07 10:56:32 +00:00
|
|
|
rid := req.Header.Get(config.TargetHeader)
|
2017-02-10 16:21:32 +09:00
|
|
|
if rid == "" {
|
2017-03-18 01:22:58 +09:00
|
|
|
rid = config.Generator()
|
2017-02-10 16:21:32 +09:00
|
|
|
}
|
2021-12-07 10:56:32 +00:00
|
|
|
res.Header().Set(config.TargetHeader, rid)
|
2021-06-23 01:04:23 +09:00
|
|
|
if config.RequestIDHandler != nil {
|
|
|
|
config.RequestIDHandler(c, rid)
|
|
|
|
}
|
2017-02-10 16:21:32 +09:00
|
|
|
|
|
|
|
return next(c)
|
|
|
|
}
|
2021-07-15 23:34:01 +03:00
|
|
|
}, nil
|
2017-03-06 13:16:05 -08:00
|
|
|
}
|