Add ServerRequestCache

Fixes: gh-4789

web/src/main/java/org/springframework/security/web/server/authentication/RedirectServerAuthenticationEntryPoint.java

@@ -20,6 +20,8 @@ import java.net.URI;
 
 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 reactor.core.publisher.Mono;
 
 import org.springframework.security.core.AuthenticationException;
@@ -39,14 +41,22 @@ public class RedirectServerAuthenticationEntryPoint
 
 	private ServerRedirectStrategy serverRedirectStrategy = new DefaultServerRedirectStrategy();
 
+	private ServerRequestCache requestCache = new WebSessionServerRequestCache();
+
 	public RedirectServerAuthenticationEntryPoint(String location) {
 		Assert.notNull(location, "location cannot be null");
 		this.location = URI.create(location);
 	}
 
+	public void setRequestCache(ServerRequestCache requestCache) {
+		Assert.notNull(requestCache, "requestCache cannot be null");
+		this.requestCache = requestCache;
+	}
+
 	@Override
 	public Mono<Void> commence(ServerWebExchange exchange, AuthenticationException e) {
-		return this.serverRedirectStrategy.sendRedirect(exchange, this.location);
+		return this.requestCache.saveRequest(exchange)
+			.then(this.serverRedirectStrategy.sendRedirect(exchange, this.location));
 	}
 
 	/**

web/src/main/java/org/springframework/security/web/server/authentication/RedirectServerAuthenticationSuccessHandler.java

@@ -20,6 +20,8 @@ import org.springframework.security.core.Authentication;
 import org.springframework.security.web.server.DefaultServerRedirectStrategy;
 import org.springframework.security.web.server.ServerRedirectStrategy;
 import org.springframework.security.web.server.WebFilterExchange;
+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 reactor.core.publisher.Mono;
@@ -36,17 +38,28 @@ public class RedirectServerAuthenticationSuccessHandler
 
 	private ServerRedirectStrategy serverRedirectStrategy = new DefaultServerRedirectStrategy();
 
+	private ServerRequestCache requestCache = new WebSessionServerRequestCache();
+
 	public RedirectServerAuthenticationSuccessHandler() {}
 
 	public RedirectServerAuthenticationSuccessHandler(String location) {
 		this.location = URI.create(location);
 	}
 
+	public void setRequestCache(ServerRequestCache requestCache) {
+		Assert.notNull(requestCache, "requestCache cannot be null");
+		this.requestCache = requestCache;
+	}
+
 	@Override
 	public Mono<Void> onAuthenticationSuccess(WebFilterExchange webFilterExchange,
 		Authentication authentication) {
 		ServerWebExchange exchange = webFilterExchange.getExchange();
-		return this.serverRedirectStrategy.sendRedirect(exchange, this.location);
+		return this.requestCache.getRequest(exchange)
+			.map(r -> r.getPath().pathWithinApplication().value())
+			.map(URI::create)
+			.defaultIfEmpty(this.location)
+			.flatMap(location -> this.serverRedirectStrategy.sendRedirect(exchange, location));
 	}
 
 	/**

web/src/main/java/org/springframework/security/web/server/savedrequest/NoOpServerRequestCache.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2002-2017 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.security.web.server.savedrequest;
+
+import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.web.server.ServerWebExchange;
+import reactor.core.publisher.Mono;
+
+/**
+ * @author Rob Winch
+ * @since 5.0
+ */
+public class NoOpServerRequestCache implements ServerRequestCache {
+	@Override
+	public Mono<Void> saveRequest(ServerWebExchange exchange) {
+		return Mono.empty();
+	}
+
+	@Override
+	public Mono<ServerHttpRequest> getRequest(ServerWebExchange exchange) {
+		return Mono.empty();
+	}
+
+	@Override
+	public Mono<ServerHttpRequest> getMatchingRequest(
+		ServerWebExchange exchange) {
+		return Mono.empty();
+	}
+
+	@Override
+	public Mono<ServerHttpRequest> removeRequest(ServerWebExchange exchange) {
+		return Mono.empty();
+	}
+
+	public static NoOpServerRequestCache getInstance() {
+		return new NoOpServerRequestCache();
+	}
+
+	private NoOpServerRequestCache() {}
+}

