From 751553a84552308441784b867b665c1ffff13c46 Mon Sep 17 00:00:00 2001
From: thylong <theotime@protonmail.com>
Date: Sat, 5 Mar 2022 17:00:03 +0100
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Cache=20middleware:=20Add=20flag=20?=
 =?UTF-8?q?for=20e2e=20headers.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Set flag to prevent e2e headers caching to
be the default behavior of the cache middleware.
This would otherwise change quite a lot the
experience for cache middleware current users.
---
 middleware/cache/cache.go      | 18 ++++++++++--------
 middleware/cache/cache_test.go |  4 +++-
 middleware/cache/config.go     |  6 ++++++
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/middleware/cache/cache.go b/middleware/cache/cache.go
index d8f7130365f..11ea341c783 100644
--- a/middleware/cache/cache.go
+++ b/middleware/cache/cache.go
@@ -153,14 +153,16 @@ func New(config ...Config) fiber.Handler {
 
 		// Store all end-to-end headers
 		// (more: https://datatracker.ietf.org/doc/html/rfc2616#section-13.5.1)
-		c.Response().Header.VisitAll(
-			func(key []byte, value []byte) {
-				keyS := string(key)
-				if _, isHopbyHop := ignoreHeaders[keyS]; !isHopbyHop {
-					e.e2eHeaders[keyS] = value
-				}
-			},
-		)
+		if cfg.E2EHeaders {
+			c.Response().Header.VisitAll(
+				func(key []byte, value []byte) {
+					keyS := string(key)
+					if _, isHopbyHop := ignoreHeaders[keyS]; !isHopbyHop {
+						e.e2eHeaders[keyS] = value
+					}
+				},
+			)
+		}
 
 		// default cache expiration
 		expiration := uint64(cfg.Expiration.Seconds())
diff --git a/middleware/cache/cache_test.go b/middleware/cache/cache_test.go
index 2d6027d2e33..59714b3b398 100644
--- a/middleware/cache/cache_test.go
+++ b/middleware/cache/cache_test.go
@@ -304,7 +304,9 @@ func Test_CustomExpiration(t *testing.T) {
 
 func Test_AdditionalE2EResponseHeaders(t *testing.T) {
 	app := fiber.New()
-	app.Use(New())
+	app.Use(New(Config{
+		E2EHeaders: true,
+	}))
 
 	app.Get("/", func(c *fiber.Ctx) error {
 		c.Response().Header.Add("X-Foobar", "foobar")
diff --git a/middleware/cache/config.go b/middleware/cache/config.go
index cd25c4ade5a..12f78040bba 100644
--- a/middleware/cache/config.go
+++ b/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
+	E2EHeaders bool
 }
 
 // ConfigDefault is the default config
@@ -66,6 +71,7 @@ var ConfigDefault = Config{
 		return utils.CopyString(c.Path())
 	},
 	ExpirationGenerator: nil,
+	E2EHeaders:          false,
 	Storage:             nil,
 }