Add support for adding cookies as headers in MockHttpServletResponse

Issue: SPR-17110

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

@@ -0,0 +1,126 @@
+/*
+ * Copyright 2002-2018 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
+ *
+ *      http://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.mock.web;
+
+import javax.servlet.http.Cookie;
+
+import org.springframework.lang.Nullable;
+
+/**
+ * A {@code Cookie} subclass with the additional cookie directives as defined in the
+ * <a href="https://tools.ietf.org/html/rfc6265">RFC 6265</a>.
+ *
+ * @author Vedran Pavic
+ * @since 5.1
+ */
+public class MockCookie extends Cookie {
+
+	private static final long serialVersionUID = 4312531139502726325L;
+
+	@Nullable
+	private String sameSite;
+
+	/**
+	 * Constructs a {@code MockCookie} instance with the specified name and value.
+	 *
+	 * @param name  the cookie name
+	 * @param value the cookie value
+	 * @see Cookie#Cookie(String, String)
+	 */
+	public MockCookie(String name, String value) {
+		super(name, value);
+	}
+
+	/**
+	 * Factory method create {@code MockCookie} instance from Set-Cookie header value.
+	 *
+	 * @param setCookieHeader the Set-Cookie header value
+	 * @return the created cookie instance
+	 */
+	public static MockCookie parse(String setCookieHeader) {
+		String[] cookieParts = setCookieHeader.split("\\s*=\\s*", 2);
+		if (cookieParts.length != 2) {
+			throw new IllegalArgumentException("Invalid Set-Cookie header value");
+		}
+		String name = cookieParts[0];
+		String[] valueAndDirectives = cookieParts[1].split("\\s*;\\s*", 2);
+		String value = valueAndDirectives[0];
+		String[] directives = valueAndDirectives[1].split("\\s*;\\s*");
+		String domain = null;
+		int maxAge = -1;
+		String path = null;
+		boolean secure = false;
+		boolean httpOnly = false;
+		String sameSite = null;
+		for (String directive : directives) {
+			if (directive.startsWith("Domain")) {
+				domain = directive.split("=")[1];
+			}
+			else if (directive.startsWith("Max-Age")) {
+				maxAge = Integer.parseInt(directive.split("=")[1]);
+			}
+			else if (directive.startsWith("Path")) {
+				path = directive.split("=")[1];
+			}
+			else if (directive.startsWith("Secure")) {
+				secure = true;
+			}
+			else if (directive.startsWith("HttpOnly")) {
+				httpOnly = true;
+			}
+			else if (directive.startsWith("SameSite")) {
+				sameSite = directive.split("=")[1];
+			}
+		}
+		MockCookie cookie = new MockCookie(name, value);
+		if (domain != null) {
+			cookie.setDomain(domain);
+		}
+		cookie.setMaxAge(maxAge);
+		cookie.setPath(path);
+		cookie.setSecure(secure);
+		cookie.setHttpOnly(httpOnly);
+		cookie.setSameSite(sameSite);
+		return cookie;
+	}
+
+	/**
+	 * Return the cookie "SameSite" attribute, or {@code null} if not set.
+	 * <p>
+	 * This limits the scope of the cookie such that it will only be attached to same site
+	 * requests if {@code "Strict"} or cross-site requests if {@code "Lax"}.
+	 *
+	 * @see <a href="https://tools.ietf.org/html/draft-ietf-httpbis-rfc6265bis#section-4.1.2.7">RFC6265 bis</a>
+	 */
+	@Nullable
+	public String getSameSite() {
+		return this.sameSite;
+	}
+
+	/**
+	 * Add the "SameSite" attribute to the cookie.
+	 * <p>
+	 * This limits the scope of the cookie such that it will only be attached to same site
+	 * requests if {@code "Strict"} or cross-site requests if {@code "Lax"}.
+	 *
+	 * @see <a href="https://tools.ietf.org/html/draft-ietf-httpbis-rfc6265bis#section-4.1.2.7">RFC6265 bis</a>
+	 */
+	public void setSameSite(@Nullable String sameSite) {
+		this.sameSite = sameSite;
+	}
+
+}

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

@@ -54,6 +54,7 @@ import org.springframework.web.util.WebUtils;
  * @author Juergen Hoeller
  * @author Rod Johnson
  * @author Brian Clozel
