Add Undertow WebSocket auto-configuration

Fixes gh-2028

Author: philwebb

spring-boot-autoconfigure/pom.xml

@@ -170,6 +170,11 @@
 			<artifactId>undertow-servlet</artifactId>
 			<optional>true</optional>
 		</dependency>
+		<dependency>
+			<groupId>io.undertow</groupId>
+			<artifactId>undertow-websockets-jsr</artifactId>
+			<optional>true</optional>
+		</dependency>
 		<dependency>
 			<groupId>org.freemarker</groupId>
 			<artifactId>freemarker</artifactId>

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/websocket/UndertowWebSocketContainerCustomizer.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2012-2014 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
+ *
+ *      http://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;
+
+import io.undertow.servlet.api.DeploymentInfo;
+import io.undertow.websockets.jsr.WebSocketDeploymentInfo;
+
+import org.springframework.boot.context.embedded.undertow.UndertowDeploymentInfoCustomizer;
+import org.springframework.boot.context.embedded.undertow.UndertowEmbeddedServletContainerFactory;
+
+/**
+ * {@link WebSocketContainerCustomizer} for
+ * {@link UndertowEmbeddedServletContainerFactory}.
+ *
+ * @author Phillip Webb
+ * @since 1.2.0
+ */
+public class UndertowWebSocketContainerCustomizer extends
+		WebSocketContainerCustomizer<UndertowEmbeddedServletContainerFactory> {
+
+	@Override
+	protected void doCustomize(UndertowEmbeddedServletContainerFactory container) {
+		WebsocketDeploymentInfoCustomizer customizer = new WebsocketDeploymentInfoCustomizer();
+		container.addDeploymentInfoCustomizers(customizer);
+	}
+
+	private static class WebsocketDeploymentInfoCustomizer implements
+			UndertowDeploymentInfoCustomizer {
+
+		@Override
+		public void customize(DeploymentInfo deploymentInfo) {
+			WebSocketDeploymentInfo info = new WebSocketDeploymentInfo();
+			deploymentInfo.addServletContextAttribute(
+					WebSocketDeploymentInfo.ATTRIBUTE_NAME, info);
+		}
+
+	}
+}

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/websocket/WebSocketAutoConfiguration.java

@@ -75,4 +75,16 @@ public EmbeddedServletContainerCustomizer websocketContainerCustomizer() {
 
 	}
 
+	@Configuration
+	@ConditionalOnClass(io.undertow.websockets.jsr.Bootstrap.class)
+	static class UndertowWebSocketConfiguration {
+
+		@Bean
+		@ConditionalOnMissingBean(name = "websocketContainerCustomizer")
+		public EmbeddedServletContainerCustomizer websocketContainerCustomizer() {
+			return new UndertowWebSocketContainerCustomizer();
+		}
+
+	}
+
 }

spring-boot-dependencies/pom.xml

@@ -563,6 +563,11 @@
 				<artifactId>undertow-servlet</artifactId>
 				<version>${undertow.version}</version>
 			</dependency>
+			<dependency>
+				<groupId>io.undertow</groupId>
+				<artifactId>undertow-websockets-jsr</artifactId>
+				<version>${undertow.version}</version>
+			</dependency>
 			<dependency>
 				<groupId>javax.cache</groupId>
 				<artifactId>cache-api</artifactId>

spring-boot-starters/spring-boot-starter-undertow/pom.xml

@@ -33,6 +33,10 @@
 				</exclusion>
 			</exclusions>
 		</dependency>
+		<dependency>
+			<groupId>io.undertow</groupId>
+			<artifactId>undertow-websockets-jsr</artifactId>
+		</dependency>
 		<dependency>
 			<groupId>javax.servlet</groupId>
 			<artifactId>javax.servlet-api</artifactId>

spring-boot/src/main/java/org/springframework/boot/context/embedded/undertow/UndertowDeploymentInfoCustomizer.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2012-2014 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
+ *
+ *      http://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.context.embedded.undertow;
+
+import io.undertow.servlet.api.DeploymentInfo;
+
+/**
+ * Callback interface that can be used to customize an Undertow {@link DeploymentInfo}.
+ *
+ * @author Phillip Webb
+ * @since 1.2.0
+ * @see UndertowEmbeddedServletContainerFactory
+ */
+public interface UndertowDeploymentInfoCustomizer {
+
+	/**
+	 * @param deploymentInfo the {@code DeploymentInfo} to customize
+	 */
+	void customize(DeploymentInfo deploymentInfo);
+
+}

spring-boot/src/main/java/org/springframework/boot/context/embedded/undertow/UndertowEmbeddedServletContainerFactory.java

