ServerEndpointExporter can initialize itself based on a late-provided ServletContext as well (for Boot)

Also allows for direct setting of a ServerContainer and for custom triggering of endpoint registration.

Issue: SPR-12109

spring-websocket/src/main/java/org/springframework/web/socket/server/standard/ServerEndpointExporter.java

@@ -16,27 +16,21 @@
 
 package org.springframework.web.socket.server.standard;
 
-import java.lang.reflect.Method;
-import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.LinkedHashSet;
 import java.util.List;
-import java.util.Map;
+import java.util.Set;
+import javax.servlet.ServletContext;
 import javax.websocket.DeploymentException;
 import javax.websocket.server.ServerContainer;
 import javax.websocket.server.ServerEndpoint;
 import javax.websocket.server.ServerEndpointConfig;
 
-import org.apache.commons.logging.Log;
-import org.apache.commons.logging.LogFactory;
-
-import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.context.ApplicationContext;
-import org.springframework.context.ApplicationContextAware;
 import org.springframework.util.Assert;
-import org.springframework.util.ClassUtils;
+import org.springframework.web.context.support.WebApplicationObjectSupport;
-import org.springframework.util.ReflectionUtils;
 
 /**
  * Detects beans of type {@link javax.websocket.server.ServerEndpointConfig} and registers
@@ -50,24 +44,36 @@ import org.springframework.util.ReflectionUtils;
  * done with the help of the {@code <absolute-ordering>} element in web.xml.
  *
  * @author Rossen Stoyanchev
+ * @author Juergen Hoeller
  * @since 4.0
  * @see ServerEndpointRegistration
  * @see SpringConfigurator
  * @see ServletServerContainerFactoryBean
  */