web/src/main/java/org/springframework/security/web/server/savedrequest/ServerRequestCache.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2002-2017 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.security.web.server.savedrequest;
+
+import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.web.server.ServerWebExchange;
+import reactor.core.publisher.Mono;
+
+/**
+ * Saves a {@link ServerHttpRequest} so it can be "replayed" later. This is useful for
+ * when a page was requested and authentication is necessary.
+ *
+ * @author Rob Winch
+ * @since 5.0
+ */
+public interface ServerRequestCache {
+
+	/**
+	 * Save the {@link ServerHttpRequest}
+	 * @param exchange the exchange to save
+	 * @return Return a {@code Mono<Void>} which only replays complete and error signals
+	 * from this {@link Mono}.
+	 */
+	Mono<Void> saveRequest(ServerWebExchange exchange);
+
+	/**
+	 * Get the saved {@link ServerHttpRequest}
+	 * @param exchange the exchange to obtain the saved {@link ServerHttpRequest} from
+	 * @return the {@link ServerHttpRequest}
+	 */
+	Mono<ServerHttpRequest> getRequest(ServerWebExchange exchange);
+
+	/**
+	 * If the provided {@link ServerWebExchange} matches the saved {@link ServerHttpRequest}
+	 * gets the saved {@link ServerHttpRequest}
+	 * @param exchange the exchange to obtain the request from
+	 * @return the {@link ServerHttpRequest}
+	 */
+	Mono<ServerHttpRequest> getMatchingRequest(ServerWebExchange exchange);
+
+	/**
+	 * If the {@link ServerWebExchange} contains a saved {@link ServerHttpRequest} remove
+	 * and return it.
+	 *
+	 * @param exchange the {@link ServerWebExchange} to obtain and remove the
+	 * {@link ServerHttpRequest}
+	 * @return the {@link ServerHttpRequest}
+	 */
+	Mono<ServerHttpRequest> removeRequest(ServerWebExchange exchange);
+}

web/src/main/java/org/springframework/security/web/server/savedrequest/ServerRequestCacheWebFilter.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2002-2017 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.security.web.server.savedrequest;
+
+import org.springframework.util.Assert;
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.WebFilter;
+import org.springframework.web.server.WebFilterChain;
+import reactor.core.publisher.Mono;
+
+/**
+ * A {@link WebFilter} that replays any matching request in {@link ServerRequestCache}
+ *
+ * @author Rob Winch
+ * @since 5.0
+ */
+public class ServerRequestCacheWebFilter implements WebFilter {
+	private ServerRequestCache requestCache = new WebSessionServerRequestCache();
+
+	@Override
+	public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
+		return this.requestCache.getMatchingRequest(exchange)
+			.flatMap(r -> this.requestCache.removeRequest(exchange))
+			.map(r -> exchange.mutate().request(r).build())
+			.defaultIfEmpty(exchange)
+			.flatMap(e -> chain.filter(e));
+	}
+
+	public void setRequestCache(ServerRequestCache requestCache) {
+		Assert.notNull(requestCache, "requestCache cannot be null");
+		this.requestCache = requestCache;
+	}
+}

web/src/main/java/org/springframework/security/web/server/savedrequest/WebSessionServerRequestCache.java

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2002-2017 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.security.web.server.savedrequest;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.http.HttpMethod;
+import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.security.web.server.util.matcher.ServerWebExchangeMatcher;
+import org.springframework.security.web.server.util.matcher.ServerWebExchangeMatchers;
+import org.springframework.util.Assert;
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.WebSession;
+import reactor.core.publisher.Mono;
+
+import java.net.URI;
+
+/**
+ * An implementation of {@link ServerRequestCache} that saves the
+ * {@link ServerHttpRequest} in the {@link WebSession}.
+ *
+ * The current implementation only saves the URL that was requested.
+ *
+ * @author Rob Winch
+ * @since 5.0
+ */
+public class WebSessionServerRequestCache implements ServerRequestCache {
+	private static final String DEFAULT_SAVED_REQUEST_ATTR = "SPRING_SECURITY_SAVED_REQUEST";
+
+	protected final Log logger = LogFactory.getLog(this.getClass());
+
+	private String sessionAttrName = DEFAULT_SAVED_REQUEST_ATTR;
+
+	private ServerWebExchangeMatcher saveRequestMatcher = ServerWebExchangeMatchers.pathMatchers(
+		HttpMethod.GET, "/**");
+
+	/**
+	 * Sets the matcher to determine if the request should be saved. The default is to match
+	 * on any GET request.
+	 *
+	 * @param saveRequestMatcher
+	 */
+	public void setSaveRequestMatcher(ServerWebExchangeMatcher saveRequestMatcher) {
+		Assert.notNull(saveRequestMatcher, "saveRequestMatcher cannot be null");
+		this.saveRequestMatcher = saveRequestMatcher;
+	}
+
+	@Override
+	public Mono<Void> saveRequest(ServerWebExchange exchange) {
+		return this.saveRequestMatcher.matches(exchange)
+			.filter(m -> m.isMatch())
+			.flatMap(m -> exchange.getSession())
+			.map(WebSession::getAttributes)
+			.doOnNext(attrs -> attrs.put(this.sessionAttrName, pathInApplication(exchange.getRequest())))
+			.then();
+	}
+
+	@Override
+	public Mono<ServerHttpRequest> getRequest(ServerWebExchange exchange) {
+		return exchange.getSession()
+			.flatMap(session -> Mono.justOrEmpty(session.<String>getAttribute(this.sessionAttrName)))
+			.map(path -> exchange.getRequest().mutate().path(path).build());
+	}
+
+	@Override
+	public Mono<ServerHttpRequest> getMatchingRequest(
+		ServerWebExchange exchange) {
+		return getRequest(exchange)
+			.filter( request -> pathInApplication(request).equals(
+				pathInApplication(exchange.getRequest())));
+	}
+
+	@Override
+	public Mono<ServerHttpRequest> removeRequest(ServerWebExchange exchange) {
+		return exchange.getSession()
+			.map(WebSession::getAttributes)
+			.flatMap(attrs -> Mono.justOrEmpty(attrs.remove(this.sessionAttrName)))
+			.cast(String.class)
+			.map(path -> exchange.getRequest().mutate().path(path).build());
+	}
+
+	private static String pathInApplication(ServerHttpRequest request) {
+		return request.getPath().pathWithinApplication().value();
+	}
+}

