From 07455c8a27fbbc095402657d325384b7719f8528 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@vmware.com>
Date: Tue, 1 Dec 2020 17:44:57 +0000
Subject: [PATCH] ContentCachingResponseWrapper skips contentLength for chunked
 responses

Closes gh-26182
---
 .../util/ContentCachingResponseWrapper.java   | 15 +++-
 .../ContentCachingResponseWrapperTests.java   | 72 +++++++++++++++++++
 2 files changed, 86 insertions(+), 1 deletion(-)
 create mode 100644 spring-web/src/test/java/org/springframework/web/filter/ContentCachingResponseWrapperTests.java

diff --git a/spring-web/src/main/java/org/springframework/web/util/ContentCachingResponseWrapper.java b/spring-web/src/main/java/org/springframework/web/util/ContentCachingResponseWrapper.java
index 1a814bea786b..98003b19599a 100644
--- a/spring-web/src/main/java/org/springframework/web/util/ContentCachingResponseWrapper.java
+++ b/spring-web/src/main/java/org/springframework/web/util/ContentCachingResponseWrapper.java
@@ -26,6 +26,8 @@
 import javax.servlet.http.HttpServletResponse;
 import javax.servlet.http.HttpServletResponseWrapper;
 
+import org.springframework.http.HttpHeaders;
+import org.springframework.util.ClassUtils;
 import org.springframework.util.FastByteArrayOutputStream;
 
 /**
@@ -41,6 +43,11 @@
  */
 public class ContentCachingResponseWrapper extends HttpServletResponseWrapper {
 
+	/** Checking for Servlet 3.0+ HttpServletResponse.getHeader(String) */
+	private static final boolean servlet3Present =
+			ClassUtils.hasMethod(HttpServletResponse.class, "getHeader", String.class);
+
+
 	private final FastByteArrayOutputStream content = new FastByteArrayOutputStream(1024);
 
 	private final ServletOutputStream outputStream = new ResponseServletOutputStream();
@@ -214,7 +221,13 @@ protected void copyBodyToResponse(boolean complete) throws IOException {
 		if (this.content.size() > 0) {
 			HttpServletResponse rawResponse = (HttpServletResponse) getResponse();
 			if ((complete || this.contentLength != null) && !rawResponse.isCommitted()) {
-				rawResponse.setContentLength(complete ? this.content.size() : this.contentLength);
+				boolean hasTransferEncoding = false;
+				if (servlet3Present) {
+					hasTransferEncoding = (rawResponse.getHeader(HttpHeaders.TRANSFER_ENCODING) != null);
+				}
+				if (!hasTransferEncoding) {
+					rawResponse.setContentLength(complete ? this.content.size() : this.contentLength);
+				}
 				this.contentLength = null;
 			}
 			this.content.writeTo(rawResponse.getOutputStream());
diff --git a/spring-web/src/test/java/org/springframework/web/filter/ContentCachingResponseWrapperTests.java b/spring-web/src/test/java/org/springframework/web/filter/ContentCachingResponseWrapperTests.java
new file mode 100644
index 000000000000..28ba42bdf752
--- /dev/null
+++ b/spring-web/src/test/java/org/springframework/web/filter/ContentCachingResponseWrapperTests.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright 2002-2020 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.web.filter;
+
+import java.nio.charset.StandardCharsets;
+
+import javax.servlet.http.HttpServletResponse;
+
+import org.junit.Test;
+
+import org.springframework.http.HttpHeaders;
+import org.springframework.mock.web.test.MockHttpServletResponse;
+import org.springframework.util.FileCopyUtils;
+import org.springframework.web.util.ContentCachingResponseWrapper;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * Unit tests for {@link ContentCachingResponseWrapper}.
+ * @author Rossen Stoyanchev
+ */
+public class ContentCachingResponseWrapperTests {
+
+	@Test
+	public void copyBodyToResponse() throws Exception {
+		byte[] responseBody = "Hello World".getBytes(StandardCharsets.UTF_8);
+		MockHttpServletResponse response = new MockHttpServletResponse();
+
+		ContentCachingResponseWrapper responseWrapper = new ContentCachingResponseWrapper(response);
+		responseWrapper.setStatus(HttpServletResponse.SC_OK);
+		FileCopyUtils.copy(responseBody, responseWrapper.getOutputStream());
+		responseWrapper.copyBodyToResponse();
+
+		assertEquals(200, response.getStatus());
+		assertTrue(response.getContentLength() > 0);
+		assertArrayEquals(responseBody, response.getContentAsByteArray());
+	}
+
+	@Test
+	public void copyBodyToResponseWithTransferEncoding() throws Exception {
+		byte[] responseBody = "6\r\nHello 5\r\nWorld0\r\n\r\n".getBytes(StandardCharsets.UTF_8);
+		MockHttpServletResponse response = new MockHttpServletResponse();
+
+		ContentCachingResponseWrapper responseWrapper = new ContentCachingResponseWrapper(response);
+		responseWrapper.setStatus(HttpServletResponse.SC_OK);
+		responseWrapper.setHeader(HttpHeaders.TRANSFER_ENCODING, "chunked");
+		FileCopyUtils.copy(responseBody, responseWrapper.getOutputStream());
+		responseWrapper.copyBodyToResponse();
+
+		assertEquals(200, response.getStatus());
+		assertEquals("chunked", response.getHeader(HttpHeaders.TRANSFER_ENCODING));
+		assertNull(response.getHeader(HttpHeaders.CONTENT_LENGTH));
+		assertArrayEquals(responseBody, response.getContentAsByteArray());
+	}
+
+}