+ * @author Vedran Pavic
  * @since 1.0.2
  */
 public class MockHttpServletResponse implements HttpServletResponse {
@@ -353,6 +354,12 @@ public class MockHttpServletResponse implements HttpServletResponse {
 		if (cookie.isHttpOnly()) {
 			buf.append("; HttpOnly");
 		}
+		if (cookie instanceof MockCookie) {
+			MockCookie mockCookie = (MockCookie) cookie;
+			if (StringUtils.hasText(mockCookie.getSameSite())) {
+				buf.append("; SameSite=").append(mockCookie.getSameSite());
+			}
+		}
 		return buf.toString();
 	}
 
@@ -596,6 +603,11 @@ public class MockHttpServletResponse implements HttpServletResponse {
 			this.locale = language != null ? language : Locale.getDefault();
 			return true;
 		}
+		else if (HttpHeaders.SET_COOKIE.equalsIgnoreCase(name)) {
+			MockCookie cookie = MockCookie.parse(value.toString());
+			addCookie(cookie);
+			return true;
+		}
 		else {
 			return false;
 		}

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

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2002-2018 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
+ *
+ *      http://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.mock.web;
+
+import org.junit.Test;
+
+import static org.junit.Assert.*;
+
+/**
+ * Unit tests for {@link MockCookie}.
+ *
+ * @author Vedran Pavic
+ */
+public class MockCookieTests {
+
+	@Test
+	public void constructCookie() {
+		MockCookie cookie = new MockCookie("SESSION", "123");
+
+		assertEquals("SESSION", cookie.getName());
+		assertEquals("123", cookie.getValue());
+	}
+
+	@Test
+	public void setSameSite() {
+		MockCookie cookie = new MockCookie("SESSION", "123");
+		cookie.setSameSite("Strict");
+
+		assertEquals("Strict", cookie.getSameSite());
+	}
+
+	@Test
+	public void parseValidHeader() {
+		MockCookie cookie = MockCookie.parse(
+				"SESSION=123; Domain=example.com; Max-Age=60; Path=/; Secure; HttpOnly; SameSite=Lax");
+
+		assertEquals("SESSION", cookie.getName());
+		assertEquals("123", cookie.getValue());
+		assertEquals("example.com", cookie.getDomain());
+		assertEquals(60, cookie.getMaxAge());
+		assertEquals("/", cookie.getPath());
+		assertTrue(cookie.getSecure());
+		assertTrue(cookie.isHttpOnly());
+		assertEquals("Lax", cookie.getSameSite());
+	}
+
+	@Test(expected = IllegalArgumentException.class)
+	public void parseInvalidHeader() {
+		MockCookie.parse("invalid");
+	}
+
+}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -322,4 +322,35 @@ public class MockHttpServletResponseTests {
 		assertEquals(HttpServletResponse.SC_NOT_FOUND, response.getStatus());
 	}
 
+	@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());
+	}
+
+	@Test
+	public void addMockCookie() {
+		MockCookie mockCookie = new MockCookie("SESSION", "123");
+		mockCookie.setPath("/");
+		mockCookie.setDomain("example.com");
+		mockCookie.setMaxAge(0);
+		mockCookie.setSecure(true);
+		mockCookie.setHttpOnly(true);
+		mockCookie.setSameSite("Lax");
+
+		response.addCookie(mockCookie);
+
+		assertEquals("SESSION=123; Path=/; Domain=example.com; Max-Age=0; " +
+				"Expires=Thu, 1 Jan 1970 00:00:00 GMT; Secure; HttpOnly; SameSite=Lax",
+				response.getHeader(HttpHeaders.SET_COOKIE));
+	}
+
 }

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

