Allow Set-Cookie header to be overwritten in MockHttpServletResponse

Prior to this commit, there was no way to replace the Set-Cookie header
via MockHttpServletResponse. Specifically, an invocation of setHeader()
for the Set-Cookie header resulted in an additional Set-Cookie header
instead of replacing the existing one, which is in violation of the
contract for javax.servlet.http.HttpServletResponse.setHeader(...).

This commit refactors the internals of MockHttpServletResponse to ensure
that an existing Set-Cookie header is overwritten when set via an
invocation of setHeader(). This commit also verifies the expected
behavior for addHeader() and addCookie() with regard to multiple cookies.

Closes gh-23512

spring-test/src/main/java/org/springframework/mock/web/MockHttpServletResponse.java

@@ -55,6 +55,7 @@ import org.springframework.web.util.WebUtils;
  * @author Rod Johnson
  * @author Brian Clozel
  * @author Vedran Pavic
+ * @author Sam Brannen
  * @since 1.0.2
  */
 public class MockHttpServletResponse implements HttpServletResponse {
@@ -573,20 +574,22 @@ public class MockHttpServletResponse implements HttpServletResponse {
 	}
 
 	private void setHeaderValue(String name, Object value) {
-		if (setSpecialHeader(name, value)) {
+		boolean replaceHeader = true;
+		if (setSpecialHeader(name, value, replaceHeader)) {
 			return;
 		}
-		doAddHeaderValue(name, value, true);
+		doAddHeaderValue(name, value, replaceHeader);
 	}
 
 	private void addHeaderValue(String name, Object value) {
-		if (setSpecialHeader(name, value)) {
+		boolean replaceHeader = false;
+		if (setSpecialHeader(name, value, replaceHeader)) {
 			return;
 		}
-		doAddHeaderValue(name, value, false);
+		doAddHeaderValue(name, value, replaceHeader);
 	}
 
-	private boolean setSpecialHeader(String name, Object value) {
+	private boolean setSpecialHeader(String name, Object value, boolean replaceHeader) {
 		if (HttpHeaders.CONTENT_TYPE.equalsIgnoreCase(name)) {
 			setContentType(value.toString());
 			return true;
@@ -605,7 +608,12 @@ public class MockHttpServletResponse implements HttpServletResponse {
 		}
 		else if (HttpHeaders.SET_COOKIE.equalsIgnoreCase(name)) {
 			MockCookie cookie = MockCookie.parse(value.toString());
-			addCookie(cookie);
+			if (replaceHeader) {
+				setCookie(cookie);
+			}
+			else {
+				addCookie(cookie);
+			}
 			return true;
 		}
 		else {
@@ -628,6 +636,20 @@ public class MockHttpServletResponse implements HttpServletResponse {
 		}
 	}
 
+	/**
+	 * Set the {@code Set-Cookie} header to the supplied {@link Cookie},
+	 * overwriting any previous cookies.
+	 * @param cookie the {@code Cookie} to set
+	 * @since 5.1.10
+	 * @see #addCookie(Cookie)
+	 */
+	private void setCookie(Cookie cookie) {
+		Assert.notNull(cookie, "Cookie must not be null");
+		this.cookies.clear();
+		this.cookies.add(cookie);
+		doAddHeaderValue(HttpHeaders.SET_COOKIE, getCookieHeader(cookie), true);
+	}
+
 	@Override
 	public void setStatus(int status) {
 		if (!this.isCommitted()) {

spring-test/src/test/java/org/springframework/mock/web/MockHttpServletResponseTests.java

@@ -322,22 +322,36 @@ public class MockHttpServletResponseTests {
 		assertEquals(HttpServletResponse.SC_NOT_FOUND, response.getStatus());
 	}
 
+	/**
+	 * @since 5.1.10
+	 */
 	@Test
-	public void setCookieHeaderValid() {
-		response.addHeader(HttpHeaders.SET_COOKIE, "SESSION=123; Path=/; Secure; HttpOnly; SameSite=Lax");
-		Cookie cookie = response.getCookie("SESSION");
-		assertNotNull(cookie);
-		assertTrue(cookie instanceof MockCookie);
-		assertEquals("SESSION", cookie.getName());
-		assertEquals("123", cookie.getValue());
-		assertEquals("/", cookie.getPath());
-		assertTrue(cookie.getSecure());
-		assertTrue(cookie.isHttpOnly());
-		assertEquals("Lax", ((MockCookie) cookie).getSameSite());
+	public void setCookieHeader() {
+		response.setHeader(HttpHeaders.SET_COOKIE, "SESSION=123; Path=/; Secure; HttpOnly; SameSite=Lax");
+		assertNumCookies(1);
+		assertPrimarySessionCookie("123");
+
+		// Setting the Set-Cookie header a 2nd time should overwrite the previous value
+		response.setHeader(HttpHeaders.SET_COOKIE, "SESSION=999; Path=/; Secure; HttpOnly; SameSite=Lax");
+		assertNumCookies(1);
+		assertPrimarySessionCookie("999");
 	}
 
 	@Test
-	public void addMockCookie() {
+	public void addCookieHeader() {
+		response.addHeader(HttpHeaders.SET_COOKIE, "SESSION=123; Path=/; Secure; HttpOnly; SameSite=Lax");
+		assertNumCookies(1);
+		assertPrimarySessionCookie("123");
+
+		// Adding a 2nd cookie header should result in 2 cookies.
+		response.addHeader(HttpHeaders.SET_COOKIE, "SESSION=999; Path=/; Secure; HttpOnly; SameSite=Lax");
+		assertNumCookies(2);
+		assertPrimarySessionCookie("123");
+		assertCookieValues("123", "999");
+	}
+
+	@Test
+	public void addCookie() {
 		MockCookie mockCookie = new MockCookie("SESSION", "123");
 		mockCookie.setPath("/");
 		mockCookie.setDomain("example.com");
@@ -348,9 +362,35 @@ public class MockHttpServletResponseTests {
 
 		response.addCookie(mockCookie);
 
+		assertNumCookies(1);
 		assertEquals("SESSION=123; Path=/; Domain=example.com; Max-Age=0; " +
 				"Expires=Thu, 01 Jan 1970 00:00:00 GMT; Secure; HttpOnly; SameSite=Lax",
 				response.getHeader(HttpHeaders.SET_COOKIE));
+
+		// Adding a 2nd Cookie should result in 2 Cookies.
+		response.addCookie(new MockCookie("SESSION", "999"));
+		assertNumCookies(2);
+		assertCookieValues("123", "999");
+	}
+
+	private void assertNumCookies(int expected) {
+		assertEquals(expected, this.response.getCookies().length);
+	}
+
+	private void assertCookieValues(String... expected) {
+		assertArrayEquals(expected, Arrays.stream(response.getCookies()).map(Cookie::getValue).toArray());
+	}
+
+	private void assertPrimarySessionCookie(String expectedValue) {
+		Cookie cookie = this.response.getCookie("SESSION");
+		assertNotNull(cookie);
+		assertTrue(cookie instanceof MockCookie);
+		assertEquals("SESSION", cookie.getName());
+		assertEquals(expectedValue, cookie.getValue());
+		assertEquals("/", cookie.getPath());
+		assertTrue(cookie.getSecure());
+		assertTrue(cookie.isHttpOnly());
+		assertEquals("Lax", ((MockCookie) cookie).getSameSite());
 	}
 
 }

spring-web/src/test/java/org/springframework/mock/web/test/MockHttpServletResponse.java

@@ -55,6 +55,7 @@ import org.springframework.web.util.WebUtils;
  * @author Rod Johnson
  * @author Brian Clozel
  * @author Vedran Pavic
+ * @author Sam Brannen
  * @since 1.0.2
  */
 public class MockHttpServletResponse implements HttpServletResponse {
@@ -573,20 +574,22 @@ public class MockHttpServletResponse implements HttpServletResponse {
 	}
 
 	private void setHeaderValue(String name, Object value) {
-		if (setSpecialHeader(name, value)) {
+		boolean replaceHeader = true;
+		if (setSpecialHeader(name, value, replaceHeader)) {
 			return;
 		}
-		doAddHeaderValue(name, value, true);
+		doAddHeaderValue(name, value, replaceHeader);
 	}
 
 	private void addHeaderValue(String name, Object value) {
-		if (setSpecialHeader(name, value)) {
+		boolean replaceHeader = false;
+		if (setSpecialHeader(name, value, replaceHeader)) {
 			return;
 		}
-		doAddHeaderValue(name, value, false);
+		doAddHeaderValue(name, value, replaceHeader);
 	}
 
-	private boolean setSpecialHeader(String name, Object value) {
+	private boolean setSpecialHeader(String name, Object value, boolean replaceHeader) {
 		if (HttpHeaders.CONTENT_TYPE.equalsIgnoreCase(name)) {
 			setContentType(value.toString());
 			return true;
@@ -605,7 +608,12 @@ public class MockHttpServletResponse implements HttpServletResponse {
 		}
 		else if (HttpHeaders.SET_COOKIE.equalsIgnoreCase(name)) {
 			MockCookie cookie = MockCookie.parse(value.toString());
-			addCookie(cookie);
+			if (replaceHeader) {
+				setCookie(cookie);
+			}
+			else {
+				addCookie(cookie);
+			}
 			return true;
 		}
 		else {
@@ -628,6 +636,20 @@ public class MockHttpServletResponse implements HttpServletResponse {
 		}
 	}
 
+	/**
+	 * Set the {@code Set-Cookie} header to the supplied {@link Cookie},
+	 * overwriting any previous cookies.
+	 * @param cookie the {@code Cookie} to set
+	 * @since 5.1.10
+	 * @see #addCookie(Cookie)
+	 */
+	private void setCookie(Cookie cookie) {
+		Assert.notNull(cookie, "Cookie must not be null");
+		this.cookies.clear();
+		this.cookies.add(cookie);
+		doAddHeaderValue(HttpHeaders.SET_COOKIE, getCookieHeader(cookie), true);
+	}
+
 	@Override
 	public void setStatus(int status) {
 		if (!this.isCommitted()) {