Save original request on oauth2Client filter

oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/web/server/OAuth2AuthorizationRequestRedirectWebFilter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2019 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -24,6 +24,8 @@
 import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
 import org.springframework.security.web.server.DefaultServerRedirectStrategy;
 import org.springframework.security.web.server.ServerRedirectStrategy;
+import org.springframework.security.web.server.savedrequest.ServerRequestCache;
+import org.springframework.security.web.server.savedrequest.WebSessionServerRequestCache;
 import org.springframework.util.Assert;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilter;
@@ -67,6 +69,7 @@ public class OAuth2AuthorizationRequestRedirectWebFilter implements WebFilter {
 	private final ServerOAuth2AuthorizationRequestResolver authorizationRequestResolver;
 	private ServerAuthorizationRequestRepository<OAuth2AuthorizationRequest> authorizationRequestRepository =
 		new WebSessionOAuth2ServerAuthorizationRequestRepository();
+	private ServerRequestCache requestCache = new WebSessionServerRequestCache();
 
 	/**
 	 * Constructs an {@code OAuth2AuthorizationRequestRedirectFilter} using the provided parameters.
@@ -98,11 +101,23 @@ public final void setAuthorizationRequestRepository(
 		this.authorizationRequestRepository = authorizationRequestRepository;
 	}
 
+	/**
+	 * The request cache to use to save the request before sending a redirect.
+	 * @param requestCache the cache to redirect to.
+	 */
+	public void setRequestCache(ServerRequestCache requestCache) {
+		Assert.notNull(requestCache, "requestCache cannot be null");
+		this.requestCache = requestCache;
+	}
+
 	@Override
 	public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
 		return this.authorizationRequestResolver.resolve(exchange)
 			.switchIfEmpty(chain.filter(exchange).then(Mono.empty()))
-			.onErrorResume(ClientAuthorizationRequiredException.class, e -> this.authorizationRequestResolver.resolve(exchange, e.getClientRegistrationId()))
+			.onErrorResume(ClientAuthorizationRequiredException.class, e -> {
+				return this.requestCache.saveRequest(exchange)
+					.then(this.authorizationRequestResolver.resolve(exchange, e.getClientRegistrationId()));
+			})
 			.flatMap(clientRegistration -> sendRedirectForAuthorization(exchange, clientRegistration));
 	}
 

oauth2/oauth2-client/src/test/java/org/springframework/security/oauth2/client/web/server/OAuth2AuthorizationRequestRedirectWebFilterTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2019 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,6 +26,7 @@
 import org.springframework.security.oauth2.client.registration.ReactiveClientRegistrationRepository;
 import org.springframework.security.oauth2.client.registration.TestClientRegistrations;
 import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
+import org.springframework.security.web.server.savedrequest.ServerRequestCache;
 import org.springframework.test.web.reactive.server.FluxExchangeResult;
 import org.springframework.test.web.reactive.server.WebTestClient;
 import org.springframework.web.server.handler.FilteringWebHandler;
@@ -53,6 +54,9 @@ public class OAuth2AuthorizationRequestRedirectWebFilterTests {
 	@Mock
 	private ServerAuthorizationRequestRepository<OAuth2AuthorizationRequest> authzRequestRepository;
 
+	@Mock
+	private ServerRequestCache requestCache;
+
 	private ClientRegistration registration = TestClientRegistrations.clientRegistration().build();
 
 	private OAuth2AuthorizationRequestRedirectWebFilter filter;
@@ -139,4 +143,33 @@ public void filterWhenExceptionThenRedirected() {
 				.is3xxRedirection()
 				.returnResult(String.class);
 	}
+
+	@Test
+	public void filterWhenExceptionThenSaveRequestSessionAttribute() {
+		this.filter.setRequestCache(this.requestCache);
+		when(this.requestCache.saveRequest(any())).thenReturn(Mono.empty());
+		FilteringWebHandler webHandler = new FilteringWebHandler(
+				e -> Mono.error(new ClientAuthorizationRequiredException(this.registration.getRegistrationId())),
+				Arrays.asList(this.filter));
+		this.client = WebTestClient.bindToWebHandler(webHandler).build();
+		this.client.get()
+				.uri("https://example.com/foo")
+				.exchange()
+				.expectStatus()
+				.is3xxRedirection()
+				.returnResult(String.class);
+		verify(this.requestCache).saveRequest(any());
+	}
+
+	@Test
+	public void filterWhenPathMatchesThenRequestSessionAttributeNotSaved() {
+		this.filter.setRequestCache(this.requestCache);
+		this.client.get()
+				.uri("https://example.com/oauth2/authorization/registration-id")
+				.exchange()
+				.expectStatus()
+				.is3xxRedirection()
+				.returnResult(String.class);
+		verifyZeroInteractions(this.requestCache);
+	}
 }