From 81e74e65d4c82aef6d47c93dc3f45fb48708878a Mon Sep 17 00:00:00 2001
From: Rob Winch <362503+rwinch@users.noreply.github.com>
Date: Tue, 22 Oct 2024 18:17:15 -0500
Subject: [PATCH] Support ServerExchangeRejectedHandler @Bean

Closes gh-16061
---
 .../WebFluxSecurityConfiguration.java         |  5 ++-
 .../WebFluxSecurityConfigurationTests.java    | 28 ++++++++++++++++
 .../pages/reactive/exploits/firewall.adoc     | 32 +++++++++++++++++++
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/config/src/main/java/org/springframework/security/config/annotation/web/reactive/WebFluxSecurityConfiguration.java b/config/src/main/java/org/springframework/security/config/annotation/web/reactive/WebFluxSecurityConfiguration.java
index 0aba5e9719..72ea9a1a3a 100644
--- a/config/src/main/java/org/springframework/security/config/annotation/web/reactive/WebFluxSecurityConfiguration.java
+++ b/config/src/main/java/org/springframework/security/config/annotation/web/reactive/WebFluxSecurityConfiguration.java
@@ -34,6 +34,7 @@
 import org.springframework.security.web.server.ObservationWebFilterChainDecorator;
 import org.springframework.security.web.server.SecurityWebFilterChain;
 import org.springframework.security.web.server.WebFilterChainProxy;
+import org.springframework.security.web.server.firewall.ServerExchangeRejectedHandler;
 import org.springframework.security.web.server.firewall.ServerWebExchangeFirewall;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.ObjectUtils;
@@ -81,12 +82,14 @@ void setObservationRegistry(ObservationRegistry observationRegistry) {
 
 	@Bean(SPRING_SECURITY_WEBFILTERCHAINFILTER_BEAN_NAME)
 	@Order(WEB_FILTER_CHAIN_FILTER_ORDER)
-	WebFilterChainProxy springSecurityWebFilterChainFilter(ObjectProvider<ServerWebExchangeFirewall> firewall) {
+	WebFilterChainProxy springSecurityWebFilterChainFilter(ObjectProvider<ServerWebExchangeFirewall> firewall,
+			ObjectProvider<ServerExchangeRejectedHandler> rejectedHandler) {
 		WebFilterChainProxy proxy = new WebFilterChainProxy(getSecurityWebFilterChains());
 		if (!this.observationRegistry.isNoop()) {
 			proxy.setFilterChainDecorator(new ObservationWebFilterChainDecorator(this.observationRegistry));
 		}
 		firewall.ifUnique(proxy::setFirewall);
+		rejectedHandler.ifUnique(proxy::setExchangeRejectedHandler);
 		return proxy;
 	}
 
diff --git a/config/src/test/java/org/springframework/security/config/annotation/web/reactive/WebFluxSecurityConfigurationTests.java b/config/src/test/java/org/springframework/security/config/annotation/web/reactive/WebFluxSecurityConfigurationTests.java
index 4c8c249d7e..d8be15886b 100644
--- a/config/src/test/java/org/springframework/security/config/annotation/web/reactive/WebFluxSecurityConfigurationTests.java
+++ b/config/src/test/java/org/springframework/security/config/annotation/web/reactive/WebFluxSecurityConfigurationTests.java
@@ -32,6 +32,8 @@
 import org.springframework.security.config.test.SpringTestContextExtension;
 import org.springframework.security.config.users.ReactiveAuthenticationTestConfiguration;
 import org.springframework.security.web.server.WebFilterChainProxy;
+import org.springframework.security.web.server.firewall.HttpStatusExchangeRejectedHandler;
+import org.springframework.security.web.server.firewall.ServerExchangeRejectedHandler;
 import org.springframework.security.web.server.firewall.ServerWebExchangeFirewall;
 import org.springframework.web.server.handler.DefaultWebFilterChain;
 
@@ -70,6 +72,20 @@ void loadConfigWhenDefaultThenFirewalled() throws Exception {
 		assertThat(exchange.getResponse().getStatusCode()).isEqualTo(HttpStatus.BAD_REQUEST);
 	}
 
+	@Test
+	void loadConfigWhenCustomRejectedHandler() throws Exception {
+		this.spring
+			.register(ServerHttpSecurityConfiguration.class, ReactiveAuthenticationTestConfiguration.class,
+					WebFluxSecurityConfiguration.class, CustomServerExchangeRejectedHandlerConfig.class)
+			.autowire();
+		WebFilterChainProxy webFilterChainProxy = this.spring.getContext().getBean(WebFilterChainProxy.class);
+		MockServerWebExchange exchange = MockServerWebExchange.from(MockServerHttpRequest.get("/;/").build());
+		DefaultWebFilterChain chain = emptyChain();
+		webFilterChainProxy.filter(exchange, chain).block();
+		assertThat(exchange.getResponse().getStatusCode())
+			.isEqualTo(CustomServerExchangeRejectedHandlerConfig.EXPECTED_STATUS);
+	}
+
 	@Test
 	void loadConfigWhenFirewallBeanThenCustomized() throws Exception {
 		this.spring
@@ -107,6 +123,18 @@ ServerWebExchangeFirewall noOpFirewall() {
 
 	}
 
+	@Configuration
+	static class CustomServerExchangeRejectedHandlerConfig {
+
+		static HttpStatus EXPECTED_STATUS = HttpStatus.I_AM_A_TEAPOT;
+
+		@Bean
+		ServerExchangeRejectedHandler rejectedHandler() {
+			return new HttpStatusExchangeRejectedHandler(EXPECTED_STATUS);
+		}
+
+	}
+
 	@Configuration
 	static class SubclassConfig extends WebFluxSecurityConfiguration {
 
diff --git a/docs/modules/ROOT/pages/reactive/exploits/firewall.adoc b/docs/modules/ROOT/pages/reactive/exploits/firewall.adoc
index dce7fb05f7..f359ebdcf9 100644
--- a/docs/modules/ROOT/pages/reactive/exploits/firewall.adoc
+++ b/docs/modules/ROOT/pages/reactive/exploits/firewall.adoc
@@ -200,3 +200,35 @@ firewall.setAllowedHeaderValues {
 }
 ----
 ======
+
+The `ServerExchangeRejectedHandler` interface is used to handle `ServerExchangeRejectedException` throw by Spring Security's `ServerWebExchangeFirewall`.
+By default `HttpStatusExchangeRejectedHandler` is used to send an HTTP 400 response to clients when a request is rejected.
+To customize the behavior, users can expose a `ServerExchangeRejectedHandler` Bean.
+For example, the following will send an HTTP 404 when the request is rejected:
+
+
+.Send 404 on Request Rejected
+[tabs]
+======
+Java::
++
+[source,java,role="primary"]
+----
+@Bean
+ServerExchangeRejectedHandler rejectedHandler() {
+	return new HttpStatusExchangeRejectedHandler(HttpStatus.NOT_FOUND);
+}
+----
+
+Kotlin::
++
+[source,kotlin,role="secondary"]
+----
+@Bean
+fun rejectedHandler(): ServerExchangeRejectedHandler {
+    return HttpStatusExchangeRejectedHandler(HttpStatus.NOT_FOUND)
+}
+----
+======
+
+Handling can be completely customized by creating a custom `ServerExchangeRejectedHandler` implementation.