web/src/test/java/org/springframework/security/web/server/savedrequest/WebSessionServerRequestCacheTests.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright 2002-2017 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.security.web.server.savedrequest;
+
+import org.junit.Test;
+import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
+import org.springframework.mock.web.server.MockServerWebExchange;
+import org.springframework.security.web.server.util.matcher.ServerWebExchangeMatcher;
+
+import static org.assertj.core.api.Assertions.*;
+
+/**
+ * @author Rob Winch
+ * @since 5.0
+ */
+public class WebSessionServerRequestCacheTests {
+	private WebSessionServerRequestCache cache = new WebSessionServerRequestCache();
+
+	@Test
+	public void saveRequestGetRequestWhenGetThenFound() {
+		MockServerWebExchange exchange = MockServerWebExchange.from(MockServerHttpRequest.get("/secured/"));
+		this.cache.saveRequest(exchange).block();
+
+		ServerHttpRequest saved = this.cache.getRequest(exchange).block();
+
+		assertThat(saved.getURI()).isEqualTo(exchange.getRequest().getURI());
+	}
+
+	@Test
+	public void saveRequestGetRequestWhenPostThenNotFound() {
+		MockServerWebExchange exchange = MockServerWebExchange.from(MockServerHttpRequest.post("/secured/"));
+		this.cache.saveRequest(exchange).block();
+
+		assertThat(this.cache.getRequest(exchange).block()).isNull();
+	}
+
+	@Test
+	public void saveRequestGetRequestWhenPostAndCustomMatcherThenFound() {
+		this.cache.setSaveRequestMatcher(e -> ServerWebExchangeMatcher.MatchResult.match());
+		MockServerWebExchange exchange = MockServerWebExchange.from(MockServerHttpRequest.post("/secured/"));
+		this.cache.saveRequest(exchange).block();
+
+		ServerHttpRequest saved = this.cache.getRequest(exchange).block();
+
+		assertThat(saved.getURI()).isEqualTo(exchange.getRequest().getURI());
+	}
+
+	@Test
+	public void saveRequestRemoveRequestWhenThenFound() {
+		MockServerWebExchange exchange = MockServerWebExchange.from(MockServerHttpRequest.get("/secured/"));
+		this.cache.saveRequest(exchange).block();
+
+		ServerHttpRequest saved = this.cache.removeRequest(exchange).block();
+
+		assertThat(saved.getURI()).isEqualTo(exchange.getRequest().getURI());
+	}
+
+	@Test
+	public void removeRequestGetRequestWhenDefaultThenNotFound() {
+		MockServerWebExchange exchange = MockServerWebExchange.from(MockServerHttpRequest.get("/secured/"));
+		this.cache.saveRequest(exchange).block();
+
+		this.cache.removeRequest(exchange).block();
+
+		assertThat(this.cache.getRequest(exchange).block()).isNull();
+	}
+}