@@ -0,0 +1,126 @@
+/*
+ * Copyright 2002-2018 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
+ *
+ *      http://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.mock.web.test;
+
+import javax.servlet.http.Cookie;
+
+import org.springframework.lang.Nullable;
+
+/**
+ * A {@code Cookie} subclass with the additional cookie directives as defined in the
+ * <a href="https://tools.ietf.org/html/rfc6265">RFC 6265</a>.
+ *
+ * @author Vedran Pavic
+ * @since 5.1
+ */
+public class MockCookie extends Cookie {
+
+	private static final long serialVersionUID = 4312531139502726325L;
+
+	@Nullable
+	private String sameSite;
+
+	/**
+	 * Constructs a {@code MockCookie} instance with the specified name and value.
+	 *
+	 * @param name  the cookie name
+	 * @param value the cookie value
+	 * @see Cookie#Cookie(String, String)
+	 */
+	public MockCookie(String name, String value) {
+		super(name, value);
+	}
+
+	/**
+	 * Factory method create {@code MockCookie} instance from Set-Cookie header value.
+	 *
+	 * @param setCookieHeader the Set-Cookie header value
+	 * @return the created cookie instance
+	 */
+	public static MockCookie parse(String setCookieHeader) {
+		String[] cookieParts = setCookieHeader.split("\\s*=\\s*", 2);
+		if (cookieParts.length != 2) {
+			throw new IllegalArgumentException("Invalid Set-Cookie header value");
+		}
+		String name = cookieParts[0];
+		String[] valueAndDirectives = cookieParts[1].split("\\s*;\\s*", 2);
+		String value = valueAndDirectives[0];
+		String[] directives = valueAndDirectives[1].split("\\s*;\\s*");
+		String domain = null;
+		int maxAge = -1;
+		String path = null;
+		boolean secure = false;
+		boolean httpOnly = false;
+		String sameSite = null;
+		for (String directive : directives) {
+			if (directive.startsWith("Domain")) {
+				domain = directive.split("=")[1];
+			}
+			else if (directive.startsWith("Max-Age")) {
+				maxAge = Integer.parseInt(directive.split("=")[1]);
+			}
+			else if (directive.startsWith("Path")) {
+				path = directive.split("=")[1];
+			}
+			else if (directive.startsWith("Secure")) {
+				secure = true;
+			}
+			else if (directive.startsWith("HttpOnly")) {
+				httpOnly = true;
+			}
+			else if (directive.startsWith("SameSite")) {
+				sameSite = directive.split("=")[1];
+			}
+		}
+		MockCookie cookie = new MockCookie(name, value);
+		if (domain != null) {
+			cookie.setDomain(domain);
+		}
+		cookie.setMaxAge(maxAge);
+		cookie.setPath(path);
+		cookie.setSecure(secure);
+		cookie.setHttpOnly(httpOnly);
+		cookie.setSameSite(sameSite);
+		return cookie;
+	}
+
+	/**
+	 * Return the cookie "SameSite" attribute, or {@code null} if not set.
+	 * <p>
+	 * This limits the scope of the cookie such that it will only be attached to same site
+	 * requests if {@code "Strict"} or cross-site requests if {@code "Lax"}.
+	 *
+	 * @see <a href="https://tools.ietf.org/html/draft-ietf-httpbis-rfc6265bis#section-4.1.2.7">RFC6265 bis</a>
+	 */
+	@Nullable
+	public String getSameSite() {
+		return this.sameSite;
+	}
+
+	/**
+	 * Add the "SameSite" attribute to the cookie.
+	 * <p>
+	 * This limits the scope of the cookie such that it will only be attached to same site
+	 * requests if {@code "Strict"} or cross-site requests if {@code "Lax"}.
+	 *
+	 * @see <a href="https://tools.ietf.org/html/draft-ietf-httpbis-rfc6265bis#section-4.1.2.7">RFC6265 bis</a>
+	 */
+	public void setSameSite(@Nullable String sameSite) {
+		this.sameSite = sameSite;
+	}
+
+}

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

@@ -54,6 +54,7 @@ import org.springframework.web.util.WebUtils;
  * @author Juergen Hoeller
  * @author Rod Johnson
  * @author Brian Clozel
+ * @author Vedran Pavic
  * @since 1.0.2
  */
 public class MockHttpServletResponse implements HttpServletResponse {
@@ -353,6 +354,12 @@ public class MockHttpServletResponse implements HttpServletResponse {
 		if (cookie.isHttpOnly()) {
 			buf.append("; HttpOnly");
 		}
+		if (cookie instanceof MockCookie) {
+			MockCookie mockCookie = (MockCookie) cookie;
+			if (StringUtils.hasText(mockCookie.getSameSite())) {
+				buf.append("; SameSite=").append(mockCookie.getSameSite());
+			}
+		}
 		return buf.toString();
 	}
 
@@ -596,6 +603,11 @@ public class MockHttpServletResponse implements HttpServletResponse {
 			this.locale = language != null ? language : Locale.getDefault();
 			return true;
 		}
+		else if (HttpHeaders.SET_COOKIE.equalsIgnoreCase(name)) {
+			MockCookie cookie = MockCookie.parse(value.toString());
+			addCookie(cookie);
+			return true;
+		}
 		else {
 			return false;
 		}