From 85704c890e1ff05e33a43cd5a4c318affc513341 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@vmware.com>
Date: Thu, 27 Jul 2023 22:12:54 +0300
Subject: [PATCH] Use DispatchExceptionHandler for handler errors

See gh-30930
---
 .../reactive/DispatchExceptionHandler.java    |  3 +-
 .../web/reactive/DispatcherHandler.java       | 59 +++++++++----------
 2 files changed, 30 insertions(+), 32 deletions(-)

diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/DispatchExceptionHandler.java b/spring-webflux/src/main/java/org/springframework/web/reactive/DispatchExceptionHandler.java
index 4486a38bfd3..06a2ec8a12a 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/DispatchExceptionHandler.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/DispatchExceptionHandler.java
@@ -29,7 +29,8 @@ import org.springframework.web.server.ServerWebExchange;
  * a {@link HandlerAdapter} to apply its exception handling to deferred exceptions
  * from asynchronous return values, and to response rendering.
  * <li>Implemented by a {@link HandlerAdapter} in order to handle exceptions that
- * occur before a request is mapped to a handler.
+ * occur before a request is mapped to a handler, or for unhandled errors from a
+ * handler..
  * </ul>
  *
  * @author Rossen Stoyanchev
diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/DispatcherHandler.java b/spring-webflux/src/main/java/org/springframework/web/reactive/DispatcherHandler.java
index 766a1843544..c51a5f5026e 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/DispatcherHandler.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/DispatcherHandler.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -150,7 +150,7 @@ public class DispatcherHandler implements WebHandler, PreFlightRequestHandler, A
 				.concatMap(mapping -> mapping.getHandler(exchange))
 				.next()
 				.switchIfEmpty(createNotFoundError())
-				.onErrorResume(ex -> handleDispatchError(exchange, ex))
+				.onErrorResume(ex -> handleResultMono(exchange, Mono.error(ex)))
 				.flatMap(handler -> handleRequestWith(exchange, handler));
 	}
 
@@ -161,8 +161,7 @@ public class DispatcherHandler implements WebHandler, PreFlightRequestHandler, A
 		});
 	}
 
-	private Mono<Void> handleDispatchError(ServerWebExchange exchange, Throwable ex) {
-		Mono<HandlerResult> resultMono = Mono.error(ex);
+	private Mono<Void> handleResultMono(ServerWebExchange exchange, Mono<HandlerResult> resultMono) {
 		if (this.handlerAdapters != null) {
 			for (HandlerAdapter adapter : this.handlerAdapters) {
 				if (adapter instanceof DispatchExceptionHandler exceptionHandler) {
@@ -170,36 +169,19 @@ public class DispatcherHandler implements WebHandler, PreFlightRequestHandler, A
 				}
 			}
 		}
-		return resultMono.flatMap(result -> handleResult(exchange, result));
-	}
-
-	private Mono<Void> handleRequestWith(ServerWebExchange exchange, Object handler) {
-		if (ObjectUtils.nullSafeEquals(exchange.getResponse().getStatusCode(), HttpStatus.FORBIDDEN)) {
-			return Mono.empty();  // CORS rejection
-		}
-		if (this.handlerAdapters != null) {
-			for (HandlerAdapter adapter : this.handlerAdapters) {
-				if (adapter.supports(handler)) {
-					return adapter.handle(exchange, handler)
-							.flatMap(result -> handleResult(exchange, result));
-				}
+		return resultMono.flatMap(result -> {
+			Mono<Void> voidMono = handleResult(exchange, result, "Handler " + result.getHandler());
+			if (result.getExceptionHandler() != null) {
+				voidMono = voidMono.onErrorResume(ex ->
+						result.getExceptionHandler().handleError(exchange, ex).flatMap(result2 ->
+								handleResult(exchange, result2, "Exception handler " +
+										result2.getHandler() + ", error=\"" + ex.getMessage() + "\"")));
 			}
-		}
-		return Mono.error(new IllegalStateException("No HandlerAdapter: " + handler));
+			return voidMono;
+		});
 	}
 
-	private Mono<Void> handleResult(ServerWebExchange exchange, HandlerResult result) {
-		Mono<Void> resultMono = doHandleResult(exchange, result, "Handler " + result.getHandler());
-		if (result.getExceptionHandler() != null) {
-			resultMono = resultMono.onErrorResume(ex ->
-					result.getExceptionHandler().handleError(exchange, ex).flatMap(result2 ->
-							doHandleResult(exchange, result2, "Exception handler " +
-									result2.getHandler() + ", error=\"" + ex.getMessage() + "\"")));
-		}
-		return resultMono;
-	}
-
-	private Mono<Void> doHandleResult(
+	private Mono<Void> handleResult(
 			ServerWebExchange exchange, HandlerResult handlerResult, String description) {
 
 		if (this.resultHandlers != null) {
@@ -214,6 +196,21 @@ public class DispatcherHandler implements WebHandler, PreFlightRequestHandler, A
 				"No HandlerResultHandler for " + handlerResult.getReturnValue()));
 	}
 
+	private Mono<Void> handleRequestWith(ServerWebExchange exchange, Object handler) {
+		if (ObjectUtils.nullSafeEquals(exchange.getResponse().getStatusCode(), HttpStatus.FORBIDDEN)) {
+			return Mono.empty();  // CORS rejection
+		}
+		if (this.handlerAdapters != null) {
+			for (HandlerAdapter adapter : this.handlerAdapters) {
+				if (adapter.supports(handler)) {
+					Mono<HandlerResult> resultMono = adapter.handle(exchange, handler);
+					return handleResultMono(exchange, resultMono);
+				}
+			}
+		}
+		return Mono.error(new IllegalStateException("No HandlerAdapter: " + handler));
+	}
+
 	@Override
 	public Mono<Void> handlePreFlight(ServerWebExchange exchange) {
 		return Flux.fromIterable(this.handlerMappings != null ? this.handlerMappings : Collections.emptyList())