From ea4a1ec580cd84f644083cee2adb040c89e87b0e Mon Sep 17 00:00:00 2001
From: Ryan Rupp <3022260+ryanrupp@users.noreply.github.com>
Date: Thu, 5 Jan 2023 14:44:05 -0600
Subject: [PATCH] Optimize the initial buffer size when using content cache
 limiting

To avoid over-allocating the initial buffer for content caching:
1. If the content size is known and is smaller than the content limit, use that
2. If the content size is unknown use the default initial buffer size (or content limit if smaller) and allow it to group as needed

This allows for scenarios where limiting behavior is desired to avoid overly large caching but at the same time the majority of requests are not near that limit. For example, allowing a maximum of 1MB content caching but the majority of requests are in the single digit KB size. Rather than allocate the worst case 1MB byte arrays per request these can be scaled in to be more appropriately sized.
---
 .../util/ContentCachingRequestWrapper.java    |  9 ++++-
 .../ContentCachingRequestWrapperTests.java    | 40 ++++++++++++++++++-
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/spring-web/src/main/java/org/springframework/web/util/ContentCachingRequestWrapper.java b/spring-web/src/main/java/org/springframework/web/util/ContentCachingRequestWrapper.java
index 29b93e01a8b5..b7d836bb790f 100644
--- a/spring-web/src/main/java/org/springframework/web/util/ContentCachingRequestWrapper.java
+++ b/spring-web/src/main/java/org/springframework/web/util/ContentCachingRequestWrapper.java
@@ -55,6 +55,8 @@
  */
 public class ContentCachingRequestWrapper extends HttpServletRequestWrapper {
 
+	private static final int DEFAULT_INITIAL_BUFFER_SIZE_WHEN_UNKNOWN = 1024;
+
 	private final ByteArrayOutputStream cachedContent;
 
 	@Nullable
@@ -74,7 +76,8 @@ public class ContentCachingRequestWrapper extends HttpServletRequestWrapper {
 	public ContentCachingRequestWrapper(HttpServletRequest request) {
 		super(request);
 		int contentLength = request.getContentLength();
-		this.cachedContent = new ByteArrayOutputStream(contentLength >= 0 ? contentLength : 1024);
+		this.cachedContent = new ByteArrayOutputStream(contentLength >= 0 ?
+				contentLength : DEFAULT_INITIAL_BUFFER_SIZE_WHEN_UNKNOWN);
 		this.contentCacheLimit = null;
 	}
 
@@ -87,7 +90,9 @@ public ContentCachingRequestWrapper(HttpServletRequest request) {
 	 */
 	public ContentCachingRequestWrapper(HttpServletRequest request, int contentCacheLimit) {
 		super(request);
-		this.cachedContent = new ByteArrayOutputStream(contentCacheLimit);
+		int contentLength = request.getContentLength();
+		int guessedInitBufferSize = contentLength >= 0 ? contentLength : DEFAULT_INITIAL_BUFFER_SIZE_WHEN_UNKNOWN;
+		this.cachedContent = new ByteArrayOutputStream(Math.min(guessedInitBufferSize, contentCacheLimit));
 		this.contentCacheLimit = contentCacheLimit;
 	}
 
diff --git a/spring-web/src/test/java/org/springframework/web/util/ContentCachingRequestWrapperTests.java b/spring-web/src/test/java/org/springframework/web/util/ContentCachingRequestWrapperTests.java
index d9132561a26c..0adc609fcdd0 100644
--- a/spring-web/src/test/java/org/springframework/web/util/ContentCachingRequestWrapperTests.java
+++ b/spring-web/src/test/java/org/springframework/web/util/ContentCachingRequestWrapperTests.java
@@ -48,13 +48,15 @@ public class ContentCachingRequestWrapperTests {
 
 	@Test
 	void cachedContent() throws Exception {
+		byte[] contentPayload = "Hello World".getBytes(CHARSET);
 		this.request.setMethod(GET);
 		this.request.setCharacterEncoding(CHARSET);
-		this.request.setContent("Hello World".getBytes(CHARSET));
+		this.request.setContent(contentPayload);
 
 		ContentCachingRequestWrapper wrapper = new ContentCachingRequestWrapper(this.request);
 		byte[] response = FileCopyUtils.copyToByteArray(wrapper.getInputStream());
 		assertThat(wrapper.getContentAsByteArray()).isEqualTo(response);
+		assertThat(wrapper.getContentAsByteArray()).isEqualTo(contentPayload);
 	}
 
 	@Test
@@ -117,4 +119,40 @@ void inputStreamFormPostRequest() throws Exception {
 		assertThat(wrapper.getContentAsByteArray()).isEqualTo(response);
 	}
 
+	/**
+	 * When content limiting, if the content length is known and is lower than the content limit, the internal
+	 * buffer should allocate that size. This exaggerates this scenario by using a very large content limit that
+	 * would cause an OOM Error likely otherwise.
+	 */
+	@Test
+	void cachedContentWithLimitKnownLengthAvoidOverAllocation() throws Exception {
+		byte[] contentPayload = "Hello World".getBytes(CHARSET);
+		this.request.setMethod(POST);
+		this.request.setCharacterEncoding(CHARSET);
+		this.request.setContent(contentPayload);
+
+		ContentCachingRequestWrapper wrapper = new ContentCachingRequestWrapper(this.request, Integer.MAX_VALUE);
+
+		byte[] contents = FileCopyUtils.copyToByteArray(wrapper.getInputStream());
+		assertThat(contents).isEqualTo(contentPayload);
+		assertThat(wrapper.getContentAsByteArray()).isEqualTo(contentPayload);
+	}
+
+	/**
+	 * When content limiting, if the content length is unknown, rather than allocating the content limit upfront
+	 * (which may be quite a bit larger than the actual request), a smaller initial buffer should be used that is
+	 * allowed to grow as needed up to the limit.  This exaggerates this scenario by using a very large content limit
+	 * that would cause an OOM Error likely otherwise.
+	 */
+	@Test
+	void cachedContentWithLimitUnknownLengthAvoidOverAllocation() throws Exception {
+		this.request.setMethod(POST);
+		this.request.setCharacterEncoding(CHARSET);
+
+		ContentCachingRequestWrapper wrapper = new ContentCachingRequestWrapper(this.request, Integer.MAX_VALUE);
+
+		byte[] contents = FileCopyUtils.copyToByteArray(wrapper.getInputStream());
+		assertThat(contents).isEmpty();
+		assertThat(wrapper.getContentAsByteArray()).isEmpty();
+	}
 }