From 95abd18fea067b271bf50252b0a741c5850cab93 Mon Sep 17 00:00:00 2001
From: Brian Clozel <bclozel@pivotal.io>
Date: Thu, 3 Nov 2016 17:00:18 +0100
Subject: [PATCH] Fix Exception Handling result handler resolution

Prior to this commit, the wrong `HandlerResultHandler` could be
resolved when handling exceptions; this could happen only if the
original handler and exception handler had different signatures:

```
Publisher<String> originalHandler() { ... }

@ExceptionHandler(MyCustomException.class)
ResponseEntity<Mono<String>> handleException() { ... }
```

In that case, the `ResponseBodyResultHandler` would be used when
handling exceptions instead of the `ResponseEntityResultHandler`.

This commit ensures that the `HandlerResult` returned by the exception
handler is used to resolve the `HandlerResultHandler`. The latter will
process the result and use it to write to the HTTP response.

Issue: SPR-14876
---
 .../web/reactive/DispatcherHandler.java          |  2 +-
 ...MappingExceptionHandlingIntegrationTests.java | 16 +++++++++++-----
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/spring-web-reactive/src/main/java/org/springframework/web/reactive/DispatcherHandler.java b/spring-web-reactive/src/main/java/org/springframework/web/reactive/DispatcherHandler.java
index ef2e55079d6..c6a3b339295 100644
--- a/spring-web-reactive/src/main/java/org/springframework/web/reactive/DispatcherHandler.java
+++ b/spring-web-reactive/src/main/java/org/springframework/web/reactive/DispatcherHandler.java
@@ -142,7 +142,7 @@ public class DispatcherHandler implements WebHandler, ApplicationContextAware {
 	private Mono<Void> handleResult(ServerWebExchange exchange, HandlerResult result) {
 		return getResultHandler(result).handleResult(exchange, result)
 				.otherwise(ex -> result.applyExceptionHandler(ex).then(exceptionResult ->
-						getResultHandler(result).handleResult(exchange, exceptionResult)));
+						getResultHandler(exceptionResult).handleResult(exchange, exceptionResult)));
 	}
 
 	private HandlerResultHandler getResultHandler(HandlerResult handlerResult) {
diff --git a/spring-web-reactive/src/test/java/org/springframework/web/reactive/result/method/annotation/RequestMappingExceptionHandlingIntegrationTests.java b/spring-web-reactive/src/test/java/org/springframework/web/reactive/result/method/annotation/RequestMappingExceptionHandlingIntegrationTests.java
index 9c1bedc2049..2a7f12045ef 100644
--- a/spring-web-reactive/src/test/java/org/springframework/web/reactive/result/method/annotation/RequestMappingExceptionHandlingIntegrationTests.java
+++ b/spring-web-reactive/src/test/java/org/springframework/web/reactive/result/method/annotation/RequestMappingExceptionHandlingIntegrationTests.java
@@ -25,6 +25,7 @@ import org.springframework.context.annotation.AnnotationConfigApplicationContext
 import org.springframework.context.annotation.ComponentScan;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.http.HttpHeaders;
+import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.ExceptionHandler;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.RestController;
@@ -52,13 +53,13 @@ public class RequestMappingExceptionHandlingIntegrationTests extends AbstractReq
 
 	@Test
 	public void controllerThrowingException() throws Exception {
-		String expected = "Recovered from error: Boo";
+		String expected = "Recovered from error: State";
 		assertEquals(expected, performGet("/thrown-exception", new HttpHeaders(), String.class).getBody());
 	}
 
 	@Test
 	public void controllerReturnsMonoError() throws Exception {
-		String expected = "Recovered from error: Boo";
+		String expected = "Recovered from error: Argument";
 		assertEquals(expected, performGet("/mono-error", new HttpHeaders(), String.class).getBody());
 	}
 
@@ -78,19 +79,24 @@ public class RequestMappingExceptionHandlingIntegrationTests extends AbstractReq
 
 		@GetMapping("/thrown-exception")
 		public Publisher<String> handleAndThrowException() {
-			throw new IllegalStateException("Boo");
+			throw new IllegalStateException("State");
 		}
 
 		@GetMapping("/mono-error")
 		public Publisher<String> handleWithError() {
-			return Mono.error(new IllegalStateException("Boo"));
+			return Mono.error(new IllegalArgumentException("Argument"));
 		}
 
 		@ExceptionHandler
-		public Publisher<String> handleException(IllegalStateException ex) {
+		public Publisher<String> handleArgumentException(IllegalArgumentException ex) {
 			return Mono.just("Recovered from error: " + ex.getMessage());
 		}
 
+		@ExceptionHandler
+		public ResponseEntity<Publisher<String>> handleStateException(IllegalStateException ex) {
+			return ResponseEntity.ok(Mono.just("Recovered from error: " + ex.getMessage()));
+		}
+
 	}
 
 }