Improve WebFlux support for response status exceptions

Support @ResponseStatus annotated exceptions.
Supported root cause exceptions with response status information.

Issue: SPR-16567

spring-web/src/main/java/org/springframework/web/server/handler/ResponseStatusExceptionHandler.java

@@ -20,14 +20,22 @@ import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import reactor.core.publisher.Mono;
 
+import org.springframework.core.annotation.AnnotatedElementUtils;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.lang.Nullable;
+import org.springframework.web.bind.annotation.ResponseStatus;
 import org.springframework.web.server.ResponseStatusException;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebExceptionHandler;
 
 /**
- * Handle {@link ResponseStatusException} by setting the response status.
+ * Handle instances of {@link ResponseStatusException}, or of exceptions annotated
+ * with {@link ResponseStatus @ResponseStatus}, by extracting the
+ * {@code HttpStatus} from them and updating the status of the response.
+ *
+ * <p>If the response is already committed, the error remains unresolved and is
+ * propagated.
  *
  * @author Rossen Stoyanchev
  * @author Sebastien Deleuze
@@ -40,24 +48,37 @@ public class ResponseStatusExceptionHandler implements WebExceptionHandler {
 
 	@Override
 	public Mono<Void> handle(ServerWebExchange exchange, Throwable ex) {
-		if (ex instanceof ResponseStatusException) {
+		HttpStatus status = resolveHttpStatus(ex);
-			HttpStatus status = ((ResponseStatusException) ex).getStatus();
+		if (status != null && exchange.getResponse().setStatusCode(status)) {
-			if (exchange.getResponse().setStatusCode(status)) {
+			if (status.is5xxServerError()) {
-				if (status.is5xxServerError()) {
+				logger.error(buildMessage(exchange.getRequest(), ex));
-					logger.error(buildMessage(exchange.getRequest(), ex));
-				}
-				else if (status == HttpStatus.BAD_REQUEST) {
-					logger.warn(buildMessage(exchange.getRequest(), ex));
-				}
-				else {
-					logger.trace(buildMessage(exchange.getRequest(), ex));
-				}
-				return exchange.getResponse().setComplete();
 			}
+			else if (status == HttpStatus.BAD_REQUEST) {
+				logger.warn(buildMessage(exchange.getRequest(), ex));
+			}
+			else {
+				logger.trace(buildMessage(exchange.getRequest(), ex));
+			}
+			return exchange.getResponse().setComplete();
 		}
 		return Mono.error(ex);
 	}
 
+	@Nullable
+	private HttpStatus resolveHttpStatus(Throwable ex) {
+		if (ex instanceof ResponseStatusException) {
+			return ((ResponseStatusException) ex).getStatus();
+		}
+		ResponseStatus status = AnnotatedElementUtils.findMergedAnnotation(ex.getClass(), ResponseStatus.class);
+		if (status != null) {
+			return status.code();
+		}
+		if (ex.getCause() != null) {
+			return resolveHttpStatus(ex.getCause());
+		}
+		return null;
+	}
+
 	private String buildMessage(ServerHttpRequest request, Throwable ex) {
 		return "Failed to handle request [" + request.getMethod() + " "
 				+ request.getURI() + "]: " + ex.getMessage();

spring-web/src/test/java/org/springframework/web/server/handler/ResponseStatusExceptionHandlerTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -25,6 +25,7 @@ import reactor.test.StepVerifier;
 import org.springframework.http.HttpStatus;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpRequest;
 import org.springframework.mock.web.test.server.MockServerWebExchange;
+import org.springframework.web.bind.annotation.ResponseStatus;
 import org.springframework.web.server.ResponseStatusException;
 
 import static org.junit.Assert.assertEquals;
@@ -42,21 +43,42 @@ public class ResponseStatusExceptionHandlerTests {
 
 
 	@Test
-	public void handleException() throws Exception {
+	public void handleResponseStatusException() {
 		Throwable ex = new ResponseStatusException(HttpStatus.BAD_REQUEST, "");
 		this.handler.handle(this.exchange, ex).block(Duration.ofSeconds(5));
 		assertEquals(HttpStatus.BAD_REQUEST, this.exchange.getResponse().getStatusCode());
 	}
 
 	@Test
-	public void unresolvedException() throws Exception {
+	public void handleAnnotatedException() {
+		Throwable ex = new CustomException();
+		this.handler.handle(this.exchange, ex).block(Duration.ofSeconds(5));
+		assertEquals(HttpStatus.I_AM_A_TEAPOT, this.exchange.getResponse().getStatusCode());
+	}
+
+	@Test
+	public void handleNestedResponseStatusException() {
+		Throwable ex = new Exception(new ResponseStatusException(HttpStatus.BAD_REQUEST, ""));
+		this.handler.handle(this.exchange, ex).block(Duration.ofSeconds(5));
+		assertEquals(HttpStatus.BAD_REQUEST, this.exchange.getResponse().getStatusCode());
+	}
+
+	@Test
+	public void handleNestedAnnotatedException() {
+		Throwable ex = new Exception(new CustomException());
+		this.handler.handle(this.exchange, ex).block(Duration.ofSeconds(5));
+		assertEquals(HttpStatus.I_AM_A_TEAPOT, this.exchange.getResponse().getStatusCode());
+	}
+
+	@Test
+	public void unresolvedException() {
 		Throwable expected = new IllegalStateException();
 		Mono<Void> mono = this.handler.handle(this.exchange, expected);
 		StepVerifier.create(mono).consumeErrorWith(actual -> assertSame(expected, actual)).verify();
 	}
 
 	@Test // SPR-16231
-	public void responseCommitted() throws Exception {
+	public void responseCommitted() {
 		Throwable ex = new ResponseStatusException(HttpStatus.INTERNAL_SERVER_ERROR, "Oops");
 		this.exchange.getResponse().setStatusCode(HttpStatus.CREATED);
 		Mono<Void> mono = this.exchange.getResponse().setComplete()
@@ -64,4 +86,9 @@ public class ResponseStatusExceptionHandlerTests {
 		StepVerifier.create(mono).consumeErrorWith(actual -> assertSame(ex, actual)).verify();
 	}
 
+
+	@ResponseStatus(HttpStatus.I_AM_A_TEAPOT)
+	private static class CustomException extends Exception {
+	}
+
 }