Add cookie customizer to CookieRequestCache

Issue spring-projectsgh-15286
fb64 · Aug 23, 2024 · db19bc7 · db19bc7
1 parent e90a6b6
commit db19bc7
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 0 deletions.
diff --git a/web/src/main/java/org/springframework/security/web/savedrequest/CookieRequestCache.java b/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,13 @@ public void setRequestMatcher(RequestMatcher requestMatcher) {
 		this.requestMatcher = requestMatcher;
 	}
 
+	/**
+	 * Sets the {@link Consumer}, allowing customization of cookie.
+	 * @param cookieCustomizer customize for cookie
+	 */
+	public void setCookieCustomizer(Consumer<Cookie> cookieCustomizer) {
+		Assert.notNull(cookieCustomizer, "cookieCustomizer cannot be null");
+		this.cookieCustomizer = cookieCustomizer;
+	}
+
 }
diff --git a/web/src/test/java/org/springframework/security/web/savedrequest/CookieRequestCacheTests.java b/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());
 	}