Initialize WebSocket infrastructure when using WebFlux and Jetty

In Spring Framework 5.x with Jetty 9, the reactive
JettyRequestUpgradeStrategy was able to initialize Jetty's WebSocket
infrastructure itself. With Jetty 10 this is no longer possible and
Boot must perform the initialization as part of preparing the
reactive JettyWebServer.

This commit updates the reactive WebSocket auto-configuration to
initialize Jetty's WebSocket infrastructure as part of creating the
reactive JettyWebServer.

Fixes gh-33347

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/websocket/reactive/JettyWebSocketReactiveWebServerCustomizer.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2012-2023 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.websocket.reactive;
+
+import jakarta.servlet.ServletContext;
+import org.eclipse.jetty.server.Handler;
+import org.eclipse.jetty.server.handler.HandlerCollection;
+import org.eclipse.jetty.server.handler.HandlerWrapper;
+import org.eclipse.jetty.servlet.ServletContextHandler;
+import org.eclipse.jetty.websocket.core.server.WebSocketMappings;
+import org.eclipse.jetty.websocket.core.server.WebSocketServerComponents;
+import org.eclipse.jetty.websocket.jakarta.server.internal.JakartaWebSocketServerContainer;
+import org.eclipse.jetty.websocket.server.JettyWebSocketServerContainer;
+import org.eclipse.jetty.websocket.servlet.WebSocketUpgradeFilter;
+
+import org.springframework.boot.web.embedded.jetty.JettyReactiveWebServerFactory;
+import org.springframework.boot.web.server.WebServerFactoryCustomizer;
+import org.springframework.core.Ordered;
+
+/**
+ * WebSocket customizer for {@link JettyReactiveWebServerFactory}.
+ *
+ * @author Andy Wilkinson
+ * @since 3.0.8
+ */
+public class JettyWebSocketReactiveWebServerCustomizer
+		implements WebServerFactoryCustomizer<JettyReactiveWebServerFactory>, Ordered {
+
+	@Override
+	public void customize(JettyReactiveWebServerFactory factory) {
+		factory.addServerCustomizers((server) -> {
+			ServletContextHandler servletContextHandler = findServletContextHandler(server);
+			if (servletContextHandler != null) {
+				ServletContext servletContext = servletContextHandler.getServletContext();
+				if (JettyWebSocketServerContainer.getContainer(servletContext) == null) {
+					WebSocketServerComponents.ensureWebSocketComponents(server, servletContext);
+					JettyWebSocketServerContainer.ensureContainer(servletContext);
+				}
+				if (JakartaWebSocketServerContainer.getContainer(servletContext) == null) {
+					WebSocketServerComponents.ensureWebSocketComponents(server, servletContext);
+					WebSocketUpgradeFilter.ensureFilter(servletContext);
+					WebSocketMappings.ensureMappings(servletContext);
+					JakartaWebSocketServerContainer.ensureContainer(servletContext);
+				}
+			}
+		});
+	}
+
+	private ServletContextHandler findServletContextHandler(Handler handler) {
+		if (handler instanceof ServletContextHandler servletContextHandler) {
+			return servletContextHandler;
+		}
+		if (handler instanceof HandlerWrapper handlerWrapper) {
+			return findServletContextHandler(handlerWrapper.getHandler());
+		}
+		if (handler instanceof HandlerCollection handlerCollection) {
+			for (Handler contained : handlerCollection.getHandlers()) {
+				ServletContextHandler servletContextHandler = findServletContextHandler(contained);
+				if (servletContextHandler != null) {
+					return servletContextHandler;
+				}
+			}
+		}
+		return null;
+	}
+
+	@Override
+	public int getOrder() {
+		return 0;
+	}
+
+}

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/websocket/reactive/WebSocketReactiveAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2022 the original author or authors.
+ * Copyright 2012-2023 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.
@@ -20,6 +20,7 @@ import jakarta.servlet.Servlet;
 import jakarta.websocket.server.ServerContainer;
 import org.apache.catalina.startup.Tomcat;
 import org.apache.tomcat.websocket.server.WsSci;
+import org.eclipse.jetty.websocket.jakarta.server.config.JakartaWebSocketServletContainerInitializer;
 
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -57,4 +58,16 @@ public class WebSocketReactiveAutoConfiguration {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	@ConditionalOnClass(JakartaWebSocketServletContainerInitializer.class)
+	static class JettyWebSocketConfiguration {
+
+		@Bean
+		@ConditionalOnMissingBean(name = "websocketReactiveWebServerCustomizer")
+		JettyWebSocketReactiveWebServerCustomizer websocketServletWebServerCustomizer() {
+			return new JettyWebSocketReactiveWebServerCustomizer();
+		}
+
+	}
+
 }

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/websocket/reactive/WebSocketReactiveAutoConfigurationTests.java

