Fix issue with StompSubProtocolHandler initialization

This change ensures that StompSubProtocolHandler is injected with an
ApplicationEventPublisher for both the Java and XML config.

Issue: SPR-11825

spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistry.java

@@ -21,6 +21,9 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 
+import org.springframework.beans.BeansException;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextAware;
 import org.springframework.messaging.simp.user.UserSessionRegistry;
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.util.Assert;
@@ -38,9 +41,10 @@ import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
  * {@link SimpleUrlHandlerMapping} for use in Spring MVC.
  *
  * @author Rossen Stoyanchev
+ * @author Artem Bilan
  * @since 4.0
  */
-public class WebMvcStompEndpointRegistry implements StompEndpointRegistry {
+public class WebMvcStompEndpointRegistry implements StompEndpointRegistry, ApplicationContextAware {
 
 	private final WebSocketHandler webSocketHandler;
 
@@ -84,6 +88,11 @@ public class WebMvcStompEndpointRegistry implements StompEndpointRegistry {
 		this.sockJsScheduler = defaultSockJsTaskScheduler;
 	}
 
+	@Override
+	public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
+		this.stompHandler.setApplicationEventPublisher(applicationContext);
+	}
+
 	private static SubProtocolWebSocketHandler unwrapSubProtocolWebSocketHandler(WebSocketHandler wsHandler) {
 		WebSocketHandler actual = WebSocketHandlerDecorator.unwrap(wsHandler);
 		Assert.isInstanceOf(SubProtocolWebSocketHandler.class, actual, "No SubProtocolWebSocketHandler in " + wsHandler);

spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupport.java

@@ -16,6 +16,8 @@
 
 package org.springframework.web.socket.config.annotation;
 
+import java.util.Collections;
+
 import org.springframework.beans.factory.config.CustomScopeConfigurer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.messaging.simp.SimpSessionScope;
@@ -26,8 +28,6 @@ import org.springframework.web.servlet.HandlerMapping;
 import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
 
-import java.util.Collections;
-
 /**
  * Extends {@link AbstractMessageBrokerConfiguration} and adds configuration for
  * receiving and responding to STOMP messages from WebSocket clients.
@@ -37,6 +37,7 @@ import java.util.Collections;
  * also be extended directly.
  *
  * @author Rossen Stoyanchev
+ * @author Artem Bilan
  * @since 4.0
  */
 public abstract class WebSocketMessageBrokerConfigurationSupport extends AbstractMessageBrokerConfiguration {
@@ -58,6 +59,8 @@ public abstract class WebSocketMessageBrokerConfigurationSupport extends Abstrac
 		WebMvcStompEndpointRegistry registry = new WebMvcStompEndpointRegistry(
 				webSocketHandler, transportRegistration, sessionRegistry, taskScheduler);
 
+		registry.setApplicationContext(getApplicationContext());
+
 		registerStompEndpoints(registry);
 
 		return registry.getHandlerMapping();

spring-websocket/src/main/java/org/springframework/web/socket/messaging/SubProtocolWebSocketHandler.java

@@ -28,8 +28,6 @@ import java.util.concurrent.ConcurrentHashMap;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
-import org.springframework.context.ApplicationEventPublisher;
-import org.springframework.context.ApplicationEventPublisherAware;
 import org.springframework.context.SmartLifecycle;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
@@ -60,10 +58,11 @@ import org.springframework.web.socket.handler.SessionLimitExceededException;
  *
  * @author Rossen Stoyanchev
  * @author Andy Wilkinson
+ * @author Artem Bilan
  * @since 4.0
  */
 public class SubProtocolWebSocketHandler implements WebSocketHandler,
-		SubProtocolCapable, MessageHandler, SmartLifecycle, ApplicationEventPublisherAware {
+		SubProtocolCapable, MessageHandler, SmartLifecycle {
 
 	private final Log logger = LogFactory.getLog(SubProtocolWebSocketHandler.class);
 
@@ -82,12 +81,10 @@ public class SubProtocolWebSocketHandler implements WebSocketHandler,
 
 	private int sendBufferSizeLimit = 512 * 1024;
 
-	private Object lifecycleMonitor = new Object();
+	private final Object lifecycleMonitor = new Object();
 
 	private volatile boolean running = false;
 
-	private ApplicationEventPublisher eventPublisher;
-
 
 	public SubProtocolWebSocketHandler(MessageChannel clientInboundChannel, SubscribableChannel clientOutboundChannel) {
 		Assert.notNull(clientInboundChannel, "ClientInboundChannel must not be null");
@@ -132,10 +129,6 @@ public class SubProtocolWebSocketHandler implements WebSocketHandler,
 						+ " to protocol '" + protocol + "', it is already mapped to handler " + replaced);
 			}
 		}
-
-		if (handler instanceof ApplicationEventPublisherAware) {
-			((ApplicationEventPublisherAware) handler).setApplicationEventPublisher(this.eventPublisher);
-		}
 	}
 
 	/**
@@ -188,11 +181,6 @@ public class SubProtocolWebSocketHandler implements WebSocketHandler,
 		return sendBufferSizeLimit;
 	}
 
-	@Override
-	public void setApplicationEventPublisher(ApplicationEventPublisher eventPublisher) {
-		this.eventPublisher = eventPublisher;
-	}
-
 	@Override
 	public boolean isAutoStartup() {
 		return true;

spring-websocket/src/test/java/org/springframework/web/socket/config/MessageBrokerBeanDefinitionParserTests.java

@@ -16,6 +16,8 @@
 
 package org.springframework.web.socket.config;
 
+import static org.junit.Assert.*;
+
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -39,12 +41,12 @@ import org.springframework.messaging.converter.MessageConverter;
 import org.springframework.messaging.converter.StringMessageConverter;
 import org.springframework.messaging.simp.SimpMessagingTemplate;
 import org.springframework.messaging.simp.annotation.support.SimpAnnotationMethodMessageHandler;
-import org.springframework.messaging.simp.user.DefaultUserDestinationResolver;
 import org.springframework.messaging.simp.broker.SimpleBrokerMessageHandler;
+import org.springframework.messaging.simp.stomp.StompBrokerRelayMessageHandler;
+import org.springframework.messaging.simp.user.DefaultUserDestinationResolver;
 import org.springframework.messaging.simp.user.UserDestinationMessageHandler;
 import org.springframework.messaging.simp.user.UserDestinationResolver;
 import org.springframework.messaging.simp.user.UserSessionRegistry;
-import org.springframework.messaging.simp.stomp.StompBrokerRelayMessageHandler;
 import org.springframework.messaging.support.AbstractSubscribableChannel;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
@@ -64,13 +66,12 @@ import org.springframework.web.socket.sockjs.transport.TransportType;
 import org.springframework.web.socket.sockjs.transport.handler.DefaultSockJsService;
 import org.springframework.web.socket.sockjs.transport.handler.WebSocketTransportHandler;
 
-import static org.junit.Assert.*;
-
 /**
  * Test fixture for MessageBrokerBeanDefinitionParser.
  * See test configuration files websocket-config-broker-*.xml.
  *
  * @author Brian Clozel
+ * @author Artem Bilan
  */
 public class MessageBrokerBeanDefinitionParserTests {
 
@@ -118,6 +119,8 @@ public class MessageBrokerBeanDefinitionParserTests {
 		assertNotNull(stompHandler);
 		assertEquals(128 * 1024, stompHandler.getMessageSizeLimit());
 
+		assertNotNull(new DirectFieldAccessor(stompHandler).getPropertyValue("eventPublisher"));
+
 		httpRequestHandler = (HttpRequestHandler) suhm.getUrlMap().get("/test/**");
 		assertNotNull(httpRequestHandler);
 		assertThat(httpRequestHandler, Matchers.instanceOf(SockJsHttpRequestHandler.class));
@@ -365,4 +368,4 @@ public class MessageBrokerBeanDefinitionParserTests {
 				((WebSocketHandlerDecorator) handler).getLastHandler() : handler;
 	}
 
-}
+}