AuthenticationSuccessHandler uses WebFilterExchange

Issue gh-4532

webflux/src/main/java/org/springframework/security/web/server/authentication/AuthenticationSuccessHandler.java

@@ -19,8 +19,7 @@
 package org.springframework.security.web.server.authentication;
 
 import org.springframework.security.core.Authentication;
-import org.springframework.web.server.ServerWebExchange;
-import org.springframework.web.server.WebFilterChain;
+import org.springframework.security.web.server.WebFilterExchange;
 import reactor.core.publisher.Mono;
 
 /**
@@ -28,5 +27,5 @@ import reactor.core.publisher.Mono;
  * @since 5.0
  */
 public interface AuthenticationSuccessHandler {
-	Mono<Void> success(Authentication authentication, ServerWebExchange exchange, WebFilterChain chain);
+	Mono<Void> success(Authentication authentication, WebFilterExchange webFilterExchange);
 }

webflux/src/main/java/org/springframework/security/web/server/authentication/AuthenticationWebFilter.java

@@ -27,6 +27,7 @@ import org.springframework.security.core.Authentication;
 import org.springframework.security.core.context.SecurityContextImpl;
 import org.springframework.security.web.server.AuthenticationEntryPoint;
 import org.springframework.security.web.server.HttpBasicAuthenticationConverter;
+import org.springframework.security.web.server.WebFilterExchange;
 import org.springframework.security.web.server.authentication.www.HttpBasicAuthenticationEntryPoint;
 import org.springframework.security.web.server.context.SecurityContextRepository;
 import org.springframework.security.web.server.context.SecurityContextRepositoryServerWebExchange;
@@ -87,7 +88,7 @@ public class AuthenticationWebFilter implements WebFilter {
 		SecurityContextImpl securityContext = new SecurityContextImpl();
 		securityContext.setAuthentication(authentication);
 		return this.securityContextRepository.save(exchange, securityContext)
-			.then(this.authenticationSuccessHandler.success(authentication, exchange, chain));
+			.then(this.authenticationSuccessHandler.success(authentication, new WebFilterExchange(exchange, chain)));
 	}
 
 	public void setSecurityContextRepository(

webflux/src/main/java/org/springframework/security/web/server/authentication/RedirectAuthenticationSuccessHandler.java

@@ -21,6 +21,7 @@ package org.springframework.security.web.server.authentication;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.web.server.DefaultRedirectStrategy;
 import org.springframework.security.web.server.RedirectStrategy;
+import org.springframework.security.web.server.WebFilterExchange;
 import org.springframework.util.Assert;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilterChain;
@@ -38,8 +39,8 @@ public class RedirectAuthenticationSuccessHandler implements AuthenticationSucce
 	private RedirectStrategy redirectStrategy = new DefaultRedirectStrategy();
 
 	@Override
-	public Mono<Void> success(Authentication authentication,
-		ServerWebExchange exchange, WebFilterChain chain) {
+	public Mono<Void> success(Authentication authentication, WebFilterExchange webFilterExchange) {
+		ServerWebExchange exchange = webFilterExchange.getExchange();
 		return this.redirectStrategy.sendRedirect(exchange, this.location);
 	}
 

webflux/src/main/java/org/springframework/security/web/server/authentication/WebFilterChainAuthenticationSuccessHandler.java

@@ -19,6 +19,7 @@
 package org.springframework.security.web.server.authentication;
 
 import org.springframework.security.core.Authentication;
+import org.springframework.security.web.server.WebFilterExchange;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilterChain;
 import reactor.core.publisher.Mono;
@@ -29,7 +30,8 @@ import reactor.core.publisher.Mono;
  */
 public class WebFilterChainAuthenticationSuccessHandler implements AuthenticationSuccessHandler {
 	@Override
-	public Mono<Void> success(Authentication authentication, ServerWebExchange exchange, WebFilterChain chain) {
-		return chain.filter(exchange);
+	public Mono<Void> success(Authentication authentication, WebFilterExchange webFilterExchange) {
+		ServerWebExchange exchange = webFilterExchange.getExchange();
+		return webFilterExchange.getChain().filter(exchange);
 	}
 }

webflux/src/test/java/org/springframework/security/web/server/authentication/AuthenticationWebFilterTests.java

@@ -150,7 +150,7 @@ public class AuthenticationWebFilterTests {
 			.bindToWebFilters(this.filter)
 			.build();
 
-		EntityExchangeResult<String> result = client
+		client
 			.get()
 			.uri("/")
 			.exchange()
@@ -188,7 +188,7 @@ public class AuthenticationWebFilterTests {
 		Mono<Authentication> authentication = Mono.just(new TestingAuthenticationToken("test", "this", "ROLE_USER"));
 		when(this.authenticationConverter.apply(any())).thenReturn(authentication);
 		when(this.authenticationManager.authenticate(any())).thenReturn(authentication);
-		when(this.successHandler.success(any(),any(),any())).thenReturn(Mono.empty());
+		when(this.successHandler.success(any(),any())).thenReturn(Mono.empty());
 		when(this.securityContextRepository.save(any(),any())).thenAnswer( a -> Mono.just(a.getArguments()[0]));
 
 		WebTestClient client = WebTestClientBuilder
@@ -202,7 +202,7 @@ public class AuthenticationWebFilterTests {
 			.expectStatus().isOk()
 			.expectBody().isEmpty();
 
-		verify(this.successHandler).success(eq(authentication.block()), any(), any());
+		verify(this.successHandler).success(eq(authentication.block()), any());
 		verify(this.securityContextRepository).save(any(), any());
 		verifyZeroInteractions(this.entryPoint);
 	}

webflux/src/test/java/org/springframework/security/web/server/authentication/RedirectAuthenticationSuccessHandlerTests.java

@@ -28,6 +28,7 @@ import org.springframework.security.authentication.AuthenticationCredentialsNotF
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.AuthenticationException;
 import org.springframework.security.web.server.RedirectStrategy;
+import org.springframework.security.web.server.WebFilterExchange;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilterChain;
 import reactor.core.publisher.Mono;
@@ -70,8 +71,8 @@ public class RedirectAuthenticationSuccessHandlerTests {
 
 	@Test
 	public void successWhenNoSubscribersThenNoActions() {
-		this.handler.success(this.authentication, this.exchange,
-			this.chain);
+		this.handler.success(this.authentication, new WebFilterExchange(this.exchange,
+			this.chain));
 
 		verifyZeroInteractions(this.exchange);
 	}
@@ -80,8 +81,8 @@ public class RedirectAuthenticationSuccessHandlerTests {
 	public void successWhenSubscribeThenStatusAndLocationSet() {
 		this.exchange = MockServerHttpRequest.get("/").toExchange();
 
-		this.handler.success(this.authentication, this.exchange,
-			this.chain).block();
+		this.handler.success(this.authentication, new WebFilterExchange(this.exchange,
+			this.chain)).block();
 
 		assertThat(this.exchange.getResponse().getStatusCode()).isEqualTo(
 			HttpStatus.FOUND);
@@ -95,8 +96,8 @@ public class RedirectAuthenticationSuccessHandlerTests {
 		this.handler.setRedirectStrategy(this.redirectStrategy);
 		this.exchange = MockServerHttpRequest.get("/").toExchange();
 
-		assertThat(this.handler.success(this.authentication, this.exchange,
-			this.chain)).isEqualTo(result);
+		assertThat(this.handler.success(this.authentication, new WebFilterExchange(this.exchange,
+			this.chain))).isEqualTo(result);
 		verify(this.redirectStrategy).sendRedirect(any(), eq(this.location));
 	}