Reimplemented ServerEndpointExporter to avoid BeanPostProcessor role

Issue: SPR-12340
(cherry picked from commit 10328f1)

Author: jhoeller

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

@@ -19,6 +19,7 @@
 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;
@@ -27,7 +28,6 @@
 import javax.websocket.server.ServerEndpointConfig;
 
 import org.springframework.beans.factory.InitializingBean;
-import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.context.ApplicationContext;
 import org.springframework.util.Assert;
 import org.springframework.web.context.support.WebApplicationObjectSupport;
@@ -41,7 +41,7 @@
  *
  * <p>When this class is used, by declaring it in Spring configuration, it should be
  * possible to turn off a Servlet container's scan for WebSocket endpoints. This can be
- * done with the help of the {@code <absolute-ordering>} element in web.xml.
+ * done with the help of the {@code <absolute-ordering>} element in {@code web.xml}.
  *
  * @author Rossen Stoyanchev
  * @author Juergen Hoeller
@@ -50,19 +50,26 @@
  * @see SpringConfigurator
  * @see ServletServerContainerFactoryBean
  */
-public class ServerEndpointExporter extends WebApplicationObjectSupport implements BeanPostProcessor, InitializingBean {
-
-	private ServerContainer serverContainer;
+public class ServerEndpointExporter extends WebApplicationObjectSupport implements InitializingBean {
 
 	private List<Class<?>> annotatedEndpointClasses;
 
-	private Set<Class<?>> annotatedEndpointBeanTypes;
+	private ServerContainer serverContainer;
+
 
+	/**
+	 * Explicitly list annotated endpoint types that should be registered on startup. This
+	 * can be done if you wish to turn off a Servlet container's scan for endpoints, which
+	 * goes through all 3rd party jars in the, and rely on Spring configuration instead.
+	 * @param annotatedEndpointClasses {@link ServerEndpoint}-annotated types
+	 */
+	public void setAnnotatedEndpointClasses(Class<?>... annotatedEndpointClasses) {
+		this.annotatedEndpointClasses = Arrays.asList(annotatedEndpointClasses);
+	}
 
 	/**
 	 * 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;
@@ -75,33 +82,6 @@ protected ServerContainer getServerContainer() {
 		return this.serverContainer;
 	}
 
-	/**
-	 * Explicitly list annotated endpoint types that should be registered on startup. This
-	 * can be done if you wish to turn off a Servlet container's scan for endpoints, which
-	 * goes through all 3rd party jars in the, and rely on Spring configuration instead.
-	 * @param annotatedEndpointClasses {@link ServerEndpoint}-annotated types
- 	 */
-	public void setAnnotatedEndpointClasses(Class<?>... annotatedEndpointClasses) {
-		this.annotatedEndpointClasses = Arrays.asList(annotatedEndpointClasses);
-	}
-
-	@Override
-	protected void initApplicationContext(ApplicationContext context) {
-		// Initializes ServletContext given a WebApplicationContext
-		super.initApplicationContext(context);
-
-		// Retrieve beans which are annotated with @ServerEndpoint
-		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");
-			}
-			this.annotatedEndpointBeanTypes.add(beanType);
-		}
-	}
-
 	@Override
 	protected void initServletContext(ServletContext servletContext) {
 		if (this.serverContainer == null) {
@@ -110,64 +90,69 @@ protected void initServletContext(ServletContext servletContext) {
 		}
 	}
 
+	@Override
+	protected boolean isContextRequired() {
+		return false;
+	}
 
 	@Override
 	public void afterPropertiesSet() {
 		Assert.state(getServerContainer() != null, "javax.websocket.server.ServerContainer not available");
 		registerEndpoints();
 	}
 
+
 	/**
 	 * Actually register the endpoints. Called by {@link #afterPropertiesSet()}.
-	 * @since 4.1
 	 */
 	protected void registerEndpoints() {
 		Set<Class<?>> endpointClasses = new LinkedHashSet<Class<?>>();
 		if (this.annotatedEndpointClasses != null) {
 			endpointClasses.addAll(this.annotatedEndpointClasses);
 		}
-		if (this.annotatedEndpointBeanTypes != null) {
-			endpointClasses.addAll(this.annotatedEndpointBeanTypes);
+
+		ApplicationContext context = getApplicationContext();
+		if (context != null) {
+			String[] endpointBeanNames = context.getBeanNamesForAnnotation(ServerEndpoint.class);
+			for (String beanName : endpointBeanNames) {
+				endpointClasses.add(context.getType(beanName));
+			}
 		}
+
 		for (Class<?> endpointClass : endpointClasses) {
 			registerEndpoint(endpointClass);
 		}
+
+		if (context != null) {
+			Map<String, ServerEndpointConfig> endpointConfigMap = context.getBeansOfType(ServerEndpointConfig.class);
+			for (ServerEndpointConfig endpointConfig : endpointConfigMap.values()) {
+				registerEndpoint(endpointConfig);
+			}
+		}
 	}
 
 	private void registerEndpoint(Class<?> endpointClass) {
 		try {
 			if (logger.isInfoEnabled()) {
-				logger.info("Registering @ServerEndpoint type: " + endpointClass);
+				logger.info("Registering @ServerEndpoint class: " + endpointClass);
 			}
 			getServerContainer().addEndpoint(endpointClass);
 		}
 		catch (DeploymentException ex) {
-			throw new IllegalStateException("Failed to register @ServerEndpoint type " + endpointClass, ex);
+			throw new IllegalStateException("Failed to register @ServerEndpoint class: " + endpointClass, ex);
 		}
 	}
 
-
-	@Override
-	public Object postProcessBeforeInitialization(Object bean, String beanName) {
-		return bean;
-	}
-
-	@Override
-	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);
+	private void registerEndpoint(ServerEndpointConfig endpointConfig) {
+		try {
+			if (logger.isInfoEnabled()) {
+				logger.info("Registering ServerEndpointConfig: " + endpointConfig);
 			}
+			getServerContainer().addEndpoint(endpointConfig);
+		}
+		catch (DeploymentException ex) {
+			throw new IllegalStateException("Failed to register ServerEndpointConfig: " + endpointConfig, ex);
 		}
-		return bean;
 	}
 
 }

