Skip to content

Add cookie customizer to CookieRequestCache #15685

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 3, 2024
Merged

Add cookie customizer to CookieRequestCache #15685

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions web/src/main/java/org/springframework/security/web/savedrequest/CookieRequestCache.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)) {
Expand All @@ -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);
}

Expand Down Expand Up @@ -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;
}

}
32 changes: 24 additions & 8 deletions ...n/java/org/springframework/security/web/server/savedrequest/CookieServerRequestCache.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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.
Expand All @@ -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();
Expand All @@ -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) {
Expand Down
17 changes: 17 additions & 0 deletions web/src/test/java/org/springframework/security/web/savedrequest/CookieRequestCacheTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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());
}
Expand Down
16 changes: 16 additions & 0 deletions ...a/org/springframework/security/web/server/savedrequest/CookieServerRequestCacheTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
}

}