From 4392d4790887ca8efcbb52bb02f90545736c2c19 Mon Sep 17 00:00:00 2001
From: Rob Winch <rwinch@users.noreply.github.com>
Date: Tue, 5 Sep 2017 15:16:48 -0500
Subject: [PATCH] RedirectAuthenticationEntryPoint uses RedirectStrategy

Issue gh-4529
---
 .../RedirectAuthenticationEntryPoint.java     | 21 ++++++++-----------
 ...RedirectAuthenticationEntryPointTests.java | 19 ++++++++++-------
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/webflux/src/main/java/org/springframework/security/web/server/authentication/RedirectAuthenticationEntryPoint.java b/webflux/src/main/java/org/springframework/security/web/server/authentication/RedirectAuthenticationEntryPoint.java
index 823494cd37..3d01b2a622 100644
--- a/webflux/src/main/java/org/springframework/security/web/server/authentication/RedirectAuthenticationEntryPoint.java
+++ b/webflux/src/main/java/org/springframework/security/web/server/authentication/RedirectAuthenticationEntryPoint.java
@@ -20,6 +20,8 @@ package org.springframework.security.web.server.authentication;
 
 import java.net.URI;
 
+import org.springframework.security.web.server.DefaultRedirectStrategy;
+import org.springframework.security.web.server.RedirectStrategy;
 import reactor.core.publisher.Mono;
 
 import org.springframework.http.HttpStatus;
@@ -38,7 +40,7 @@ import org.springframework.web.server.ServerWebExchange;
 public class RedirectAuthenticationEntryPoint implements AuthenticationEntryPoint {
 	private final URI location;
 
-	private HttpStatus httpStatus = HttpStatus.FOUND;
+	private RedirectStrategy redirectStrategy = new DefaultRedirectStrategy();
 
 	public RedirectAuthenticationEntryPoint(String location) {
 		Assert.notNull(location, "location cannot be null");
@@ -47,20 +49,15 @@ public class RedirectAuthenticationEntryPoint implements AuthenticationEntryPoin
 
 	@Override
 	public Mono<Void> commence(ServerWebExchange exchange, AuthenticationException e) {
-		return Mono.fromRunnable(() -> {
-			ServerHttpResponse response = exchange.getResponse();
-			response.setStatusCode(this.httpStatus);
-			response.getHeaders().setLocation(this.location);
-		});
+		return this.redirectStrategy.sendRedirect(exchange, this.location);
 	}
 
 	/**
-	 * Sets the {@link HttpStatus}.
-	 *
-	 * @param httpStatus the status to use. The default is {@code HttpStatus.FOUND}
+	 * Sets the RedirectStrategy to use.
+	 * @param redirectStrategy the strategy to use. Default is DefaultRedirectStrategy.
 	 */
-	public void setHttpStatus(HttpStatus httpStatus) {
-		Assert.notNull(httpStatus, "httpStatus cannot be null");
-		this.httpStatus = httpStatus;
+	public void setRedirectStrategy(RedirectStrategy redirectStrategy) {
+		Assert.notNull(redirectStrategy, "redirectStrategy cannot be null");
+		this.redirectStrategy = redirectStrategy;
 	}
 }
diff --git a/webflux/src/test/java/org/springframework/security/web/server/authentication/RedirectAuthenticationEntryPointTests.java b/webflux/src/test/java/org/springframework/security/web/server/authentication/RedirectAuthenticationEntryPointTests.java
index d408a76669..b1909848b8 100644
--- a/webflux/src/test/java/org/springframework/security/web/server/authentication/RedirectAuthenticationEntryPointTests.java
+++ b/webflux/src/test/java/org/springframework/security/web/server/authentication/RedirectAuthenticationEntryPointTests.java
@@ -27,10 +27,14 @@ import org.springframework.http.HttpStatus;
 import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
 import org.springframework.security.authentication.AuthenticationCredentialsNotFoundException;
 import org.springframework.security.core.AuthenticationException;
+import org.springframework.security.web.server.RedirectStrategy;
 import org.springframework.web.server.ServerWebExchange;
+import reactor.core.publisher.Mono;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.mockito.Mockito.when;
 
 /**
  * @author Rob Winch
@@ -41,6 +45,8 @@ public class RedirectAuthenticationEntryPointTests {
 
 	@Mock
 	private ServerWebExchange exchange;
+	@Mock
+	private RedirectStrategy redirectStrategy;
 
 	private String location = "/login";
 
@@ -76,18 +82,17 @@ public class RedirectAuthenticationEntryPointTests {
 
 	@Test
 	public void commenceWhenCustomStatusThenStatusSet() {
+		Mono<Void> result = Mono.empty();
+		when(this.redirectStrategy.sendRedirect(any(), any())).thenReturn(result);
 		HttpStatus status = HttpStatus.MOVED_PERMANENTLY;
-		this.entryPoint.setHttpStatus(status);
+		this.entryPoint.setRedirectStrategy(this.redirectStrategy);
 		this.exchange = MockServerHttpRequest.get("/").toExchange();
 
-		this.entryPoint.commence(this.exchange, this.exception).block();
-
-		assertThat(this.exchange.getResponse().getStatusCode()).isEqualTo(status);
-		assertThat(this.exchange.getResponse().getHeaders().getLocation()).hasPath(this.location);
+		assertThat(this.entryPoint.commence(this.exchange, this.exception)).isEqualTo(result);
 	}
 
 	@Test(expected = IllegalArgumentException.class)
-	public void setHttpStatusWhenNullLocationThenException() {
-		this.entryPoint.setHttpStatus(null);
+	public void setRedirectStrategyWhenNullThenException() {
+		this.entryPoint.setRedirectStrategy(null);
 	}
 }