Diff for: spring-websocket/src/main/java/org/springframework/web/socket/server/standard/ServerEndpointRegistration.java

@@ -5,7 +5,7 @@
  * 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
+ *      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,
@@ -73,7 +73,7 @@ public class ServerEndpointRegistration extends ServerEndpointConfig.Configurato
 
 	/**
 	 * Create a new {@link ServerEndpointRegistration} instance from an
-	 * {@code javax.webscoket.Endpoint} class.
+	 * {@code javax.websocket.Endpoint} class.
 	 * @param path the endpoint path
 	 * @param endpointClass the endpoint class
 	 */
@@ -202,4 +202,9 @@ public List<Extension> getNegotiatedExtensions(List<Extension> installed, List<E
 		return super.getNegotiatedExtensions(installed, requested);
 	}
 
+
+	@Override
+	public String toString() {
+		return "ServerEndpointRegistration for path '" + getPath() + "': " + getEndpointClass();
+	}
 }

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

@@ -5,7 +5,7 @@
  * 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
+ *      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,
@@ -67,7 +67,7 @@ public void setup() {
 
 
 	@Test
-	public void addAnnotatedEndpointBeans() throws Exception {
+	public void addAnnotatedEndpointClasses() throws Exception {
 		this.exporter.setAnnotatedEndpointClasses(AnnotatedDummyEndpoint.class);
 		this.exporter.setApplicationContext(this.webAppContext);
 		this.exporter.afterPropertiesSet();
@@ -77,7 +77,7 @@ public void addAnnotatedEndpointBeans() throws Exception {
 	}
 
 	@Test
-	public void addAnnotatedEndpointBeansWithServletContextOnly() throws Exception {
+	public void addAnnotatedEndpointClassesWithServletContextOnly() throws Exception {
 		this.exporter.setAnnotatedEndpointClasses(AnnotatedDummyEndpoint.class, AnnotatedDummyEndpointBean.class);
 		this.exporter.setServletContext(this.servletContext);
 		this.exporter.afterPropertiesSet();
@@ -87,7 +87,7 @@ public void addAnnotatedEndpointBeansWithServletContextOnly() throws Exception {
 	}
 
 	@Test
-	public void addAnnotatedEndpointBeansWithServerContainerOnly() throws Exception {
+	public void addAnnotatedEndpointClassesWithExplicitServerContainerOnly() throws Exception {
 		this.exporter.setAnnotatedEndpointClasses(AnnotatedDummyEndpoint.class, AnnotatedDummyEndpointBean.class);
 		this.exporter.setServerContainer(this.serverContainer);
 		this.exporter.afterPropertiesSet();
@@ -98,31 +98,37 @@ public void addAnnotatedEndpointBeansWithServerContainerOnly() throws Exception
 
 	@Test
 	public void addServerEndpointConfigBean() throws Exception {
+		ServerEndpointRegistration endpointRegistration = new ServerEndpointRegistration("/dummy", new DummyEndpoint());
+		this.webAppContext.getBeanFactory().registerSingleton("dummyEndpoint", endpointRegistration);
+
 		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 {
+	public void addServerEndpointConfigBeanWithExplicitServletContext() throws Exception {
+		ServerEndpointRegistration endpointRegistration = new ServerEndpointRegistration("/dummy", new DummyEndpoint());
+		this.webAppContext.getBeanFactory().registerSingleton("dummyEndpoint", endpointRegistration);
+
 		this.exporter.setServletContext(this.servletContext);
+		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 addServerEndpointConfigBeanWithServerContainerOnly() throws Exception {
+	public void addServerEndpointConfigBeanWithExplicitServerContainer() throws Exception {
+		ServerEndpointRegistration endpointRegistration = new ServerEndpointRegistration("/dummy", new DummyEndpoint());
+		this.webAppContext.getBeanFactory().registerSingleton("dummyEndpoint", endpointRegistration);
+		this.servletContext.removeAttribute("javax.websocket.server.ServerContainer");
+
 		this.exporter.setServerContainer(this.serverContainer);
+		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);
 	}