request_id.go 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. package middleware
  2. import (
  3. "github.com/labstack/echo"
  4. "github.com/labstack/gommon/random"
  5. )
  6. type (
  7. // RequestIDConfig defines the config for RequestID middleware.
  8. RequestIDConfig struct {
  9. // Skipper defines a function to skip middleware.
  10. Skipper Skipper
  11. // Generator defines a function to generate an ID.
  12. // Optional. Default value random.String(32).
  13. Generator func() string
  14. }
  15. )
  16. var (
  17. // DefaultRequestIDConfig is the default RequestID middleware config.
  18. DefaultRequestIDConfig = RequestIDConfig{
  19. Skipper: DefaultSkipper,
  20. Generator: generator,
  21. }
  22. )
  23. // RequestID returns a X-Request-ID middleware.
  24. func RequestID() echo.MiddlewareFunc {
  25. return RequestIDWithConfig(DefaultRequestIDConfig)
  26. }
  27. // RequestIDWithConfig returns a X-Request-ID middleware with config.
  28. func RequestIDWithConfig(config RequestIDConfig) echo.MiddlewareFunc {
  29. // Defaults
  30. if config.Skipper == nil {
  31. config.Skipper = DefaultRequestIDConfig.Skipper
  32. }
  33. if config.Generator == nil {
  34. config.Generator = generator
  35. }
  36. return func(next echo.HandlerFunc) echo.HandlerFunc {
  37. return func(c echo.Context) error {
  38. if config.Skipper(c) {
  39. return next(c)
  40. }
  41. req := c.Request()
  42. res := c.Response()
  43. rid := req.Header.Get(echo.HeaderXRequestID)
  44. if rid == "" {
  45. rid = config.Generator()
  46. }
  47. res.Header().Set(echo.HeaderXRequestID, rid)
  48. return next(c)
  49. }
  50. }
  51. }
  52. func generator() string {
  53. return random.String(32)
  54. }