@@ -0,0 +1,139 @@
+/*
+ * Copyright 2012-2023 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.websocket.reactive;
+
+import java.util.function.Function;
+import java.util.stream.Stream;
+
+import jakarta.servlet.ServletContext;
+import jakarta.websocket.server.ServerContainer;
+import org.apache.catalina.Container;
+import org.apache.catalina.Context;
+import org.apache.catalina.startup.Tomcat;
+import org.eclipse.jetty.servlet.ServletContextHandler;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import org.springframework.boot.testsupport.classpath.ForkedClassPath;
+import org.springframework.boot.testsupport.web.servlet.DirtiesUrlFactories;
+import org.springframework.boot.testsupport.web.servlet.Servlet5ClassPathOverrides;
+import org.springframework.boot.web.embedded.jetty.JettyReactiveWebServerFactory;
+import org.springframework.boot.web.embedded.jetty.JettyWebServer;
+import org.springframework.boot.web.embedded.tomcat.TomcatReactiveWebServerFactory;
+import org.springframework.boot.web.embedded.tomcat.TomcatWebServer;
+import org.springframework.boot.web.reactive.context.AnnotationConfigReactiveWebServerApplicationContext;
+import org.springframework.boot.web.reactive.server.ReactiveWebServerFactory;
+import org.springframework.boot.web.server.WebServerFactoryCustomizerBeanPostProcessor;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.http.server.reactive.HttpHandler;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for {@link WebSocketReactiveAutoConfiguration}.
+ *
+ * @author Andy Wilkinson
+ */
+@DirtiesUrlFactories
+class WebSocketReactiveAutoConfigurationTests {
+
+	@ParameterizedTest(name = "{0}")
+	@MethodSource("testConfiguration")
+	@ForkedClassPath
+	void serverContainerIsAvailableFromTheServletContext(String server,
+			Function<AnnotationConfigReactiveWebServerApplicationContext, ServletContext> servletContextAccessor,
+			Class<?>... configuration) {
+		try (AnnotationConfigReactiveWebServerApplicationContext context = new AnnotationConfigReactiveWebServerApplicationContext(
+				configuration)) {
+			Object serverContainer = servletContextAccessor.apply(context)
+				.getAttribute("jakarta.websocket.server.ServerContainer");
+			assertThat(serverContainer).isInstanceOf(ServerContainer.class);
+		}
+	}
+
+	static Stream<Arguments> testConfiguration() {
+		return Stream.of(Arguments.of("Jetty",
+				(Function<AnnotationConfigReactiveWebServerApplicationContext, ServletContext>) WebSocketReactiveAutoConfigurationTests::getJettyServletContext,
+				new Class<?>[] { JettyConfiguration.class,
+						WebSocketReactiveAutoConfiguration.JettyWebSocketConfiguration.class }),
+				Arguments.of("Tomcat",
+						(Function<AnnotationConfigReactiveWebServerApplicationContext, ServletContext>) WebSocketReactiveAutoConfigurationTests::getTomcatServletContext,
+						new Class<?>[] { TomcatConfiguration.class,
+								WebSocketReactiveAutoConfiguration.TomcatWebSocketConfiguration.class }));
+	}
+
+	private static ServletContext getJettyServletContext(AnnotationConfigReactiveWebServerApplicationContext context) {
+		return ((ServletContextHandler) ((JettyWebServer) context.getWebServer()).getServer().getHandler())
+			.getServletContext();
+	}
+
+	private static ServletContext getTomcatServletContext(AnnotationConfigReactiveWebServerApplicationContext context) {
+		return findContext(((TomcatWebServer) context.getWebServer()).getTomcat()).getServletContext();
+	}
+
+	private static Context findContext(Tomcat tomcat) {
+		for (Container child : tomcat.getHost().findChildren()) {
+			if (child instanceof Context context) {
+				return context;
+			}
+		}
+		throw new IllegalStateException("The host does not contain a Context");
+	}
+
+	@Configuration(proxyBeanMethods = false)
+	static class CommonConfiguration {
+
+		@Bean
+		static WebServerFactoryCustomizerBeanPostProcessor webServerFactoryCustomizerBeanPostProcessor() {
+			return new WebServerFactoryCustomizerBeanPostProcessor();
+		}
+
+		@Bean
+		HttpHandler echoHandler() {
+			return (request, response) -> response.writeWith(request.getBody());
+		}
+
+	}
+
+	@Configuration(proxyBeanMethods = false)
+	static class TomcatConfiguration extends CommonConfiguration {
+
+		@Bean
+		ReactiveWebServerFactory webServerFactory() {
+			TomcatReactiveWebServerFactory factory = new TomcatReactiveWebServerFactory();
+			factory.setPort(0);
+			return factory;
+		}
+
+	}
+
+	@Servlet5ClassPathOverrides
+	@Configuration(proxyBeanMethods = false)
+	static class JettyConfiguration extends CommonConfiguration {
+
+		@Bean
+		ReactiveWebServerFactory webServerFactory() {
+			JettyReactiveWebServerFactory factory = new JettyReactiveWebServerFactory();
+			factory.setPort(0);
+			return factory;
+		}
+
+	}
+
+}