✨ Cache middleware: Store e2e headers.

middleware/cache/README.md

@@ -10,6 +10,7 @@ Cache middleware for [Fiber](https://github.com/gofiber/fiber) designed to inter
 	- [Examples](#examples)
 		- [Default Config](#default-config)
 		- [Custom Config](#custom-config)
+		- [Custom Cache Key Or Expiration](#custom-cache-key-or-expiration)
 		- [Config](#config)
 		- [Default Config](#default-config-1)
 
@@ -112,6 +113,11 @@ type Config struct {
 	//
 	// Default: an in memory store for this process only
 	Storage fiber.Storage
+
+	// allows you to store additional headers generated by next middlewares & handler
+	//
+	// Default: false
+	StoreResponseHeaders bool
 }
 ```
 
@@ -128,6 +134,7 @@ var ConfigDefault = Config{
 		return utils.CopyString(c.Path())
 	},
 	ExpirationGenerator : nil,
+	StoreResponseHeaders: false,
 	Storage:      nil,
 }
 ```

middleware/cache/cache.go

@@ -27,6 +27,19 @@ const (
 	cacheMiss        = "miss"
 )
 
+var ignoreHeaders = map[string]interface{}{
+	"Connection":          nil,
+	"Keep-Alive":          nil,
+	"Proxy-Authenticate":  nil,
+	"Proxy-Authorization": nil,
+	"TE":                  nil,
+	"Trailers":            nil,
+	"Transfer-Encoding":   nil,
+	"Upgrade":             nil,
+	"Content-Type":        nil, // already stored explicitely by the cache manager
+	"Content-Encoding":    nil, // already stored explicitely by the cache manager
+}
+
 // New creates a new middleware handler
 func New(config ...Config) fiber.Handler {
 	// Set default config
@@ -96,6 +109,9 @@ func New(config ...Config) fiber.Handler {
 			if len(e.cencoding) > 0 {
 				c.Response().Header.SetBytesV(fiber.HeaderContentEncoding, e.cencoding)
 			}
+			for k, v := range e.headers {
+				c.Response().Header.AddBytesV(k, v)
+			}
 			// Set Cache-Control header if enabled
 			if cfg.CacheControl {
 				maxAge := strconv.FormatUint(e.exp-ts, 10)
@@ -133,6 +149,20 @@ func New(config ...Config) fiber.Handler {
 		e.status = c.Response().StatusCode()
 		e.ctype = utils.CopyBytes(c.Response().Header.ContentType())
 		e.cencoding = utils.CopyBytes(c.Response().Header.Peek(fiber.HeaderContentEncoding))
+		e.headers = make(map[string][]byte)
+
+		// Store all response headers
+		// (more: https://datatracker.ietf.org/doc/html/rfc2616#section-13.5.1)
+		if cfg.StoreResponseHeaders {
+			c.Response().Header.VisitAll(
+				func(key []byte, value []byte) {
+					keyS := string(key)
+					if _, isHopbyHop := ignoreHeaders[keyS]; !isHopbyHop {
+						e.headers[keyS] = value
+					}
+				},
+			)
+		}
 
 		// default cache expiration
 		expiration := uint64(cfg.Expiration.Seconds())

middleware/cache/cache_test.go

@@ -302,6 +302,28 @@ func Test_CustomExpiration(t *testing.T) {
 	utils.AssertEqual(t, 6000, newCacheTime)
 }
 
+func Test_AdditionalE2EResponseHeaders(t *testing.T) {
+	app := fiber.New()
+	app.Use(New(Config{
+		StoreResponseHeaders: true,
+	}))
+
+	app.Get("/", func(c *fiber.Ctx) error {
+		c.Response().Header.Add("X-Foobar", "foobar")
+		return c.SendString("hi")
+	})
+
+	req := httptest.NewRequest("GET", "/", nil)
+	resp, err := app.Test(req)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, "foobar", resp.Header.Get("X-Foobar"))
+
+	req = httptest.NewRequest("GET", "/", nil)
+	resp, err = app.Test(req)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, "foobar", resp.Header.Get("X-Foobar"))
+}
+
 func Test_CacheHeader(t *testing.T) {
 	app := fiber.New()
 
@@ -475,3 +497,31 @@ func Benchmark_Cache_Storage(b *testing.B) {
 	utils.AssertEqual(b, fiber.StatusTeapot, fctx.Response.Header.StatusCode())
 	utils.AssertEqual(b, true, len(fctx.Response.Body()) > 30000)
 }
+
+func Benchmark_Cache_AdditionalHeaders(b *testing.B) {
+	app := fiber.New()
+	app.Use(New(Config{
+		StoreResponseHeaders: true,
+	}))
+
+	app.Get("/demo", func(c *fiber.Ctx) error {
+		c.Response().Header.Add("X-Foobar", "foobar")
+		return c.SendStatus(418)
+	})
+
+	h := app.Handler()
+
+	fctx := &fasthttp.RequestCtx{}
+	fctx.Request.Header.SetMethod("GET")
+	fctx.Request.SetRequestURI("/demo")
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for n := 0; n < b.N; n++ {
+		h(fctx)
+	}
+
+	utils.AssertEqual(b, fiber.StatusTeapot, fctx.Response.Header.StatusCode())
+	utils.AssertEqual(b, []byte("foobar"), fctx.Response.Header.Peek("X-Foobar"))
+}

middleware/cache/config.go

@@ -54,6 +54,11 @@ type Config struct {
 
 	// Deprecated, use KeyGenerator instead
 	Key func(*fiber.Ctx) string
+
+	// allows you to store additional headers generated by next middlewares & handler
+	//
+	// Default: false
+	StoreResponseHeaders bool
 }
 
 // ConfigDefault is the default config
@@ -65,8 +70,9 @@ var ConfigDefault = Config{
 	KeyGenerator: func(c *fiber.Ctx) string {
 		return utils.CopyString(c.Path())
 	},
-	ExpirationGenerator: nil,
-	Storage:             nil,
+	ExpirationGenerator:  nil,
+	StoreResponseHeaders: false,
+	Storage:              nil,
 }
 
 // Helper function to set default values

middleware/cache/manager.go

@@ -18,6 +18,7 @@ type item struct {
 	cencoding []byte
 	status    int
 	exp       uint64
+	headers   map[string][]byte
 }
 
 //msgp:ignore manager
@@ -61,6 +62,7 @@ func (m *manager) release(e *item) {
 	e.ctype = nil
 	e.status = 0
 	e.exp = 0
+	e.headers = nil
 	m.pool.Put(e)
 }