From 2374e7d4eb3c4bbfcf89becda7c90c722e77c4b2 Mon Sep 17 00:00:00 2001
From: Scott Frederick <scott.frederick@broadcom.com>
Date: Tue, 6 Feb 2024 15:25:39 -0600
Subject: [PATCH] Fix invalid request handling for WebFlux actuator endpoints

Fixes gh-39236
---
 ...AbstractWebFluxEndpointHandlerMapping.java | 25 +++++++++-
 .../AbstractWebEndpointIntegrationTests.java  | 46 +++++++++++++++++++
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/reactive/AbstractWebFluxEndpointHandlerMapping.java b/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/reactive/AbstractWebFluxEndpointHandlerMapping.java
index 3621a5b94a5..3cb0cf56e0a 100644
--- a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/reactive/AbstractWebFluxEndpointHandlerMapping.java
+++ b/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/reactive/AbstractWebFluxEndpointHandlerMapping.java
@@ -83,6 +83,7 @@ import org.springframework.web.util.pattern.PathPattern;
  * @author Madhura Bhave
  * @author Phillip Webb
  * @author Brian Clozel
+ * @author Scott Frederick
  * @since 2.0.0
  */
 @ImportRuntimeHints(AbstractWebFluxEndpointHandlerMappingRuntimeHints.class)
@@ -260,6 +261,26 @@ public abstract class AbstractWebFluxEndpointHandlerMapping extends RequestMappi
 
 	}
 
+	protected static final class ExceptionCapturingInvoker implements OperationInvoker {
+
+		private final OperationInvoker invoker;
+
+		public ExceptionCapturingInvoker(OperationInvoker invoker) {
+			this.invoker = invoker;
+		}
+
+		@Override
+		public Object invoke(InvocationContext context) {
+			try {
+				return this.invoker.invoke(context);
+			}
+			catch (Exception ex) {
+				return Mono.error(ex);
+			}
+		}
+
+	}
+
 	/**
 	 * Reactive handler providing actuator links at the root endpoint.
 	 */
@@ -303,9 +324,9 @@ public abstract class AbstractWebFluxEndpointHandlerMapping extends RequestMappi
 		private OperationInvoker getInvoker(WebOperation operation) {
 			OperationInvoker invoker = operation::invoke;
 			if (operation.isBlocking()) {
-				invoker = new ElasticSchedulerInvoker(invoker);
+				return new ElasticSchedulerInvoker(invoker);
 			}
-			return invoker;
+			return new ExceptionCapturingInvoker(invoker);
 		}
 
 		private Supplier<Mono<? extends SecurityContext>> getSecurityContextSupplier() {
diff --git a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/web/annotation/AbstractWebEndpointIntegrationTests.java b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/web/annotation/AbstractWebEndpointIntegrationTests.java
index 23f370bac33..cc2286270fd 100644
--- a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/web/annotation/AbstractWebEndpointIntegrationTests.java
+++ b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/web/annotation/AbstractWebEndpointIntegrationTests.java
@@ -226,6 +226,31 @@ public abstract class AbstractWebEndpointIntegrationTests<T extends Configurable
 					.isEqualTo("1 2"));
 	}
 
+	@Test
+	void readOperationWithQueryParametersMissing() {
+		load(QueryEndpointConfiguration.class,
+				(client) -> client.get().uri("/query").exchange().expectStatus().isBadRequest());
+	}
+
+	@Test
+	void reactiveReadOperationWithSingleQueryParameters() {
+		load(ReactiveQueryEndpointConfiguration.class,
+				(client) -> client.get()
+					.uri("/query?param=test")
+					.exchange()
+					.expectStatus()
+					.isOk()
+					.expectBody()
+					.jsonPath("query")
+					.isEqualTo("test"));
+	}
+
+	@Test
+	void reactiveReadOperationWithQueryParametersMissing() {
+		load(ReactiveQueryEndpointConfiguration.class,
+				(client) -> client.get().uri("/query").exchange().expectStatus().isBadRequest());
+	}
+
 	@Test
 	void readOperationWithSingleQueryParametersAndMultipleValues() {
 		load(QueryEndpointConfiguration.class,
@@ -732,6 +757,17 @@ public abstract class AbstractWebEndpointIntegrationTests<T extends Configurable
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	@Import(BaseConfiguration.class)
+	static class ReactiveQueryEndpointConfiguration {
+
+		@Bean
+		ReactiveQueryEndpoint reactiveQueryEndpoint() {
+			return new ReactiveQueryEndpoint();
+		}
+
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	@Import(BaseConfiguration.class)
 	static class VoidWriteResponseEndpointConfiguration {
@@ -974,6 +1010,16 @@ public abstract class AbstractWebEndpointIntegrationTests<T extends Configurable
 
 	}
 
+	@Endpoint(id = "query")
+	static class ReactiveQueryEndpoint {
+
+		@ReadOperation
+		Mono<Map<String, String>> query(String param) {
+			return Mono.just(Collections.singletonMap("query", param));
+		}
+
+	}
+
 	@Endpoint(id = "voidwrite")
 	static class VoidWriteResponseEndpoint {