Skip to content

Commit 8189c90

Browse files
sbrannensbrannen
committed
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
1 parent 9d2a874 commit 8189c90

File tree

3 files changed

+107
-23
lines changed

3 files changed

+107
-23
lines changed

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

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
* @author Rod Johnson
5656
* @author Brian Clozel
5757
* @author Vedran Pavic
58+
* @author Sam Brannen
5859
* @since 1.0.2
5960
*/
6061
public class MockHttpServletResponse implements HttpServletResponse {
@@ -573,20 +574,22 @@ public void addIntHeader(String name, int value) {
573574
}
574575

575576
private void setHeaderValue(String name, Object value) {
576-
if (setSpecialHeader(name, value)) {
577+
boolean replaceHeader = true;
578+
if (setSpecialHeader(name, value, replaceHeader)) {
577579
return;
578580
}
579-
doAddHeaderValue(name, value, true);
581+
doAddHeaderValue(name, value, replaceHeader);
580582
}
581583

582584
private void addHeaderValue(String name, Object value) {
583-
if (setSpecialHeader(name, value)) {
585+
boolean replaceHeader = false;
586+
if (setSpecialHeader(name, value, replaceHeader)) {
584587
return;
585588
}
586-
doAddHeaderValue(name, value, false);
589+
doAddHeaderValue(name, value, replaceHeader);
587590
}
588591

589-
private boolean setSpecialHeader(String name, Object value) {
592+
private boolean setSpecialHeader(String name, Object value, boolean replaceHeader) {
590593
if (HttpHeaders.CONTENT_TYPE.equalsIgnoreCase(name)) {
591594
setContentType(value.toString());
592595
return true;
@@ -605,7 +608,12 @@ else if (HttpHeaders.CONTENT_LANGUAGE.equalsIgnoreCase(name)) {
605608
}
606609
else if (HttpHeaders.SET_COOKIE.equalsIgnoreCase(name)) {
607610
MockCookie cookie = MockCookie.parse(value.toString());
608-
addCookie(cookie);
611+
if (replaceHeader) {
612+
setCookie(cookie);
613+
}
614+
else {
615+
addCookie(cookie);
616+
}
609617
return true;
610618
}
611619
else {
@@ -628,6 +636,20 @@ private void doAddHeaderValue(String name, Object value, boolean replace) {
628636
}
629637
}
630638

639+
/**
640+
* Set the {@code Set-Cookie} header to the supplied {@link Cookie},
641+
* overwriting any previous cookies.
642+
* @param cookie the {@code Cookie} to set
643+
* @since 5.1.10
644+
* @see #addCookie(Cookie)
645+
*/
646+
private void setCookie(Cookie cookie) {
647+
Assert.notNull(cookie, "Cookie must not be null");
648+
this.cookies.clear();
649+
this.cookies.add(cookie);
650+
doAddHeaderValue(HttpHeaders.SET_COOKIE, getCookieHeader(cookie), true);
651+
}
652+
631653
@Override
632654
public void setStatus(int status) {
633655
if (!this.isCommitted()) {

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

Lines changed: 51 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -322,22 +322,36 @@ public void modifyStatusMessageAfterSendError() throws IOException {
322322
assertEquals(HttpServletResponse.SC_NOT_FOUND, response.getStatus());
323323
}
324324

325+
/**
326+
* @since 5.1.10
327+
*/
325328
@Test
326-
public void setCookieHeaderValid() {
329+
public void setCookieHeader() {
330+
response.setHeader(HttpHeaders.SET_COOKIE, "SESSION=123; Path=/; Secure; HttpOnly; SameSite=Lax");
331+
assertNumCookies(1);
332+
assertPrimarySessionCookie("123");
333+
334+
// Setting the Set-Cookie header a 2nd time should overwrite the previous value
335+
response.setHeader(HttpHeaders.SET_COOKIE, "SESSION=999; Path=/; Secure; HttpOnly; SameSite=Lax");
336+
assertNumCookies(1);
337+
assertPrimarySessionCookie("999");
338+
}
339+
340+
@Test
341+
public void addCookieHeader() {
327342
response.addHeader(HttpHeaders.SET_COOKIE, "SESSION=123; Path=/; Secure; HttpOnly; SameSite=Lax");
328-
Cookie cookie = response.getCookie("SESSION");
329-
assertNotNull(cookie);
330-
assertTrue(cookie instanceof MockCookie);
331-
assertEquals("SESSION", cookie.getName());
332-
assertEquals("123", cookie.getValue());
333-
assertEquals("/", cookie.getPath());
334-
assertTrue(cookie.getSecure());
335-
assertTrue(cookie.isHttpOnly());
336-
assertEquals("Lax", ((MockCookie) cookie).getSameSite());
343+
assertNumCookies(1);
344+
assertPrimarySessionCookie("123");
345+
346+
// Adding a 2nd cookie header should result in 2 cookies.
347+
response.addHeader(HttpHeaders.SET_COOKIE, "SESSION=999; Path=/; Secure; HttpOnly; SameSite=Lax");
348+
assertNumCookies(2);
349+
assertPrimarySessionCookie("123");
350+
assertCookieValues("123", "999");
337351
}
338352

339353
@Test
340-
public void addMockCookie() {
354+
public void addCookie() {
341355
MockCookie mockCookie = new MockCookie("SESSION", "123");
342356
mockCookie.setPath("/");
343357
mockCookie.setDomain("example.com");
@@ -348,9 +362,35 @@ public void addMockCookie() {
348362

349363
response.addCookie(mockCookie);
350364

365+
assertNumCookies(1);
351366
assertEquals("SESSION=123; Path=/; Domain=example.com; Max-Age=0; " +
352367
"Expires=Thu, 01 Jan 1970 00:00:00 GMT; Secure; HttpOnly; SameSite=Lax",
353368
response.getHeader(HttpHeaders.SET_COOKIE));
369+
370+
// Adding a 2nd Cookie should result in 2 Cookies.
371+
response.addCookie(new MockCookie("SESSION", "999"));
372+
assertNumCookies(2);
373+
assertCookieValues("123", "999");
374+
}
375+
376+
private void assertNumCookies(int expected) {
377+
assertEquals(expected, this.response.getCookies().length);
378+
}
379+
380+
private void assertCookieValues(String... expected) {
381+
assertArrayEquals(expected, Arrays.stream(response.getCookies()).map(Cookie::getValue).toArray());
382+
}
383+
384+
private void assertPrimarySessionCookie(String expectedValue) {
385+
Cookie cookie = this.response.getCookie("SESSION");
386+
assertNotNull(cookie);
387+
assertTrue(cookie instanceof MockCookie);
388+
assertEquals("SESSION", cookie.getName());
389+
assertEquals(expectedValue, cookie.getValue());
390+
assertEquals("/", cookie.getPath());
391+
assertTrue(cookie.getSecure());
392+
assertTrue(cookie.isHttpOnly());
393+
assertEquals("Lax", ((MockCookie) cookie).getSameSite());
354394
}
355395

356396
}

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

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
* @author Rod Johnson
5656
* @author Brian Clozel
5757
* @author Vedran Pavic
58+
* @author Sam Brannen
5859
* @since 1.0.2
5960
*/
6061
public class MockHttpServletResponse implements HttpServletResponse {
@@ -573,20 +574,22 @@ public void addIntHeader(String name, int value) {
573574
}
574575

575576
private void setHeaderValue(String name, Object value) {
576-
if (setSpecialHeader(name, value)) {
577+
boolean replaceHeader = true;
578+
if (setSpecialHeader(name, value, replaceHeader)) {
577579
return;
578580
}
579-
doAddHeaderValue(name, value, true);
581+
doAddHeaderValue(name, value, replaceHeader);
580582
}
581583

582584
private void addHeaderValue(String name, Object value) {
583-
if (setSpecialHeader(name, value)) {
585+
boolean replaceHeader = false;
586+
if (setSpecialHeader(name, value, replaceHeader)) {
584587
return;
585588
}
586-
doAddHeaderValue(name, value, false);
589+
doAddHeaderValue(name, value, replaceHeader);
587590
}
588591

589-
private boolean setSpecialHeader(String name, Object value) {
592+
private boolean setSpecialHeader(String name, Object value, boolean replaceHeader) {
590593
if (HttpHeaders.CONTENT_TYPE.equalsIgnoreCase(name)) {
591594
setContentType(value.toString());
592595
return true;
@@ -605,7 +608,12 @@ else if (HttpHeaders.CONTENT_LANGUAGE.equalsIgnoreCase(name)) {
605608
}
606609
else if (HttpHeaders.SET_COOKIE.equalsIgnoreCase(name)) {
607610
MockCookie cookie = MockCookie.parse(value.toString());
608-
addCookie(cookie);
611+
if (replaceHeader) {
612+
setCookie(cookie);
613+
}
614+
else {
615+
addCookie(cookie);
616+
}
609617
return true;
610618
}
611619
else {
@@ -628,6 +636,20 @@ private void doAddHeaderValue(String name, Object value, boolean replace) {
628636
}
629637
}
630638

639+
/**
640+
* Set the {@code Set-Cookie} header to the supplied {@link Cookie},
641+
* overwriting any previous cookies.
642+
* @param cookie the {@code Cookie} to set
643+
* @since 5.1.10
644+
* @see #addCookie(Cookie)
645+
*/
646+
private void setCookie(Cookie cookie) {
647+
Assert.notNull(cookie, "Cookie must not be null");
648+
this.cookies.clear();
649+
this.cookies.add(cookie);
650+
doAddHeaderValue(HttpHeaders.SET_COOKIE, getCookieHeader(cookie), true);
651+
}
652+
631653
@Override
632654
public void setStatus(int status) {
633655
if (!this.isCommitted()) {

0 commit comments

Comments
 (0)