Add cookie customizer to CookieRequestCache and CookieServerRequestCache

Issue gh-15204
This commit is contained in:
Florian Bernard 2024-08-23 18:55:50 +02:00 committed by Marcus Hert Da Coregio
parent 820ce4ea7a
commit 008cbc2cae
4 changed files with 72 additions and 8 deletions

View File

@ -18,6 +18,7 @@ package org.springframework.security.web.savedrequest;
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 class CookieRequestCache implements RequestCache {
savedCookie.setSecure(request.isSecure());
savedCookie.setPath(getCookiePath(request));
savedCookie.setHttpOnly(true);
this.cookieCustomizer.accept(savedCookie);
response.addCookie(savedCookie);
}
@ -152,4 +157,14 @@ public class CookieRequestCache implements RequestCache {
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;
}
}

View File

@ -20,6 +20,7 @@ import java.net.URI;
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 class CookieServerRequestCache implements ServerRequestCache {
.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 class CookieServerRequestCache implements ServerRequestCache {
.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) {

View File

@ -20,6 +20,7 @@ import java.util.Arrays;
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 class CookieRequestCacheTests {
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());
}

View File

@ -138,4 +138,20 @@ public class CookieServerRequestCacheTests {
"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");
}
}