compress.go 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125
  1. package compress
  2. import (
  3. "strings"
  4. "github.com/gofiber/fiber/v3"
  5. "github.com/gofiber/fiber/v3/middleware/etag"
  6. "github.com/gofiber/utils/v2"
  7. "github.com/valyala/fasthttp"
  8. )
  9. func hasToken(header, token string) bool {
  10. for part := range strings.SplitSeq(header, ",") {
  11. if utils.EqualFold(utils.TrimSpace(part), token) {
  12. return true
  13. }
  14. }
  15. return false
  16. }
  17. func shouldSkip(c fiber.Ctx) bool {
  18. if c.Method() == fiber.MethodHead {
  19. return true
  20. }
  21. status := c.Response().StatusCode()
  22. if status < 200 ||
  23. status == fiber.StatusNoContent ||
  24. status == fiber.StatusResetContent ||
  25. status == fiber.StatusNotModified ||
  26. status == fiber.StatusPartialContent ||
  27. len(c.Response().Body()) == 0 ||
  28. c.Get(fiber.HeaderRange) != "" ||
  29. hasToken(c.Get(fiber.HeaderCacheControl), "no-transform") ||
  30. hasToken(c.GetRespHeader(fiber.HeaderCacheControl), "no-transform") {
  31. return true
  32. }
  33. return false
  34. }
  35. func appendVaryAcceptEncoding(c fiber.Ctx) {
  36. vary := c.GetRespHeader(fiber.HeaderVary)
  37. if vary == "" {
  38. c.Set(fiber.HeaderVary, fiber.HeaderAcceptEncoding)
  39. return
  40. }
  41. if hasToken(vary, "*") || hasToken(vary, fiber.HeaderAcceptEncoding) {
  42. return
  43. }
  44. c.Set(fiber.HeaderVary, vary+", "+fiber.HeaderAcceptEncoding)
  45. }
  46. // New creates a new middleware handler
  47. func New(config ...Config) fiber.Handler {
  48. // Set default config
  49. cfg := configDefault(config...)
  50. // Setup request handlers
  51. var (
  52. fctx = func(_ *fasthttp.RequestCtx) {}
  53. compressor fasthttp.RequestHandler
  54. )
  55. // Setup compression algorithm
  56. switch cfg.Level {
  57. case LevelDefault:
  58. // LevelDefault
  59. compressor = fasthttp.CompressHandlerBrotliLevel(fctx,
  60. fasthttp.CompressBrotliDefaultCompression,
  61. fasthttp.CompressDefaultCompression,
  62. )
  63. case LevelBestSpeed:
  64. // LevelBestSpeed
  65. compressor = fasthttp.CompressHandlerBrotliLevel(fctx,
  66. fasthttp.CompressBrotliBestSpeed,
  67. fasthttp.CompressBestSpeed,
  68. )
  69. case LevelBestCompression:
  70. // LevelBestCompression
  71. compressor = fasthttp.CompressHandlerBrotliLevel(fctx,
  72. fasthttp.CompressBrotliBestCompression,
  73. fasthttp.CompressBestCompression,
  74. )
  75. default:
  76. // LevelDisabled
  77. return func(c fiber.Ctx) error {
  78. return c.Next()
  79. }
  80. }
  81. // Return new handler
  82. return func(c fiber.Ctx) error {
  83. // Don't execute middleware if Next returns true
  84. if cfg.Next != nil && cfg.Next(c) {
  85. return c.Next()
  86. }
  87. // Continue stack
  88. if err := c.Next(); err != nil {
  89. return err
  90. }
  91. if shouldSkip(c) {
  92. appendVaryAcceptEncoding(c)
  93. return nil
  94. }
  95. if c.GetRespHeader(fiber.HeaderContentEncoding) != "" {
  96. appendVaryAcceptEncoding(c)
  97. return nil
  98. }
  99. compressor(c.RequestCtx())
  100. if tag := c.GetRespHeader(fiber.HeaderETag); tag != "" && !strings.HasPrefix(tag, "W/") {
  101. if c.GetRespHeader(fiber.HeaderContentEncoding) != "" {
  102. c.Set(fiber.HeaderETag, string(etag.Generate(c.Response().Body())))
  103. }
  104. }
  105. appendVaryAcceptEncoding(c)
  106. return nil
  107. }
  108. }