@@ -86,7 +86,9 @@
 public class UndertowEmbeddedServletContainerFactory extends
 		AbstractEmbeddedServletContainerFactory implements ResourceLoaderAware {
 
-	private List<UndertowBuilderCustomizer> undertowBuilderCustomizers = new ArrayList<UndertowBuilderCustomizer>();
+	private List<UndertowBuilderCustomizer> builderCustomizers = new ArrayList<UndertowBuilderCustomizer>();
+
+	private List<UndertowDeploymentInfoCustomizer> deploymentInfoCustomizers = new ArrayList<UndertowDeploymentInfoCustomizer>();
 
 	private ResourceLoader resourceLoader;
 
@@ -132,35 +134,64 @@ public UndertowEmbeddedServletContainerFactory(String contextPath, int port) {
 	/**
 	 * Set {@link UndertowBuilderCustomizer}s that should be applied to the Undertow
 	 * {@link Builder}. Calling this method will replace any existing customizers.
-	 * @param undertowBuilderCustomizers the customizers to set
+	 * @param customizers the customizers to set
 	 */
-	public void setUndertowBuilderCustomizers(
-			Collection<? extends UndertowBuilderCustomizer> undertowBuilderCustomizers) {
-		Assert.notNull(undertowBuilderCustomizers,
-				"undertowBuilderCustomizers must not be null");
-		this.undertowBuilderCustomizers = new ArrayList<UndertowBuilderCustomizer>(
-				undertowBuilderCustomizers);
+	public void setBuilderCustomizers(
+			Collection<? extends UndertowBuilderCustomizer> customizers) {
+		Assert.notNull(customizers, "Customizers must not be null");
+		this.builderCustomizers = new ArrayList<UndertowBuilderCustomizer>(customizers);
 	}
 
 	/**
 	 * Returns a mutable collection of the {@link UndertowBuilderCustomizer}s that will be
 	 * applied to the Undertow {@link Builder} .
 	 * @return the customizers that will be applied
 	 */
-	public Collection<UndertowBuilderCustomizer> getUndertowBuilderCustomizers() {
-		return this.undertowBuilderCustomizers;
+	public Collection<UndertowBuilderCustomizer> getBuilderCustomizers() {
+		return this.builderCustomizers;
 	}
 
 	/**
 	 * Add {@link UndertowBuilderCustomizer}s that should be used to customize the
 	 * Undertow {@link Builder}.
-	 * @param undertowBuilderCustomizers the customizers to add
+	 * @param customizers the customizers to add
+	 */
+	public void addBuilderCustomizers(UndertowBuilderCustomizer... customizers) {
+		Assert.notNull(customizers, "Customizers must not be null");
+		this.builderCustomizers.addAll(Arrays.asList(customizers));
+	}
+
+	/**
+	 * Set {@link UndertowDeploymentInfoCustomizer}s that should be applied to the
+	 * Undertow {@link DeploymentInfo}. Calling this method will replace any existing
+	 * customizers.
+	 * @param customizers the customizers to set
 	 */
-	public void addUndertowBuilderCustomizers(
-			UndertowBuilderCustomizer... undertowBuilderCustomizers) {
-		Assert.notNull(undertowBuilderCustomizers,
-				"undertowBuilderCustomizers must not be null");
-		this.undertowBuilderCustomizers.addAll(Arrays.asList(undertowBuilderCustomizers));
+	public void setDeploymentInfoCustomizers(
+			Collection<? extends UndertowDeploymentInfoCustomizer> customizers) {
+		Assert.notNull(customizers, "Customizers must not be null");
+		this.deploymentInfoCustomizers = new ArrayList<UndertowDeploymentInfoCustomizer>(
+				customizers);
+	}
+
+	/**
+	 * Returns a mutable collection of the {@link UndertowDeploymentInfoCustomizer}s that
+	 * will be applied to the Undertow {@link DeploymentInfo} .
+	 * @return the customizers that will be applied
+	 */
+	public Collection<UndertowDeploymentInfoCustomizer> getDeploymentInfoCustomizers() {
+		return this.deploymentInfoCustomizers;
+	}
+
+	/**
+	 * Add {@link UndertowDeploymentInfoCustomizer}s that should be used to customize the
+	 * Undertow {@link DeploymentInfo}.
+	 * @param customizers the customizers to add
+	 */
+	public void addDeploymentInfoCustomizers(
+			UndertowDeploymentInfoCustomizer... customizers) {
+		Assert.notNull(customizers, "UndertowDeploymentInfoCustomizers must not be null");
+		this.deploymentInfoCustomizers.addAll(Arrays.asList(customizers));
 	}
 
 	@Override
@@ -199,7 +230,7 @@ private Builder createBuilder(int port) {
 		else {
 			configureSsl(port, builder);
 		}
-		for (UndertowBuilderCustomizer customizer : this.undertowBuilderCustomizers) {
+		for (UndertowBuilderCustomizer customizer : this.builderCustomizers) {
 			customizer.customize(builder);
 		}
 		return builder;
@@ -298,6 +329,9 @@ private DeploymentManager createDeploymentManager(
 		deployment.setServletStackTraces(ServletStackTraces.NONE);
 		deployment.setResourceManager(getDocumentRootResourceManager());
 		configureMimeMappings(deployment);
+		for (UndertowDeploymentInfoCustomizer customizer : this.deploymentInfoCustomizers) {
+			customizer.customize(deployment);
+		}
 		DeploymentManager manager = Servlets.defaultContainer().addDeployment(deployment);
 		manager.deploy();
 		SessionManager sessionManager = manager.getDeployment().getSessionManager();

spring-boot/src/test/java/org/springframework/boot/context/embedded/undertow/UndertowEmbeddedServletContainerFactoryTests.java

@@ -17,6 +17,7 @@
 package org.springframework.boot.context.embedded.undertow;
 
 import io.undertow.Undertow.Builder;
+import io.undertow.servlet.api.DeploymentInfo;
 
 import java.util.Arrays;
 
@@ -62,19 +63,19 @@ public void errorPage404() throws Exception {
 	}
 
 	@Test
-	public void setNullUndertowBuilderCustomizersThrows() {
+	public void setNullBuilderCustomizersThrows() {
 		UndertowEmbeddedServletContainerFactory factory = getFactory();
 		this.thrown.expect(IllegalArgumentException.class);
-		this.thrown.expectMessage("undertowBuilderCustomizers must not be null");
-		factory.setUndertowBuilderCustomizers(null);
+		this.thrown.expectMessage("Customizers must not be null");
+		factory.setBuilderCustomizers(null);
 	}
 
 	@Test
-	public void addNullContextCustomizersThrows() {
+	public void addNullAddBuilderCustomizersThrows() {
 		UndertowEmbeddedServletContainerFactory factory = getFactory();
 		this.thrown.expect(IllegalArgumentException.class);
-		this.thrown.expectMessage("undertowBuilderCustomizers must not be null");
-		factory.addUndertowBuilderCustomizers((UndertowBuilderCustomizer[]) null);
+		this.thrown.expectMessage("Customizers must not be null");
+		factory.addBuilderCustomizers((UndertowBuilderCustomizer[]) null);
 	}
 
 	@Test
@@ -84,16 +85,48 @@ public void builderCustomizers() throws Exception {
 		for (int i = 0; i < customizers.length; i++) {
 			customizers[i] = mock(UndertowBuilderCustomizer.class);
 		}
-		factory.setUndertowBuilderCustomizers(Arrays.asList(customizers[0],
-				customizers[1]));
-		factory.addUndertowBuilderCustomizers(customizers[2], customizers[3]);
+		factory.setBuilderCustomizers(Arrays.asList(customizers[0], customizers[1]));
+		factory.addBuilderCustomizers(customizers[2], customizers[3]);
 		this.container = factory.getEmbeddedServletContainer();
 		InOrder ordered = inOrder((Object[]) customizers);
 		for (UndertowBuilderCustomizer customizer : customizers) {
 			ordered.verify(customizer).customize((Builder) anyObject());
 		}
 	}
 
+	@Test
+	public void setNullDeploymentInfoCustomizersThrows() {
+		UndertowEmbeddedServletContainerFactory factory = getFactory();
+		this.thrown.expect(IllegalArgumentException.class);
+		this.thrown.expectMessage("Customizers must not be null");
+		factory.setDeploymentInfoCustomizers(null);
+	}
+
+	@Test
+	public void addNullAddDeploymentInfoCustomizersThrows() {
+		UndertowEmbeddedServletContainerFactory factory = getFactory();
+		this.thrown.expect(IllegalArgumentException.class);
+		this.thrown.expectMessage("Customizers must not be null");
+		factory.addDeploymentInfoCustomizers((UndertowDeploymentInfoCustomizer[]) null);
+	}
+
+	@Test
+	public void deploymentInfo() throws Exception {
+		UndertowEmbeddedServletContainerFactory factory = getFactory();
+		UndertowDeploymentInfoCustomizer[] customizers = new UndertowDeploymentInfoCustomizer[4];
+		for (int i = 0; i < customizers.length; i++) {
+			customizers[i] = mock(UndertowDeploymentInfoCustomizer.class);
+		}
+		factory.setDeploymentInfoCustomizers(Arrays
+				.asList(customizers[0], customizers[1]));
+		factory.addDeploymentInfoCustomizers(customizers[2], customizers[3]);
+		this.container = factory.getEmbeddedServletContainer();
+		InOrder ordered = inOrder((Object[]) customizers);
+		for (UndertowDeploymentInfoCustomizer customizer : customizers) {
+			ordered.verify(customizer).customize((DeploymentInfo) anyObject());
+		}
+	}
+
 	@Test
 	public void basicSslClasspathKeyStore() throws Exception {
 		testBasicSllWithKeystore("classpath:test.jks");