From 33e802ca6c8b488c5c654d5fac48084ea9908b41 Mon Sep 17 00:00:00 2001 From: Florian Bernard Date: Fri, 23 Aug 2024 18:55:50 +0200 Subject: [PATCH] Add cookie customizer to CookieRequestCache and CookieServerRequestCache Issue gh-15204 --- .../web/savedrequest/CookieRequestCache.java | 15 +++++++++ .../CookieServerRequestCache.java | 32 ++++++++++++++----- .../savedrequest/CookieRequestCacheTests.java | 17 ++++++++++ .../CookieServerRequestCacheTests.java | 16 ++++++++++ 4 files changed, 72 insertions(+), 8 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 index 8f245eb7ea0..a51960532f1 100644 --- 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 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 cookieCustomizer) { + Assert.notNull(cookieCustomizer, "cookieCustomizer cannot be null"); + this.cookieCustomizer = cookieCustomizer; + } + } diff --git a/web/src/main/java/org/springframework/security/web/server/savedrequest/CookieServerRequestCache.java b/web/src/main/java/org/springframework/security/web/server/savedrequest/CookieServerRequestCache.java index bbe16dce1f9..0c9302a70d5 100644 --- a/web/src/main/java/org/springframework/security/web/server/savedrequest/CookieServerRequestCache.java +++ b/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 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 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 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 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) { 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 index 71def1f3255..6b18e05002d 100644 --- 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 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()); } diff --git a/web/src/test/java/org/springframework/security/web/server/savedrequest/CookieServerRequestCacheTests.java b/web/src/test/java/org/springframework/security/web/server/savedrequest/CookieServerRequestCacheTests.java index 0c4ce943e0e..759579b0b34 100644 --- a/web/src/test/java/org/springframework/security/web/server/savedrequest/CookieServerRequestCacheTests.java +++ b/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 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"); + } + }