From 69df27a99ff191cbdfc9ecf7c19f3ac3e87888af Mon Sep 17 00:00:00 2001
From: Ingmar van Dijk <ingmar.van.dijk@philips.com>
Date: Mon, 6 Jul 2020 09:13:07 +0200
Subject: [PATCH] Throw 404 ResponseStatusException when no routes found

This commit makes it possible to customize 404 responses generated by
RouterFunctionWebHandler, by throwing an ResponseStatusException
instead of returning a standard 404 response.

See gh-25358
---
 .../function/server/RouterFunctions.java      | 12 +++-------
 .../function/server/RouterFunctionsTests.java | 24 +++++++++++++++++++
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctions.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctions.java
index 30157153d4a..8b241f70cd9 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctions.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctions.java
@@ -32,10 +32,12 @@ import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import org.springframework.core.io.Resource;
+import org.springframework.http.HttpStatus;
 import org.springframework.http.codec.HttpMessageWriter;
 import org.springframework.http.server.reactive.HttpHandler;
 import org.springframework.util.Assert;
 import org.springframework.web.reactive.result.view.ViewResolver;
+import org.springframework.web.server.ResponseStatusException;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebHandler;
 import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
@@ -1231,9 +1233,6 @@ public abstract class RouterFunctions {
 
 	private static class RouterFunctionWebHandler implements WebHandler {
 
-		private static final HandlerFunction<ServerResponse> NOT_FOUND_HANDLER =
-				request -> ServerResponse.notFound().build();
-
 		private final HandlerStrategies strategies;
 
 		private final RouterFunction<?> routerFunction;
@@ -1249,7 +1248,7 @@ public abstract class RouterFunctions {
 				ServerRequest request = new DefaultServerRequest(exchange, this.strategies.messageReaders());
 				addAttributes(exchange, request);
 				return this.routerFunction.route(request)
-						.defaultIfEmpty(notFound())
+						.switchIfEmpty(Mono.error(new ResponseStatusException(HttpStatus.NOT_FOUND)))
 						.flatMap(handlerFunction -> wrapException(() -> handlerFunction.handle(request)))
 						.flatMap(response -> wrapException(() -> response.writeTo(exchange,
 								new HandlerStrategiesResponseContext(this.strategies))));
@@ -1261,11 +1260,6 @@ public abstract class RouterFunctions {
 			attributes.put(REQUEST_ATTRIBUTE, request);
 		}
 
-		@SuppressWarnings("unchecked")
-		private static <T extends ServerResponse> HandlerFunction<T> notFound() {
-			return (HandlerFunction<T>) NOT_FOUND_HANDLER;
-		}
-
 		private static <T> Mono<T> wrapException(Supplier<Mono<T>> supplier) {
 			try {
 				return supplier.get();
diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/function/server/RouterFunctionsTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/function/server/RouterFunctionsTests.java
index 46b29bdf6cd..ad6bfbdca9f 100644
--- a/spring-webflux/src/test/java/org/springframework/web/reactive/function/server/RouterFunctionsTests.java
+++ b/spring-webflux/src/test/java/org/springframework/web/reactive/function/server/RouterFunctionsTests.java
@@ -16,14 +16,17 @@
 
 package org.springframework.web.reactive.function.server;
 
+import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
+import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseCookie;
@@ -191,6 +194,27 @@ public class RouterFunctionsTests {
 		assertThat(httpResponse.getStatusCode()).isEqualTo(HttpStatus.NOT_FOUND);
 	}
 
+	@Test
+	public void toHttpHandlerRouteNotFoundReturnsResponseStatusException() {
+		HandlerFunction<ServerResponse> handlerFunction = request -> ServerResponse.accepted().build();
+		RouterFunction<ServerResponse> routerFunction =
+				RouterFunctions.route(RequestPredicates.GET("/path"), handlerFunction);
+
+		HandlerStrategies handlerStrategies = HandlerStrategies.empty().exceptionHandler((exchange, ex) -> {
+			exchange.getResponse().setStatusCode(HttpStatus.NOT_FOUND);
+			DataBuffer buffer = exchange.getResponse().bufferFactory().wrap("Custom response".getBytes(StandardCharsets.UTF_8));
+			return exchange.getResponse().writeWith(Flux.just(buffer));
+		}).build();
+		HttpHandler result = RouterFunctions.toHttpHandler(routerFunction, handlerStrategies);
+		assertThat(result).isNotNull();
+
+		MockServerHttpRequest httpRequest = MockServerHttpRequest.get("https://localhost").build();
+		MockServerHttpResponse httpResponse = new MockServerHttpResponse();
+		result.handle(httpRequest, httpResponse).block();
+		assertThat(httpResponse.getStatusCode()).isEqualTo(HttpStatus.NOT_FOUND);
+		assertThat(httpResponse.getBodyAsString().block()).isEqualTo("Custom response");
+	}
+
 	@Test
 	public void toHttpHandlerHandlerReturnResponseStatusExceptionInResponseWriteTo() {
 		HandlerFunction<ServerResponse> handlerFunction =