Fix RSocket websocket config with WebFlux

In the case of a WebFlux + RSocket over websocket setup, the RSocket
auto-configuration would not set up the required routes; only the
websocket endpoint for RSocket would be available, overriding the
handler configured for WebFlux.

This commit introduces `NettyRouteProvider`. Components implementing
that interface can contribute HTTP routes to the Reactor Netty server
being built.

* if none is provided, the regular handler setup is used
* if one or more routes are provided, routes are sorted and added before
the WebFlux handler (acting as a default)

Fixes gh-16826

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/rsocket/RSocketServerAutoConfiguration.java

@@ -38,8 +38,6 @@ import org.springframework.boot.rsocket.netty.NettyRSocketServerFactory;
 import org.springframework.boot.rsocket.server.RSocketServerBootstrap;
 import org.springframework.boot.rsocket.server.RSocketServerFactory;
 import org.springframework.boot.rsocket.server.ServerRSocketFactoryCustomizer;
-import org.springframework.boot.web.embedded.netty.NettyReactiveWebServerFactory;
-import org.springframework.boot.web.server.WebServerFactoryCustomizer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Conditional;
 import org.springframework.context.annotation.Configuration;
@@ -70,12 +68,11 @@ public class RSocketServerAutoConfiguration {
 	static class WebFluxServerAutoConfiguration {
 
 		@Bean
-		public WebServerFactoryCustomizer<NettyReactiveWebServerFactory> rSocketWebsocketCustomizer(
+		public RSocketWebSocketNettyRouteProvider rSocketWebsocketRouteProvider(
 				RSocketProperties properties,
 				MessageHandlerAcceptor messageHandlerAcceptor) {
-			RSocketNettyServerCustomizer customizer = new RSocketNettyServerCustomizer(
+			return new RSocketWebSocketNettyRouteProvider(
 					properties.getServer().getMappingPath(), messageHandlerAcceptor);
-			return (factory) -> factory.addServerCustomizers(customizer);
 		}
 
 	}

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/rsocket/RSocketWebSocketNettyRouteProvider.java

@@ -19,34 +19,34 @@ package org.springframework.boot.autoconfigure.rsocket;
 import io.rsocket.RSocketFactory;
 import io.rsocket.transport.ServerTransport;
 import io.rsocket.transport.netty.server.WebsocketRouteTransport;
-import reactor.netty.http.server.HttpServer;
+import reactor.netty.http.server.HttpServerRoutes;
 
-import org.springframework.boot.web.embedded.netty.NettyServerCustomizer;
+import org.springframework.boot.web.embedded.netty.NettyRouteProvider;
 import org.springframework.messaging.rsocket.MessageHandlerAcceptor;
 
 /**
- * {@link NettyServerCustomizer} that configures an RSocket Websocket endpoint.
+ * {@link NettyRouteProvider} that configures an RSocket Websocket endpoint.
  *
  * @author Brian Clozel
  */
-class RSocketNettyServerCustomizer implements NettyServerCustomizer {
+class RSocketWebSocketNettyRouteProvider implements NettyRouteProvider {
 
 	private final String mappingPath;
 
 	private final MessageHandlerAcceptor messageHandlerAcceptor;
 
-	RSocketNettyServerCustomizer(String mappingPath,
+	RSocketWebSocketNettyRouteProvider(String mappingPath,
 			MessageHandlerAcceptor messageHandlerAcceptor) {
 		this.mappingPath = mappingPath;
 		this.messageHandlerAcceptor = messageHandlerAcceptor;
 	}
 
 	@Override
-	public HttpServer apply(HttpServer httpServer) {
+	public HttpServerRoutes apply(HttpServerRoutes httpServerRoutes) {
 		ServerTransport.ConnectionAcceptor acceptor = RSocketFactory.receive()
 				.acceptor(this.messageHandlerAcceptor).toConnectionAcceptor();
-		return httpServer.route((routes) -> routes.ws(this.mappingPath,
-				WebsocketRouteTransport.newHandler(acceptor)));
+		return httpServerRoutes.ws(this.mappingPath,
+				WebsocketRouteTransport.newHandler(acceptor));
 	}
 
 }

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/reactive/ReactiveWebServerFactoryConfiguration.java

@@ -27,6 +27,7 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean
 import org.springframework.boot.web.embedded.jetty.JettyReactiveWebServerFactory;
 import org.springframework.boot.web.embedded.jetty.JettyServerCustomizer;
 import org.springframework.boot.web.embedded.netty.NettyReactiveWebServerFactory;
+import org.springframework.boot.web.embedded.netty.NettyRouteProvider;
 import org.springframework.boot.web.embedded.tomcat.TomcatConnectorCustomizer;
 import org.springframework.boot.web.embedded.tomcat.TomcatContextCustomizer;
 import org.springframework.boot.web.embedded.tomcat.TomcatProtocolHandlerCustomizer;
@@ -65,9 +66,12 @@ abstract class ReactiveWebServerFactoryConfiguration {
 
 		@Bean
 		public NettyReactiveWebServerFactory nettyReactiveWebServerFactory(
-				ReactorResourceFactory resourceFactory) {
+				ReactorResourceFactory resourceFactory,
+				ObjectProvider<NettyRouteProvider> routes) {
 			NettyReactiveWebServerFactory serverFactory = new NettyReactiveWebServerFactory();
 			serverFactory.setResourceFactory(resourceFactory);
+			routes.orderedStream()
+					.forEach((route) -> serverFactory.addRouteProviders(route));
 			return serverFactory;
 		}
 

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/rsocket/RSocketServerAutoConfigurationTests.java

@@ -77,8 +77,7 @@ public class RSocketServerAutoConfigurationTests {
 				.withPropertyValues("spring.rsocket.server.transport=websocket",
 						"spring.rsocket.server.mapping-path=/rsocket")
 				.run((context) -> assertThat(context)
-						.getBeanNames(WebServerFactoryCustomizer.class).hasSize(1)
-						.containsOnly("rSocketWebsocketCustomizer"));
+						.hasSingleBean(RSocketWebSocketNettyRouteProvider.class));
 	}
 
 	@Test

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/rsocket/RSocketWebSocketNettyRouteProviderTests.java

@@ -0,0 +1,162 @@
+/*
+ * Copyright 2012-2019 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.autoconfigure.rsocket;
+
+import java.net.URI;
+import java.time.Duration;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.autoconfigure.context.PropertyPlaceholderAutoConfiguration;
+import org.springframework.boot.autoconfigure.http.codec.CodecsAutoConfiguration;
+import org.springframework.boot.autoconfigure.jackson.JacksonAutoConfiguration;
+import org.springframework.boot.autoconfigure.web.reactive.HttpHandlerAutoConfiguration;
+import org.springframework.boot.autoconfigure.web.reactive.WebFluxAutoConfiguration;
+import org.springframework.boot.autoconfigure.web.reactive.error.ErrorWebFluxAutoConfiguration;
+import org.springframework.boot.test.context.runner.ReactiveWebApplicationContextRunner;
+import org.springframework.boot.web.embedded.netty.NettyReactiveWebServerFactory;
+import org.springframework.boot.web.reactive.context.AnnotationConfigReactiveWebServerApplicationContext;
+import org.springframework.boot.web.reactive.context.ReactiveWebServerApplicationContext;
+import org.springframework.boot.web.server.WebServer;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.http.MediaType;
+import org.springframework.messaging.handler.annotation.MessageMapping;
+import org.springframework.messaging.rsocket.RSocketRequester;
+import org.springframework.stereotype.Controller;
+import org.springframework.test.web.reactive.server.WebTestClient;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.ResponseBody;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for {@link RSocketWebSocketNettyRouteProvider}.
+ *
+ * @author Brian Clozel
+ */
+
+public class RSocketWebSocketNettyRouteProviderTests {
+
+	@Test
+	public void webEndpointsShouldWork() throws Exception {
+		new ReactiveWebApplicationContextRunner(
+				AnnotationConfigReactiveWebServerApplicationContext::new)
+						.withConfiguration(
+								AutoConfigurations.of(HttpHandlerAutoConfiguration.class,
+										WebFluxAutoConfiguration.class,
+										ErrorWebFluxAutoConfiguration.class,
+										PropertyPlaceholderAutoConfiguration.class,
+										JacksonAutoConfiguration.class,
+										CodecsAutoConfiguration.class,
+										RSocketStrategiesAutoConfiguration.class,
+										RSocketServerAutoConfiguration.class,
+										RSocketMessagingAutoConfiguration.class,
+										RSocketRequesterAutoConfiguration.class))
+						.withUserConfiguration(WebConfiguration.class)
+						.withPropertyValues("spring.rsocket.server.transport=websocket",
+								"spring.rsocket.server.mapping-path=/rsocket")
+						.run((context) -> {
+							ReactiveWebServerApplicationContext serverContext = (ReactiveWebServerApplicationContext) context
+									.getSourceApplicationContext();
+							RSocketRequester requester = createRSocketRequester(context,
+									serverContext.getWebServer());
+							TestProtocol rsocketResponse = requester.route("websocket")
+									.data(new TestProtocol("rsocket"))
+									.retrieveMono(TestProtocol.class)
+									.block(Duration.ofSeconds(3));
+							assertThat(rsocketResponse.getName()).isEqualTo("rsocket");
+							WebTestClient client = createWebTestClient(
+									serverContext.getWebServer());
+							client.get().uri("/protocol").exchange().expectStatus().isOk()
+									.expectBody().jsonPath("name", "http");
+						});
+	}
+
+	private WebTestClient createWebTestClient(WebServer server) {
+		return WebTestClient.bindToServer()
+				.baseUrl("http://localhost:" + server.getPort()).build();
+	}
+
+	private RSocketRequester createRSocketRequester(ApplicationContext context,
+			WebServer server) {
+		int port = server.getPort();
+		RSocketRequester.Builder builder = context
+				.getBean(RSocketRequester.Builder.class);
+		return builder.connectWebSocket(URI.create("ws://localhost:" + port + "/rsocket"))
+				.block();
+	}
+
+	@Configuration(proxyBeanMethods = false)
+	static class WebConfiguration {
+
+		@Bean
+		public WebController webController() {
+			return new WebController();
+		}
+
+		@Bean
+		public NettyReactiveWebServerFactory customServerFactory(
+				RSocketWebSocketNettyRouteProvider routeProvider) {
+			NettyReactiveWebServerFactory serverFactory = new NettyReactiveWebServerFactory(
+					0);
+			serverFactory.addRouteProviders(routeProvider);
+			return serverFactory;
+		}
+
+	}
+
+	@Controller
+	static class WebController {
+
+		@GetMapping(path = "/protocol", produces = MediaType.APPLICATION_JSON_VALUE)
+		@ResponseBody
+		public TestProtocol testWebEndpoint() {
+			return new TestProtocol("http");
+		}
+
+		@MessageMapping("websocket")
+		public TestProtocol testRSocketEndpoint() {
+			return new TestProtocol("rsocket");
+		}
+
+	}
+
+	static class TestProtocol {
+
+		private String name;
+
+		TestProtocol() {
+		}
+
+		TestProtocol(String name) {
+			this.name = name;
+		}
+
+		public String getName() {
+			return this.name;
+		}
+
+		public void setName(String name) {
+			this.name = name;
+		}
+
+	}
+
+}

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/web/embedded/netty/NettyReactiveWebServerFactory.java

@@ -45,6 +45,8 @@ public class NettyReactiveWebServerFactory extends AbstractReactiveWebServerFact
 
 	private List<NettyServerCustomizer> serverCustomizers = new ArrayList<>();
 
+	private List<NettyRouteProvider> routeProviders = new ArrayList<>();
+
 	private Duration lifecycleTimeout;
 
 	private boolean useForwardHeaders;
@@ -63,7 +65,10 @@ public class NettyReactiveWebServerFactory extends AbstractReactiveWebServerFact
 		HttpServer httpServer = createHttpServer();
 		ReactorHttpHandlerAdapter handlerAdapter = new ReactorHttpHandlerAdapter(
 				httpHandler);
-		return new NettyWebServer(httpServer, handlerAdapter, this.lifecycleTimeout);
+		NettyWebServer webServer = new NettyWebServer(httpServer, handlerAdapter,
+				this.lifecycleTimeout);
+		webServer.setRouteProviders(this.routeProviders);
+		return webServer;
 	}
 
 	/**
@@ -95,6 +100,16 @@ public class NettyReactiveWebServerFactory extends AbstractReactiveWebServerFact
 		this.serverCustomizers.addAll(Arrays.asList(serverCustomizers));
 	}
 
+	/**
+	 * Add {@link NettyRouteProvider}s that should be applied, in order, before the the
+	 * handler for the Spring application.
+	 * @param routeProviders the route providers to add
+	 */
+	public void addRouteProviders(NettyRouteProvider... routeProviders) {
+		Assert.notNull(routeProviders, "NettyRouteProvider must not be null");
+		this.routeProviders.addAll(Arrays.asList(routeProviders));
+	}
+
 	/**
 	 * Set the maximum amount of time that should be waited when starting or stopping the
 	 * server.

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/web/embedded/netty/NettyRouteProvider.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2012-2019 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.web.embedded.netty;
+
+import java.util.function.Function;
+
+import reactor.netty.http.server.HttpServerRoutes;
+
+/**
+ * Function that can add new routes to an {@link HttpServerRoutes} instance.
+ *
+ * @author Brian Clozel
+ * @see NettyReactiveWebServerFactory
+ * @since 2.2.0
+ */
+@FunctionalInterface
+public interface NettyRouteProvider extends Function<HttpServerRoutes, HttpServerRoutes> {
+
+}

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/web/embedded/netty/NettyWebServer.java

@@ -17,6 +17,8 @@
 package org.springframework.boot.web.embedded.netty;
 
 import java.time.Duration;
+import java.util.Collections;
+import java.util.List;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -50,6 +52,8 @@ public class NettyWebServer implements WebServer {
 
 	private final Duration lifecycleTimeout;
 
+	private List<NettyRouteProvider> routeProviders = Collections.emptyList();
+
 	private DisposableServer disposableServer;
 
 	public NettyWebServer(HttpServer httpServer, ReactorHttpHandlerAdapter handlerAdapter,
@@ -61,6 +65,10 @@ public class NettyWebServer implements WebServer {
 		this.lifecycleTimeout = lifecycleTimeout;
 	}
 
+	public void setRouteProviders(List<NettyRouteProvider> routeProviders) {
+		this.routeProviders = routeProviders;
+	}
+
 	@Override
 	public void start() throws WebServerException {
 		if (this.disposableServer == null) {
@@ -80,11 +88,20 @@ public class NettyWebServer implements WebServer {
 	}
 
 	private DisposableServer startHttpServer() {
-		if (this.lifecycleTimeout != null) {
-			return this.httpServer.handle(this.handlerAdapter)
-					.bindNow(this.lifecycleTimeout);
+		HttpServer server = this.httpServer;
+		if (this.routeProviders.isEmpty()) {
+			server = server.handle(this.handlerAdapter);
 		}
-		return this.httpServer.handle(this.handlerAdapter).bindNow();
+		else {
+			server = server.route((routes) -> {
+				this.routeProviders.forEach((provider) -> provider.apply(routes));
+				routes.route((r) -> true, this.handlerAdapter);
+			});
+		}
+		if (this.lifecycleTimeout != null) {
+			return server.bindNow(this.lifecycleTimeout);
+		}
+		return server.bindNow();
 	}
 
 	private ChannelBindException findBindException(Exception ex) {