Add cookie customizer to CookieRequestCache and CookieServerRequestCache

Issue gh-15204

Author: fb64

web/src/main/java/org/springframework/security/web/savedrequest/CookieRequestCache.java

@@ -18,6 +18,7 @@
 
 import java.util.Base64;
 import java.util.Collections;
+import java.util.function.Consumer;
 
 import jakarta.servlet.http.Cookie;
 import jakarta.servlet.http.HttpServletRequest;
@@ -51,6 +52,9 @@ public class CookieRequestCache implements RequestCache {
 
 	private static final int COOKIE_MAX_AGE = -1;
 
+	private Consumer<Cookie> cookieCustomizer = (cookie) -> {
+	};
+
 	@Override
 	public void saveRequest(HttpServletRequest request, HttpServletResponse response) {
 		if (!this.requestMatcher.matches(request)) {
@@ -63,6 +67,7 @@ public void saveRequest(HttpServletRequest request, HttpServletResponse response
 		savedCookie.setSecure(request.isSecure());
 		savedCookie.setPath(getCookiePath(request));
 		savedCookie.setHttpOnly(true);
+		this.cookieCustomizer.accept(savedCookie);
 		response.addCookie(savedCookie);
 	}
 
@@ -152,4 +157,14 @@ public void setRequestMatcher(RequestMatcher requestMatcher) {
 		this.requestMatcher = requestMatcher;
 	}
 
+	/**
+	 * Sets the {@link Consumer}, allowing customization of cookie.
+	 * @param cookieCustomizer customize for cookie
+	 * @since 6.4
+	 */
+	public void setCookieCustomizer(Consumer<Cookie> cookieCustomizer) {
+		Assert.notNull(cookieCustomizer, "cookieCustomizer cannot be null");
+		this.cookieCustomizer = cookieCustomizer;
+	}
+
 }

web/src/main/java/org/springframework/security/web/server/savedrequest/CookieServerRequestCache.java

@@ -20,6 +20,7 @@
 import java.time.Duration;
 import java.util.Base64;
 import java.util.Collections;
+import java.util.function.Consumer;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -59,6 +60,9 @@ public class CookieServerRequestCache implements ServerRequestCache {
 
 	private ServerWebExchangeMatcher saveRequestMatcher = createDefaultRequestMatcher();
 
+	private Consumer<ResponseCookie.ResponseCookieBuilder> cookieCustomizer = (cookieBuilder) -> {
+	};
+
 	/**
 	 * Sets the matcher to determine if the request should be saved. The default is to
 	 * match on any GET request.
@@ -77,8 +81,10 @@ public Mono<Void> saveRequest(ServerWebExchange exchange) {
 			.map((m) -> exchange.getResponse())
 			.map(ServerHttpResponse::getCookies)
 			.doOnNext((cookies) -> {
-				ResponseCookie redirectUriCookie = createRedirectUriCookie(exchange.getRequest());
-				cookies.add(REDIRECT_URI_COOKIE_NAME, redirectUriCookie);
+				ResponseCookie.ResponseCookieBuilder redirectUriCookie = createRedirectUriCookieBuilder(
+						exchange.getRequest());
+				this.cookieCustomizer.accept(redirectUriCookie);
+				cookies.add(REDIRECT_URI_COOKIE_NAME, redirectUriCookie.build());
 				logger.debug(LogMessage.format("Request added to Cookie: %s", redirectUriCookie));
 			})
 			.then();
@@ -103,25 +109,35 @@ public Mono<ServerHttpRequest> removeMatchingRequest(ServerWebExchange exchange)
 			.thenReturn(exchange.getRequest());
 	}
 
-	private static ResponseCookie createRedirectUriCookie(ServerHttpRequest request) {
+	/**
+	 * Sets the {@link Consumer}, allowing customization of cookie.
+	 * @param cookieCustomizer customize for cookie
+	 * @since 6.4
+	 */
+	public void setCookieCustomizer(Consumer<ResponseCookie.ResponseCookieBuilder> cookieCustomizer) {
+		Assert.notNull(cookieCustomizer, "cookieCustomizer cannot be null");
+		this.cookieCustomizer = cookieCustomizer;
+	}
+
+	private static ResponseCookie.ResponseCookieBuilder createRedirectUriCookieBuilder(ServerHttpRequest request) {
 		String path = request.getPath().pathWithinApplication().value();
 		String query = request.getURI().getRawQuery();
 		String redirectUri = path + ((query != null) ? "?" + query : "");
-		return createResponseCookie(request, encodeCookie(redirectUri), COOKIE_MAX_AGE);
+		return createResponseCookieBuilder(request, encodeCookie(redirectUri), COOKIE_MAX_AGE);
 	}
 
 	private static ResponseCookie invalidateRedirectUriCookie(ServerHttpRequest request) {
-		return createResponseCookie(request, null, Duration.ZERO);
+		return createResponseCookieBuilder(request, null, Duration.ZERO).build();
 	}
 
-	private static ResponseCookie createResponseCookie(ServerHttpRequest request, String cookieValue, Duration age) {
+	private static ResponseCookie.ResponseCookieBuilder createResponseCookieBuilder(ServerHttpRequest request,
+			String cookieValue, Duration age) {
 		return ResponseCookie.from(REDIRECT_URI_COOKIE_NAME, cookieValue)
 			.path(request.getPath().contextPath().value() + "/")
 			.maxAge(age)
 			.httpOnly(true)
 			.secure("https".equalsIgnoreCase(request.getURI().getScheme()))
-			.sameSite("Lax")
-			.build();
+			.sameSite("Lax");
 	}
 
 	private static String encodeCookie(String cookieValue) {

web/src/test/java/org/springframework/security/web/savedrequest/CookieRequestCacheTests.java

@@ -20,6 +20,7 @@
 import java.util.Base64;
 import java.util.Collections;
 import java.util.Locale;
+import java.util.function.Consumer;
 
 import jakarta.servlet.http.Cookie;
 import jakarta.servlet.http.HttpServletRequest;
@@ -204,6 +205,22 @@ public void matchingRequestWhenMatchThenKeepOriginalRequestLocale() {
 		assertThat(Collections.list(matchingRequest.getLocales())).contains(Locale.FRENCH, Locale.GERMANY);
 	}
 
+	@Test
+	public void setCookieCustomizer() {
+		Consumer<Cookie> cookieCustomizer = (cookie) -> {
+			cookie.setAttribute("SameSite", "Strict");
+			cookie.setAttribute("CustomAttribute", "CustomValue");
+		};
+		CookieRequestCache cookieRequestCache = new CookieRequestCache();
+		cookieRequestCache.setCookieCustomizer(cookieCustomizer);
+		MockHttpServletResponse response = new MockHttpServletResponse();
+		cookieRequestCache.saveRequest(new MockHttpServletRequest(), response);
+		Cookie savedCookie = response.getCookie(DEFAULT_COOKIE_NAME);
+		assertThat(savedCookie).isNotNull();
+		assertThat(savedCookie.getAttribute("SameSite")).isEqualTo("Strict");
+		assertThat(savedCookie.getAttribute("CustomAttribute")).isEqualTo("CustomValue");
+	}
+
 	private static String encodeCookie(String cookieValue) {
 		return Base64.getEncoder().encodeToString(cookieValue.getBytes());
 	}

web/src/test/java/org/springframework/security/web/server/savedrequest/CookieServerRequestCacheTests.java

@@ -138,4 +138,20 @@ public void removeMatchingRequestThenRedirectUriCookieExpired() {
 				"REDIRECT_URI=; Path=/; Max-Age=0; Expires=Thu, 01 Jan 1970 00:00:00 GMT; HttpOnly; SameSite=Lax");
 	}
 
+	@Test
+	public void saveRequestWithCookieCustomizerThenSameSiteStrict() {
+		MockServerWebExchange exchange = MockServerWebExchange
+			.from(MockServerHttpRequest.get("/secured/").accept(MediaType.TEXT_HTML));
+		CookieServerRequestCache cacheWithCustomizer = new CookieServerRequestCache();
+		cacheWithCustomizer.setCookieCustomizer(((cookieBuilder) -> cookieBuilder.sameSite("Strict")));
+		cacheWithCustomizer.saveRequest(exchange).block();
+		MultiValueMap<String, ResponseCookie> cookies = exchange.getResponse().getCookies();
+		assertThat(cookies).hasSize(1);
+		ResponseCookie cookie = cookies.getFirst("REDIRECT_URI");
+		assertThat(cookie).isNotNull();
+		String encodedRedirectUrl = Base64.getEncoder().encodeToString("/secured/".getBytes());
+		assertThat(cookie.toString())
+			.isEqualTo("REDIRECT_URI=" + encodedRedirectUrl + "; Path=/; HttpOnly; SameSite=Strict");
+	}
+
 }