feat(middleware): Add a MaxBodyBytes configuration option (#442)

* feat(middleware): Add configuration option MaxBodyBytes * fix(configuration): handle max_body_bytes in configuration with fallback on global config for caddy. Add unit test for that * feat(documentation): add caddy documentation * fix(review): rename directive to --------- Co-authored-by: Vincent Jordan <vincent.jordan@platform.sh> Co-authored-by: darkweak <darkweak@protonmail.com>
darkweak · Feb 12, 2024 · 5860dd7 · 5860dd7
1 parent 678a1db
commit 5860dd7
Show file tree

Hide file tree

Showing 12 changed files with 105 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -214,6 +214,7 @@ surrogate_keys:
 | `default_cache.timeout.cache`                     | The timeout duration to consider the cache provider as unreachable                                                                          | `10ms`                                                                                                                    |
 | `default_cache.ttl`                               | The TTL duration                                                                                                                            | `120s`                                                                                                                    |
 | `default_cache.default_cache_control`             | Set the default value of `Cache-Control` response header if not set by upstream (Souin treats empty `Cache-Control` as `public` if omitted) | `no-store`                                                                                                                |
+| `default_cache.max_cachable_body_bytes`                    | Set the maximum size (in bytes) for a response body to be cached (unlimited if omited) | `1048576` (1MB) |
 | `log_level`                                       | The log level                                                                                                                               | `One of DEBUG, INFO, WARN, ERROR, DPANIC, PANIC, FATAL it's case insensitive`                                             |
 | `reverse_proxy_url`                               | The reverse-proxy's instance URL (Apache, Nginx, Træfik...)                                                                                 | - `http://yourservice` (Container way)<br/>`http://localhost:81` (Local way)<br/>`http://yourdomain.com:81` (Network way) |
 | `ssl_providers`                                   | List of your providers handling certificates                                                                                                | `- traefik`<br/><br/>`- nginx`<br/><br/>`- apache`                                                                        |

diff --git a/configurationtypes/types.go b/configurationtypes/types.go
@@ -233,6 +233,7 @@ type DefaultCache struct {
 	Timeout             Timeout       `json:"timeout" yaml:"timeout"`
 	TTL                 Duration      `json:"ttl" yaml:"ttl"`
 	DefaultCacheControl string        `json:"default_cache_control" yaml:"default_cache_control"`
+	MaxBodyBytes        uint64        `json:"max_cachable_body_bytes" yaml:"max_cachable_body_bytes"`
 }
 
 // GetAllowedHTTPVerbs returns the allowed verbs to cache
@@ -325,6 +326,11 @@ func (d *DefaultCache) GetDefaultCacheControl() string {
 	return d.DefaultCacheControl
 }
 
+// GetMaxBodyBytes returns the default maximum body size (in bytes) for storing into cache
+func (d *DefaultCache) GetMaxBodyBytes() uint64 {
+	return d.MaxBodyBytes
+}
+
 // DefaultCacheInterface interface
 type DefaultCacheInterface interface {
 	GetAllowedHTTPVerbs() []string
@@ -345,6 +351,7 @@ type DefaultCacheInterface interface {
 	GetTimeout() Timeout
 	GetTTL() time.Duration
 	GetDefaultCacheControl() string
+	GetMaxBodyBytes() uint64
 }
 
 // APIEndpoint is the minimal structure to define an endpoint

diff --git a/pkg/middleware/middleware.go b/pkg/middleware/middleware.go
@@ -243,6 +243,12 @@ func (s *SouinBaseHandler) Store(
 		if res.Header.Get("Content-Length") == "" {
 			res.Header.Set("Content-Length", fmt.Sprint(bLen))
 		}
+		respBodyMaxSize := int(s.Configuration.GetDefaultCache().GetMaxBodyBytes())
+		if respBodyMaxSize > 0 && bLen > respBodyMaxSize {
+			customWriter.Header().Set("Cache-Status", status+"; detail=UPSTREAM-RESPONSE-TOO-LARGE; key="+rfc.GetCacheKeyFromCtx(rq.Context()))
+
+			return nil
+		}
 		res.Header.Set(rfc.StoredLengthHeader, res.Header.Get("Content-Length"))
 		response, err := httputil.DumpResponse(&res, true)
 		if err == nil && bLen > 0 {

diff --git a/pkg/middleware/writer.go b/pkg/middleware/writer.go
@@ -37,7 +37,6 @@ type CustomWriter struct {
 	headersSent bool
 	mutex       *sync.Mutex
 	statusCode  int
-	// size        int
 }
 
 // Header will write the response headers

diff --git a/plugins/caddy/README.md b/plugins/caddy/README.md
@@ -379,6 +379,7 @@ What does these directives mean?
 | `cdn.service_id`                          | The service id if required, depending the provider                                                                                           | `123456_id`                                                                                                             |
 | `cdn.zone_id`                             | The zone id if required, depending the provider                                                                                              | `anywhere_zone`                                                                                                         |
 | `default_cache_control`                   | Set the default value of `Cache-Control` response header if not set by upstream (Souin treats empty `Cache-Control` as `public` if omitted)  | `no-store`                                                                                                              |
+| `max_cachable_body_bytes`                          | Set the maximum size (in bytes) for a response body to be cached (unlimited if omited)                                                       | `1048576` (1MB)                                                                                                         |
 | `key`                                     | Override the key generation with the ability to disable unecessary parts                                                                     |                                                                                                                         |
 | `key.disable_body`                        | Disable the body part in the key (GraphQL context)                                                                                           | `true`<br/><br/>`(default: false)`                                                                                      |
 | `key.disable_host`                        | Disable the host part in the key                                                                                                             | `true`<br/><br/>`(default: false)`                                                                                      |

diff --git a/plugins/caddy/configuration.go b/plugins/caddy/configuration.go
@@ -22,6 +22,8 @@ type DefaultCache struct {
 	CDN       configurationtypes.CDN `json:"cdn"`
 	// The default Cache-Control header value if none set by the upstream server.
 	DefaultCacheControl string `json:"default_cache_control"`
+	// The maximum body size (in bytes) to be stored into cache.
+	MaxBodyBytes uint64 `json:"max_cachable_body_bytes"`
 	// Redis provider configuration.
 	Distributed bool `json:"distributed"`
 	// Headers to add to the cache key if they are present.
@@ -140,6 +142,11 @@ func (d *DefaultCache) GetDefaultCacheControl() string {
 	return d.DefaultCacheControl
 }
 
+// GetMaxBodyBytes returns the maximum body size (in bytes) to be cached
+func (d *DefaultCache) GetMaxBodyBytes() uint64 {
+	return d.MaxBodyBytes
+}
+
 // Configuration holder
 type Configuration struct {
 	// Default cache to fallback on when none are redefined.
@@ -412,6 +419,14 @@ func parseConfiguration(cfg *Configuration, h *caddyfile.Dispenser, isGlobal boo
 			case "default_cache_control":
 				args := h.RemainingArgs()
 				cfg.DefaultCache.DefaultCacheControl = strings.Join(args, " ")
+			case "max_cachable_body_bytes":
+				args := h.RemainingArgs()
+				maxBodyBytes, err := strconv.ParseUint(args[0], 10, 64)
+				if err != nil {
+					return h.Errf("unsupported max_cachable_body_bytes: %s", args)
+				} else {
+					cfg.DefaultCache.MaxBodyBytes = maxBodyBytes
+				}
 			case "etcd":
 				cfg.DefaultCache.Distributed = true
 				provider := configurationtypes.CacheProvider{}

diff --git a/plugins/caddy/httpcache.go b/plugins/caddy/httpcache.go
@@ -192,6 +192,9 @@ func (s *SouinCaddyMiddleware) FromApp(app *SouinApp) error {
 	if dc.DefaultCacheControl == "" {
 		s.Configuration.DefaultCache.DefaultCacheControl = appDc.DefaultCacheControl
 	}
+	if dc.MaxBodyBytes == 0 {
+		s.Configuration.DefaultCache.MaxBodyBytes = appDc.MaxBodyBytes
+	}
 	if dc.CacheName == "" {
 		s.Configuration.DefaultCache.CacheName = appDc.CacheName
 	}

diff --git a/plugins/caddy/httpcache_test.go b/plugins/caddy/httpcache_test.go
@@ -294,6 +294,62 @@ func TestNotHandledRoute(t *testing.T) {
 	}
 }
 
+func TestMaxBodyByte(t *testing.T) {
+	tester := caddytest.NewTester(t)
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port 9080
+		https_port 9443
+		cache {
+			ttl 5s
+			max_cachable_body_bytes 30
+		}
+	}
+	localhost:9080 {
+		route /max-body-bytes-stored {
+			cache
+			respond "Hello, Max body bytes stored!"
+		}
+		route /max-body-bytes-not-stored {
+			cache
+			respond "Hello, Max body bytes not stored due to the response length!"
+		}
+	}`, "caddyfile")
+
+	respStored1, _ := tester.AssertGetResponse(`http://localhost:9080/max-body-bytes-stored`, 200, "Hello, Max body bytes stored!")
+	respStored2, _ := tester.AssertGetResponse(`http://localhost:9080/max-body-bytes-stored`, 200, "Hello, Max body bytes stored!")
+	if respStored1.Header.Get("Cache-Status") != "Souin; fwd=uri-miss; stored; key=GET-http-localhost:9080-/max-body-bytes-stored" {
+		t.Errorf("unexpected Cache-Status header value %v", respStored1.Header.Get("Cache-Status"))
+	}
+	if respStored1.Header.Get("Age") != "" {
+		t.Errorf("unexpected Age header %v", respStored1.Header.Get("Age"))
+	}
+
+	if respStored2.Header.Get("Cache-Status") != "Souin; hit; ttl=4; key=GET-http-localhost:9080-/max-body-bytes-stored" {
+		t.Errorf("unexpected Cache-Status header value %v", respStored2.Header.Get("Cache-Status"))
+	}
+	if respStored2.Header.Get("Age") == "" {
+		t.Error("Age header should be present")
+	}
+
+	respNotStored1, _ := tester.AssertGetResponse(`http://localhost:9080/max-body-bytes-not-stored`, 200, "Hello, Max body bytes not stored due to the response length!")
+	respNotStored2, _ := tester.AssertGetResponse(`http://localhost:9080/max-body-bytes-not-stored`, 200, "Hello, Max body bytes not stored due to the response length!")
+	if respNotStored1.Header.Get("Cache-Status") != "Souin; fwd=uri-miss; detail=UPSTREAM-RESPONSE-TOO-LARGE; key=GET-http-localhost:9080-/max-body-bytes-not-stored" {
+		t.Errorf("unexpected Cache-Status header value %v", respNotStored1.Header.Get("Cache-Status"))
+	}
+	if respNotStored1.Header.Get("Age") != "" {
+		t.Errorf("unexpected Age header %v", respNotStored1.Header.Get("Age"))
+	}
+
+	if respNotStored2.Header.Get("Cache-Status") != "Souin; fwd=uri-miss; detail=UPSTREAM-RESPONSE-TOO-LARGE; key=GET-http-localhost:9080-/max-body-bytes-not-stored" {
+		t.Errorf("unexpected Cache-Status header value %v", respNotStored2.Header.Get("Cache-Status"))
+	}
+	if respNotStored2.Header.Get("Age") != "" {
+		t.Errorf("unexpected Age header %v", respNotStored2.Header.Get("Age"))
+	}
+}
+
 func TestMultiProvider(t *testing.T) {
 	var wg sync.WaitGroup
 	var responses []*http.Response

diff --git a/plugins/kratos/configuration.go b/plugins/kratos/configuration.go
@@ -311,6 +311,11 @@ func parseDefaultCache(dcConfiguration map[string]config.Value) *configurationty
 			}
 		case "default_cache_control":
 			dc.DefaultCacheControl, _ = defaultCacheV.String()
+		case "max_cachable_body_bytes":
+			mbb, ok := defaultCacheV.Load().(uint64)
+			if ok {
+				dc.MaxBodyBytes = mbb
+			}
 		}
 	}
 

diff --git a/plugins/souin/agnostic/configuration_parser.go b/plugins/souin/agnostic/configuration_parser.go
@@ -270,6 +270,8 @@ func parseDefaultCache(dcConfiguration map[string]interface{}) *configurationtyp
 			}
 		case "default_cache_control":
 			dc.DefaultCacheControl, _ = defaultCacheV.(string)
+		case "max_cachable_body_bytes":
+			dc.MaxBodyBytes, _ = defaultCacheV.(uint64)
 		}
 	}
 

diff --git a/plugins/traefik/main.go b/plugins/traefik/main.go
@@ -203,6 +203,8 @@ func parseConfiguration(c map[string]interface{}) Configuration {
 					dc.Storers = parseStringSlice(defaultCacheV)
 				case "default_cache_control":
 					dc.DefaultCacheControl = defaultCacheV.(string)
+				case "max_cachable_body_bytes":
+					dc.MaxBodyBytes, _ = defaultCacheV.(uint64)
 				}
 			}
 			configuration.DefaultCache = &dc

diff --git a/plugins/traefik/vendor/github.com/darkweak/souin/configurationtypes/types.go b/plugins/traefik/vendor/github.com/darkweak/souin/configurationtypes/types.go