Support ResponseStatusException in WebFlux fn

This commit introduces support for the ResponseStatusException in the
functional web framework.

Issue: SPR-15344

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctions.java

@@ -29,6 +29,7 @@ import org.springframework.util.Assert;
 import org.springframework.web.reactive.HandlerMapping;
 import org.springframework.web.reactive.function.server.support.HandlerFunctionAdapter;
 import org.springframework.web.reactive.function.server.support.ServerResponseResultHandler;
+import org.springframework.web.server.ResponseStatusException;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebHandler;
 import org.springframework.web.server.adapter.HttpWebHandlerAdapter;
@@ -155,7 +156,7 @@ public abstract class RouterFunctions {
 	 * For instance
 	 * <pre class="code">
 	 * Resource location = new FileSystemResource("public-resources/");
-	 * RoutingFunction&lt;Resource&gt; resources = RouterFunctions.resources("/resources/**", location);
+	 * RoutingFunction&lt;ServerResponse&gt; resources = RouterFunctions.resources("/resources/**", location);
      * </pre>
 	 * @param pattern the pattern to match
 	 * @param location the location directory relative to which resources should be resolved
@@ -232,11 +233,23 @@ public abstract class RouterFunctions {
 			addAttributes(exchange, request);
 			return routerFunction.route(request)
 					.defaultIfEmpty(notFound())
-					.then(handlerFunction -> handlerFunction.handle(request))
+					.then(handlerFunction -> invokeHandler(handlerFunction, request))
+					.otherwise(ResponseStatusException.class, RouterFunctions::responseStatusFallback)
 					.then(response -> response.writeTo(exchange, strategies));
 		});
 	}
 
+	private static <T extends ServerResponse> Mono<T> invokeHandler(HandlerFunction<T> handlerFunction,
+			ServerRequest request) {
+		try {
+			return handlerFunction.handle(request);
+		}
+		catch (Throwable t) {
+			return Mono.error(t);
+		}
+	}
+
+
 	/**
 	 * Convert the given {@code RouterFunction} into a {@code HandlerMapping}.
 	 * This conversion uses {@linkplain HandlerStrategies#builder() default strategies}.
@@ -284,6 +297,11 @@ public abstract class RouterFunctions {
 		return (HandlerFunction<T>) NOT_FOUND_HANDLER;
 	}
 
+	@SuppressWarnings("unchecked")
+	private static <T extends ServerResponse> Mono<T> responseStatusFallback(ResponseStatusException ex) {
+		return (Mono<T>) ServerResponse.status(ex.getStatus()).build();
+	}
+
 	@SuppressWarnings("unchecked")
 	static <T extends ServerResponse> HandlerFunction<T> cast(HandlerFunction<?> handlerFunction) {
 		return (HandlerFunction<T>) handlerFunction;

spring-webflux/src/test/java/org/springframework/web/reactive/function/server/PublisherHandlerFunctionIntegrationTests.java

@@ -16,9 +16,11 @@
 
 package org.springframework.web.reactive.function.server;
 
+import java.io.IOException;
 import java.net.URI;
 import java.util.List;
 
+import org.junit.Before;
 import org.junit.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -28,20 +30,41 @@ import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.RequestEntity;
 import org.springframework.http.ResponseEntity;
+import org.springframework.http.client.ClientHttpResponse;
+import org.springframework.web.client.ResponseErrorHandler;
 import org.springframework.web.client.RestTemplate;
+import org.springframework.web.server.ResponseStatusException;
 
-import static org.junit.Assert.*;
-import static org.springframework.web.reactive.function.BodyExtractors.*;
-import static org.springframework.web.reactive.function.BodyInserters.*;
-import static org.springframework.web.reactive.function.server.RequestPredicates.*;
-import static org.springframework.web.reactive.function.server.RouterFunctions.*;
+import static org.junit.Assert.assertEquals;
+import static org.springframework.web.reactive.function.BodyExtractors.toMono;
+import static org.springframework.web.reactive.function.BodyInserters.fromPublisher;
+import static org.springframework.web.reactive.function.server.RequestPredicates.GET;
+import static org.springframework.web.reactive.function.server.RequestPredicates.POST;
+import static org.springframework.web.reactive.function.server.RouterFunctions.route;
 
 /**
  * @author Arjen Poutsma
  */
 public class PublisherHandlerFunctionIntegrationTests extends AbstractRouterFunctionIntegrationTests {
 
-	private final RestTemplate restTemplate = new RestTemplate();
+	private RestTemplate restTemplate;
+
+	@Before
+	public void createRestTemplate() {
+		restTemplate = new RestTemplate();
+		restTemplate.setErrorHandler(new ResponseErrorHandler() {
+			@Override
+			public boolean hasError(ClientHttpResponse response) throws IOException {
+				return false;
+			}
+
+			@Override
+			public void handleError(ClientHttpResponse response) throws IOException {
+
+			}
+		});
+
+	}
 
 
 	@Override
@@ -49,7 +72,9 @@ public class PublisherHandlerFunctionIntegrationTests extends AbstractRouterFunc
 		PersonHandler personHandler = new PersonHandler();
 		return route(GET("/mono"), personHandler::mono)
 				.and(route(POST("/mono"), personHandler::postMono))
-				.and(route(GET("/flux"), personHandler::flux));
+				.and(route(GET("/flux"), personHandler::flux))
+				.and(route(GET("/throwRSE"), personHandler::throwResponseStatusException))
+				.and(route(GET("/returnRSE"), personHandler::returnResponseStatusException));
 	}
 
 
@@ -86,6 +111,19 @@ public class PublisherHandlerFunctionIntegrationTests extends AbstractRouterFunc
 		assertEquals("Jack", result.getBody().getName());
 	}
 
+	@Test
+	public void responseStatusException() {
+		ResponseEntity<String> result =
+				restTemplate.getForEntity("http://localhost:" + port + "/throwRSE", String.class);
+
+		assertEquals(HttpStatus.BAD_REQUEST, result.getStatusCode());
+
+		result = restTemplate.getForEntity("http://localhost:" + port + "/returnRSE", String.class);
+
+		assertEquals(HttpStatus.BAD_REQUEST, result.getStatusCode());
+	}
+
+
 
 	private static class PersonHandler {
 
@@ -105,6 +143,14 @@ public class PublisherHandlerFunctionIntegrationTests extends AbstractRouterFunc
 			return ServerResponse.ok().body(
 					fromPublisher(Flux.just(person1, person2), Person.class));
 		}
+
+		public Mono<ServerResponse> throwResponseStatusException(ServerRequest request) {
+			throw new ResponseStatusException(HttpStatus.BAD_REQUEST, "Bad Request");
+		}
+
+		public Mono<ServerResponse> returnResponseStatusException(ServerRequest request) {
+			return Mono.error(new ResponseStatusException(HttpStatus.BAD_REQUEST, "Bad Request"));
+		}
 	}