-public class ServerEndpointExporter implements InitializingBean, BeanPostProcessor, ApplicationContextAware {
+public class ServerEndpointExporter extends WebApplicationObjectSupport implements BeanPostProcessor, InitializingBean {
-
-	private static final Log logger = LogFactory.getLog(ServerEndpointExporter.class);
-
-
-	private final List<Class<?>> annotatedEndpointClasses = new ArrayList<Class<?>>();
-
-	private final List<Class<?>> annotatedEndpointBeanTypes = new ArrayList<Class<?>>();
-
-	private ApplicationContext applicationContext;
 
 	private ServerContainer serverContainer;
 
+	private List<Class<?>> annotatedEndpointClasses;
+
+	private Set<Class<?>> annotatedEndpointBeanTypes;
+
+
+	/**
+	 * Set the JSR-356 {@link ServerContainer} to use for endpoint registration.
+	 * If not set, the container is going to be retrieved via the {@code ServletContext}.
+	 * @since 4.1
+	 */
+	public void setServerContainer(ServerContainer serverContainer) {
+		this.serverContainer = serverContainer;
+	}
+
+	/**
+	 * Return the JSR-356 {@link ServerContainer} to use for endpoint registration.
+	 */
+	protected ServerContainer getServerContainer() {
+		return this.serverContainer;
+	}
 
 	/**
 	 * Explicitly list annotated endpoint types that should be registered on startup. This
@@ -76,17 +82,19 @@ public class ServerEndpointExporter implements InitializingBean, BeanPostProcess
 	 * @param annotatedEndpointClasses {@link ServerEndpoint}-annotated types
  	 */
 	public void setAnnotatedEndpointClasses(Class<?>... annotatedEndpointClasses) {
-		this.annotatedEndpointClasses.clear();
+		this.annotatedEndpointClasses = Arrays.asList(annotatedEndpointClasses);
-		this.annotatedEndpointClasses.addAll(Arrays.asList(annotatedEndpointClasses));
 	}
 
 	@Override
-	public void setApplicationContext(ApplicationContext applicationContext) {
+	protected void initApplicationContext(ApplicationContext context) {
-		this.applicationContext = applicationContext;
+		// Initializes ServletContext given a WebApplicationContext
-		this.serverContainer = getServerContainer();
+		super.initApplicationContext(context);
-		Map<String, Object> beans = applicationContext.getBeansWithAnnotation(ServerEndpoint.class);
+
-		for (String beanName : beans.keySet()) {
+		// Retrieve beans which are annotated with @ServerEndpoint
-			Class<?> beanType = applicationContext.getType(beanName);
+		this.annotatedEndpointBeanTypes = new LinkedHashSet<Class<?>>();
+		String[] beanNames = context.getBeanNamesForAnnotation(ServerEndpoint.class);
+		for (String beanName : beanNames) {
+			Class<?> beanType = context.getType(beanName);
 			if (logger.isInfoEnabled()) {
 				logger.info("Detected @ServerEndpoint bean '" + beanName + "', registering it as an endpoint by type");
 			}
@@ -94,65 +102,71 @@ public class ServerEndpointExporter implements InitializingBean, BeanPostProcess
 		}
 	}
 
-	protected ServerContainer getServerContainer() {
+	@Override
-		Class<?> servletContextClass;
+	protected void initServletContext(ServletContext servletContext) {
-		try {
+		if (this.serverContainer == null) {
-			servletContextClass = ClassUtils.forName("javax.servlet.ServletContext", getClass().getClassLoader());
+			this.serverContainer =
-		}
+					(ServerContainer) servletContext.getAttribute("javax.websocket.server.ServerContainer");
-		catch (Throwable ex) {
-			return null;
-		}
-
-		try {
-			Method getter = ReflectionUtils.findMethod(this.applicationContext.getClass(), "getServletContext");
-			Object servletContext = getter.invoke(this.applicationContext);
-			Method attrMethod = ReflectionUtils.findMethod(servletContextClass, "getAttribute", String.class);
-			return (ServerContainer) attrMethod.invoke(servletContext, "javax.websocket.server.ServerContainer");
-		}
-		catch (Exception ex) {
-			throw new IllegalStateException(
-					"Failed to get javax.websocket.server.ServerContainer via ServletContext attribute", ex);
 		}
 	}
 
+
 	@Override
-	public void afterPropertiesSet() throws Exception {
+	public void afterPropertiesSet() {
-		Assert.state(this.serverContainer != null, "javax.websocket.server.ServerContainer not available");
+		Assert.state(getServerContainer() != null, "javax.websocket.server.ServerContainer not available");
+		registerEndpoints();
+	}
 
-		List<Class<?>> allClasses = new ArrayList<Class<?>>(this.annotatedEndpointClasses);
+	/**
-		allClasses.addAll(this.annotatedEndpointBeanTypes);
+	 * Actually register the endpoints. Called by {@link #afterPropertiesSet()}.
-
+	 * @since 4.1
-		for (Class<?> clazz : allClasses) {
+	 */
-			try {
+	protected void registerEndpoints() {
-				logger.info("Registering @ServerEndpoint type " + clazz);
+		Set<Class<?>> endpointClasses = new LinkedHashSet<Class<?>>();
-				this.serverContainer.addEndpoint(clazz);
+		if (this.annotatedEndpointClasses != null) {
-			}
+			endpointClasses.addAll(this.annotatedEndpointClasses);
-			catch (DeploymentException e) {
+		}
-				throw new IllegalStateException("Failed to register @ServerEndpoint type " + clazz, e);
+		if (this.annotatedEndpointBeanTypes != null) {
-			}
+			endpointClasses.addAll(this.annotatedEndpointBeanTypes);
+		}
+		for (Class<?> endpointClass : endpointClasses) {
+			registerEndpoint(endpointClass);
 		}
 	}
 
-	@Override
+	private void registerEndpoint(Class<?> endpointClass) {
-	public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
+		try {
-		if (bean instanceof ServerEndpointConfig) {
+			if (logger.isInfoEnabled()) {
-			ServerEndpointConfig sec = (ServerEndpointConfig) bean;
+				logger.info("Registering @ServerEndpoint type: " + endpointClass);
-			try {
-				if (logger.isInfoEnabled()) {
-					logger.info("Registering bean '" + beanName +
-							"' as javax.websocket.Endpoint under path " + sec.getPath());
-				}
-				getServerContainer().addEndpoint(sec);
-			}
-			catch (DeploymentException e) {
-				throw new IllegalStateException("Failed to deploy Endpoint bean " + bean, e);
 			}
+			getServerContainer().addEndpoint(endpointClass);
 		}
+		catch (DeploymentException ex) {
+			throw new IllegalStateException("Failed to register @ServerEndpoint type " + endpointClass, ex);
+		}
+	}
+
+
+	@Override
+	public Object postProcessBeforeInitialization(Object bean, String beanName) {
 		return bean;
 	}
 
 	@Override
-	public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
+	public Object postProcessAfterInitialization(Object bean, String beanName) {
+		if (bean instanceof ServerEndpointConfig) {
+			ServerEndpointConfig endpointConfig = (ServerEndpointConfig) bean;
+			try {
+				if (logger.isInfoEnabled()) {
+					logger.info("Registering bean '" + beanName +
+							"' as javax.websocket.Endpoint under path " + endpointConfig.getPath());
+				}
+				getServerContainer().addEndpoint(endpointConfig);
+			}
+			catch (DeploymentException ex) {
+				throw new IllegalStateException("Failed to deploy Endpoint bean with name '" + bean + "'", ex);
+			}
+		}
 		return bean;
 	}
 

spring-websocket/src/test/java/org/springframework/web/socket/server/standard/ServerEndpointExporterTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * Copyright 2002-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.
@@ -16,6 +16,7 @@
 
 package org.springframework.web.socket.server.standard;
 
+import javax.servlet.ServletContext;
 import javax.websocket.Endpoint;
 import javax.websocket.EndpointConfig;
 import javax.websocket.Session;
@@ -24,6 +25,7 @@ import javax.websocket.server.ServerEndpoint;
 
 import org.junit.Before;
 import org.junit.Test;
+
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.mock.web.test.MockServletContext;
@@ -35,37 +37,59 @@ import static org.mockito.Mockito.*;
  * Test fixture for {@link ServerEndpointExporter}.
  *
  * @author Rossen Stoyanchev
+ * @author Juergen Hoeller
  */
 public class ServerEndpointExporterTests {
 
 	private ServerContainer serverContainer;
 
-	private ServerEndpointExporter exporter;
+	private ServletContext servletContext;
 
 	private AnnotationConfigWebApplicationContext webAppContext;
 
+	private ServerEndpointExporter exporter;
+
 
 	@Before
 	public void setup() {
 		this.serverContainer = mock(ServerContainer.class);
 
-		MockServletContext servletContext = new MockServletContext();
+		this.servletContext = new MockServletContext();
-		servletContext.setAttribute("javax.websocket.server.ServerContainer", this.serverContainer);
+		this.servletContext.setAttribute("javax.websocket.server.ServerContainer", this.serverContainer);
 
 		this.webAppContext = new AnnotationConfigWebApplicationContext();
 		this.webAppContext.register(Config.class);
-		this.webAppContext.setServletContext(servletContext);
+		this.webAppContext.setServletContext(this.servletContext);
 		this.webAppContext.refresh();
 
 		this.exporter = new ServerEndpointExporter();
-		this.exporter.setApplicationContext(this.webAppContext);
 	}
 
 
 	@Test
-	public void addAnnotatedEndpointBean() throws Exception {
+	public void addAnnotatedEndpointBeans() throws Exception {
-
 		this.exporter.setAnnotatedEndpointClasses(AnnotatedDummyEndpoint.class);
+		this.exporter.setApplicationContext(this.webAppContext);
+		this.exporter.afterPropertiesSet();
+
+		verify(this.serverContainer).addEndpoint(AnnotatedDummyEndpoint.class);
+		verify(this.serverContainer).addEndpoint(AnnotatedDummyEndpointBean.class);
+	}
+
+	@Test
+	public void addAnnotatedEndpointBeansWithServletContextOnly() throws Exception {
+		this.exporter.setAnnotatedEndpointClasses(AnnotatedDummyEndpoint.class, AnnotatedDummyEndpointBean.class);
+		this.exporter.setServletContext(this.servletContext);
+		this.exporter.afterPropertiesSet();
+
+		verify(this.serverContainer).addEndpoint(AnnotatedDummyEndpoint.class);
+		verify(this.serverContainer).addEndpoint(AnnotatedDummyEndpointBean.class);
+	}
+
+	@Test
+	public void addAnnotatedEndpointBeansWithServerContainerOnly() throws Exception {
+		this.exporter.setAnnotatedEndpointClasses(AnnotatedDummyEndpoint.class, AnnotatedDummyEndpointBean.class);
+		this.exporter.setServerContainer(this.serverContainer);
 		this.exporter.afterPropertiesSet();
 
 		verify(this.serverContainer).addEndpoint(AnnotatedDummyEndpoint.class);
@@ -74,10 +98,31 @@ public class ServerEndpointExporterTests {
 
 	@Test
 	public void addServerEndpointConfigBean() throws Exception {
+		this.exporter.setApplicationContext(this.webAppContext);
+		this.exporter.afterPropertiesSet();
 
 		ServerEndpointRegistration endpointRegistration = new ServerEndpointRegistration("/dummy", new DummyEndpoint());
 		this.exporter.postProcessAfterInitialization(endpointRegistration, "dummyEndpoint");
+		verify(this.serverContainer).addEndpoint(endpointRegistration);
+	}
 
+	@Test
+	public void addServerEndpointConfigBeanWithServletContextOnly() throws Exception {
+		this.exporter.setServletContext(this.servletContext);
+		this.exporter.afterPropertiesSet();
+
+		ServerEndpointRegistration endpointRegistration = new ServerEndpointRegistration("/dummy", new DummyEndpoint());
+		this.exporter.postProcessAfterInitialization(endpointRegistration, "dummyEndpoint");
+		verify(this.serverContainer).addEndpoint(endpointRegistration);
+	}
+
+	@Test
+	public void addServerEndpointConfigBeanWithServerContainerOnly() throws Exception {
+		this.exporter.setServerContainer(this.serverContainer);
+		this.exporter.afterPropertiesSet();
+
+		ServerEndpointRegistration endpointRegistration = new ServerEndpointRegistration("/dummy", new DummyEndpoint());
+		this.exporter.postProcessAfterInitialization(endpointRegistration, "dummyEndpoint");
 		verify(this.serverContainer).addEndpoint(endpointRegistration);
 	}
 
@@ -89,14 +134,17 @@ public class ServerEndpointExporterTests {
 		}
 	}
 
+
 	@ServerEndpoint("/path")
 	private static class AnnotatedDummyEndpoint {
 	}
 
+
 	@ServerEndpoint("/path")
 	private static class AnnotatedDummyEndpointBean {
 	}
 
+
 	@Configuration
 	static class Config {