From 35591da3bb29249fc149e1e6c484b7f03d65e5a8 Mon Sep 17 00:00:00 2001
From: vicanso <tree.xie@outlook.com>
Date: Tue, 1 Aug 2023 20:25:44 +0800
Subject: [PATCH] feat: supports ignore headers for http cache

---
 go.mod              |  4 ++--
 go.sum              |  8 ++++----
 middleware/cache.go | 10 +++++++---
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/go.mod b/go.mod
index ca62212..cf199d0 100644
--- a/go.mod
+++ b/go.mod
@@ -4,8 +4,8 @@ go 1.17
 
 require (
 	github.com/andybalholm/brotli v1.0.5
-	github.com/stretchr/testify v1.8.2
-	github.com/tidwall/gjson v1.14.4
+	github.com/stretchr/testify v1.8.4
+	github.com/tidwall/gjson v1.15.0
 	github.com/vicanso/hes v0.6.2
 	github.com/vicanso/intranet-ip v0.1.0
 	github.com/vicanso/keygrip v1.2.1
diff --git a/go.sum b/go.sum
index 0c77d7e..485b743 100644
--- a/go.sum
+++ b/go.sum
@@ -13,10 +13,10 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
-github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
-github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
-github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM=
-github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+github.com/tidwall/gjson v1.15.0 h1:5n/pM+v3r5ujuNl4YLZLsQ+UE5jlkLVm7jMzT5Mpolw=
+github.com/tidwall/gjson v1.15.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
 github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA=
 github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
 github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs=
diff --git a/middleware/cache.go b/middleware/cache.go
index f95472a..6619794 100644
--- a/middleware/cache.go
+++ b/middleware/cache.go
@@ -50,6 +50,8 @@ type CacheConfig struct {
 	// Marshal marshal function for cache, if BodyBuffer is nil,
 	// the body will be marshaled to body buffer. The default marshal function will be json.Marshal
 	Marshal func(interface{}) ([]byte, error)
+	// IgnoreHeaders ignore the headers for cache
+	IgnoreHeaders []string
 }
 
 type CacheStatus uint8
@@ -71,7 +73,7 @@ type CacheStore interface {
 const HeaderAge = "Age"
 const HeaderXCache = "X-Cache"
 
-var ignoreHeaders = []string{
+var defaultIgnoreHeaders = []string{
 	"Content-Encoding",
 	"Content-Length",
 	"Connection",
@@ -193,13 +195,14 @@ var hitForPassData = (&CacheResponse{
 }).Bytes()
 
 // Bytes converts the cache response to bytes
-func (cp *CacheResponse) Bytes() []byte {
+func (cp *CacheResponse) Bytes(ignoreHeaders ...string) []byte {
 	// 只有hit的才需要保存后续的数据
 	if cp.Status != StatusHit {
 		return []byte{
 			byte(cp.Status),
 		}
 	}
+	ignoreHeaders = append(ignoreHeaders, defaultIgnoreHeaders...)
 	headers := NewHTTPHeaders(cp.Header, ignoreHeaders...)
 	headersLength := len(headers)
 	// 1个字节保存状态
@@ -369,6 +372,7 @@ func NewCache(config CacheConfig) elton.Handler {
 	if marshal == nil {
 		marshal = json.Marshal
 	}
+	ignoreHeaders := config.IgnoreHeaders
 	compressor := config.Compressor
 	return func(c *elton.Context) error {
 		if skipper(c) {
@@ -435,7 +439,7 @@ func NewCache(config CacheConfig) elton.Handler {
 			Header:      c.Header(),
 			Body:        buffer,
 		}
-		data = cacheResp.Bytes()
+		data = cacheResp.Bytes(ignoreHeaders...)
 		// 如果想忽略store的错误，则自定义store时，
 		// 不要返回出错则可
 		err = store.Set(ctx, key, data, time.Duration(cacheAge)*time.Second)