Avoid creating a new EmbeddedServletContainerFactory for websockets

User can now also switch off and customize the websockets customizer by adding
a bean named "websocketContainerCustomizer".

Fixes gh-479

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/aop/AopAutoConfiguration.java

@@ -18,7 +18,6 @@ package org.springframework.boot.autoconfigure.aop;
 
 import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.reflect.Advice;
-import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnExpression;
 import org.springframework.context.annotation.Configuration;
@@ -26,11 +25,14 @@ import org.springframework.context.annotation.EnableAspectJAutoProxy;
 
 /**
  * 
- * <p> {@link org.springframework.boot.autoconfigure.EnableAutoConfiguration Auto-configuration}
- * for Spring's  AOP support. Equivalent to enabling {@link org.springframework.context.annotation.EnableAspectJAutoProxy}
- * in your configuration. The configuration will not be activated if {@literal spring.aop.auto=false}.
- * The {@literal proxyTargetClass} attribute will be {@literal false}, by default, but can be overridden by
- * specifying {@literal spring.aop.proxyTargetClass=true}.
+ * <p>
+ * {@link org.springframework.boot.autoconfigure.EnableAutoConfiguration
+ * Auto-configuration} for Spring's AOP support. Equivalent to enabling
+ * {@link org.springframework.context.annotation.EnableAspectJAutoProxy} in your
+ * configuration. The configuration will not be activated if
+ * {@literal spring.aop.auto=false}. The {@literal proxyTargetClass} attribute will be
+ * {@literal false}, by default, but can be overridden by specifying
+ * {@literal spring.aop.proxyTargetClass=true}.
  * 
  * @author Dave Syer
  * @author Josh Long

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/ServerProperties.java

@@ -53,8 +53,7 @@ public class ServerProperties implements EmbeddedServletContainerCustomizer {
 
 	private Integer sessionTimeout;
 
-	@NotNull
-	private String contextPath = "";
+	private String contextPath;
 
 	@NotNull
 	private String servletPath = "/";

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

@@ -23,7 +23,11 @@ import org.apache.catalina.deploy.ApplicationListener;
 import org.apache.catalina.startup.Tomcat;
 import org.springframework.boot.autoconfigure.AutoConfigureBefore;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.web.EmbeddedServletContainerAutoConfiguration;
+import org.springframework.boot.context.embedded.ConfigurableEmbeddedServletContainer;
+import org.springframework.boot.context.embedded.EmbeddedServletContainerCustomizer;
+import org.springframework.boot.context.embedded.tomcat.TomcatContextCustomizer;
 import org.springframework.boot.context.embedded.tomcat.TomcatEmbeddedServletContainerFactory;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -47,14 +51,31 @@ public class WebSocketAutoConfiguration {
 			"org.apache.tomcat.websocket.server.WsContextListener", false);
 
 	@Bean
-	public TomcatEmbeddedServletContainerFactory tomcatEmbeddedServletContainerFactory() {
-		TomcatEmbeddedServletContainerFactory factory = new TomcatEmbeddedServletContainerFactory() {
+	@ConditionalOnMissingBean(name = "websocketContainerCustomizer")
+	public EmbeddedServletContainerCustomizer websocketContainerCustomizer() {
+
+		EmbeddedServletContainerCustomizer customizer = new EmbeddedServletContainerCustomizer() {
+
 			@Override
-			protected void postProcessContext(Context context) {
-				context.addApplicationListener(WS_APPLICATION_LISTENER);
+			public void customize(ConfigurableEmbeddedServletContainer container) {
+				if (!(container instanceof TomcatEmbeddedServletContainerFactory)) {
+					throw new IllegalStateException(
+							"Websockets are currently only supported in Tomcat (found "
+									+ container.getClass() + ")");
+				}
+				((TomcatEmbeddedServletContainerFactory) container)
+						.addContextCustomizers(new TomcatContextCustomizer() {
+							@Override
+							public void customize(Context context) {
+								context.addApplicationListener(WS_APPLICATION_LISTENER);
+							}
+						});
 			}
+
 		};
-		return factory;
+
+		return customizer;
+
 	}
 
 }

spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/ServerPropertiesTests.java

@@ -29,6 +29,7 @@ import org.springframework.boot.context.embedded.ConfigurableEmbeddedServletCont
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 /**
@@ -75,7 +76,7 @@ public class ServerPropertiesTests {
 	public void testCustomizeTomcat() throws Exception {
 		ConfigurableEmbeddedServletContainer factory = mock(ConfigurableEmbeddedServletContainer.class);
 		this.properties.customize(factory);
-		verify(factory).setContextPath("");
+		verify(factory, times(0)).setContextPath("");
 	}
 
 	@Test

spring-boot-samples/spring-boot-sample-websocket/src/test/java/samples/websocket/echo/CustomContainerWebSocketsApplicationTests.java

@@ -0,0 +1,115 @@
+/*
+ * Copyright 2012-2013 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 samples.websocket.echo;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.boot.CommandLineRunner;
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.context.embedded.EmbeddedServletContainerFactory;
+import org.springframework.boot.context.embedded.tomcat.TomcatEmbeddedServletContainerFactory;
+import org.springframework.boot.test.IntegrationTest;
+import org.springframework.boot.test.SpringApplicationConfiguration;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.test.context.web.WebAppConfiguration;
+import org.springframework.web.socket.client.WebSocketConnectionManager;
+import org.springframework.web.socket.client.standard.StandardWebSocketClient;
+
+import samples.websocket.client.GreetingService;
+import samples.websocket.client.SimpleClientWebSocketHandler;
+import samples.websocket.client.SimpleGreetingService;
+import samples.websocket.config.SampleWebSocketsApplication;
+import samples.websocket.echo.CustomContainerWebSocketsApplicationTests.CustomContainerConfiguration;
+
+@RunWith(SpringJUnit4ClassRunner.class)
+@SpringApplicationConfiguration(classes={SampleWebSocketsApplication.class, CustomContainerConfiguration.class })
+@WebAppConfiguration
+@IntegrationTest
+@DirtiesContext
+public class CustomContainerWebSocketsApplicationTests {
+
+	private static Log logger = LogFactory.getLog(CustomContainerWebSocketsApplicationTests.class);
+
+	private static final String WS_URI = "ws://localhost:9010/ws/echo/websocket";
+	
+	@Configuration
+	protected static class CustomContainerConfiguration {
+		@Bean
+		public EmbeddedServletContainerFactory embeddedServletContainerFactory() {
+			return new TomcatEmbeddedServletContainerFactory("/ws", 9010);
+		}
+	}
+
+	@Test
+	public void runAndWait() throws Exception {
+		ConfigurableApplicationContext context = SpringApplication.run(
+				ClientConfiguration.class, "--spring.main.web_environment=false");
+		long count = context.getBean(ClientConfiguration.class).latch.getCount();
+		context.close();
+		assertEquals(0, count);
+	}
+
+	@Configuration
+	static class ClientConfiguration implements CommandLineRunner {
+
+		private final CountDownLatch latch = new CountDownLatch(1);
+
+		@Override
+		public void run(String... args) throws Exception {
+			logger.info("Waiting for response: latch=" + this.latch.getCount());
+			this.latch.await(10, TimeUnit.SECONDS);
+			logger.info("Got response: latch=" + this.latch.getCount());
+		}
+
+		@Bean
+		public WebSocketConnectionManager wsConnectionManager() {
+
+			WebSocketConnectionManager manager = new WebSocketConnectionManager(client(),
+					handler(), WS_URI);
+			manager.setAutoStartup(true);
+
+			return manager;
+		}
+
+		@Bean
+		public StandardWebSocketClient client() {
+			return new StandardWebSocketClient();
+		}
+
+		@Bean
+		public SimpleClientWebSocketHandler handler() {
+			return new SimpleClientWebSocketHandler(greetingService(), this.latch);
+		}
+
+		@Bean
+		public GreetingService greetingService() {
+			return new SimpleGreetingService();
+		}
+	}
+
+}