feat: add Option DiscardHeaders and CorsHeaders

chenyahui · Nov 16, 2023 · a2f9e2b · a2f9e2b
1 parent b885ae3
commit a2f9e2b
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 4 deletions.
diff --git a/cache.go b/cache.go
@@ -113,7 +113,7 @@ func cache(
 			inFlight = true
 
 			respCache := &ResponseCache{}
-			respCache.fillWithCacheWriter(cacheWriter, cfg.withoutHeader)
+			respCache.fillWithCacheWriter(cacheWriter, cfg)
 
 			// only cache 2xx response
 			if !c.IsAborted() && cacheWriter.Status() < 300 && cacheWriter.Status() >= 200 {
@@ -214,11 +214,15 @@ type ResponseCache struct {
 	Data   []byte
 }
 
-func (c *ResponseCache) fillWithCacheWriter(cacheWriter *responseCacheWriter, withoutHeader bool) {
+func (c *ResponseCache) fillWithCacheWriter(cacheWriter *responseCacheWriter, cfg *Config) {
 	c.Status = cacheWriter.Status()
 	c.Data = cacheWriter.body.Bytes()
-	if !withoutHeader {
+	if !cfg.withoutHeader {
 		c.Header = cacheWriter.Header().Clone()
+
+		for _, headerKey := range cfg.discardHeaders {
+			c.Header.Del(headerKey)
+		}
 	}
 }
 

diff --git a/option.go b/option.go
@@ -23,6 +23,7 @@ type Config struct {
 	ignoreQueryOrder bool
 	prefixKey        string
 	withoutHeader    bool
+	discardHeaders   []string
 }
 
 func newConfigByOpts(opts ...Option) *Config {
@@ -158,4 +159,19 @@ func WithoutHeader() Option {
 	return func(c *Config) {
 		c.withoutHeader = true
 	}
-}
+}
+
+func DiscardHeaders(headers []string) Option {
+	return func(c *Config) {
+		c.discardHeaders = headers
+	}
+}
+
+func CorsHeaders() []string {
+	return []string{
+		"Access-Control-Allow-Credentials",
+		"Access-Control-Expose-Headers",
+		"Access-Control-Allow-Origin",
+		"Vary",
+	}
+}