INT-3611: Support WebSocketHandlerDecoratorFactory

JIRA: https://jira.spring.io/browse/INT-3611

spring-integration-websocket/src/main/java/org/springframework/integration/websocket/ServerWebSocketContainer.java

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 package org.springframework.integration.websocket;
 
 import java.util.Arrays;
@@ -20,10 +21,12 @@ import java.util.Arrays;
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.util.Assert;
 import org.springframework.util.ObjectUtils;
+import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.config.annotation.SockJsServiceRegistration;
 import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
 import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistration;
 import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
+import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;
 import org.springframework.web.socket.server.HandshakeHandler;
 import org.springframework.web.socket.server.HandshakeInterceptor;
 import org.springframework.web.socket.sockjs.frame.SockJsMessageCodec;
@@ -47,9 +50,11 @@ public class ServerWebSocketContainer extends IntegrationWebSocketContainer impl
 
 	private final String[] paths;
 
-	private volatile HandshakeHandler handshakeHandler;
+	private HandshakeHandler handshakeHandler;
 
-	private volatile HandshakeInterceptor[] interceptors;
+	private HandshakeInterceptor[] interceptors;
+
+	private WebSocketHandlerDecoratorFactory[] decoratorFactories;
 
 	private SockJsServiceOptions sockJsServiceOptions;
 
@@ -62,12 +67,30 @@ public class ServerWebSocketContainer extends IntegrationWebSocketContainer impl
 		return this;
 	}
 
-	public ServerWebSocketContainer setInterceptors(HandshakeInterceptor[] interceptors) {
+	public ServerWebSocketContainer setInterceptors(HandshakeInterceptor... interceptors) {
+		Assert.notNull(interceptors, "'interceptors' must not be null");
+		Assert.noNullElements(interceptors, "'interceptors' must not contain null elements");
 		this.interceptors = Arrays.copyOf(interceptors, interceptors.length);
 		return this;
 	}
 
 
+	/**
+	 * Configure one or more factories to decorate the handler used to process
+	 * WebSocket messages. This may be useful in some advanced use cases, for
+	 * example to allow Spring Security to forcibly close the WebSocket session
+	 * when the corresponding HTTP session expires.
+	 * @param factories the WebSocketHandlerDecoratorFactory array to use
+	 * @return the current ServerWebSocketContainer
+	 * @since 4.2
+	 */
+	public ServerWebSocketContainer setDecoratorFactories(WebSocketHandlerDecoratorFactory... factories) {
+		Assert.notNull(factories, "'factories' must not be null");
+		Assert.noNullElements(factories, "'factories' must not contain null elements");
+		this.decoratorFactories = Arrays.copyOf(factories, factories.length);
+		return this;
+	}
+
 	public ServerWebSocketContainer withSockJs(SockJsServiceOptions... sockJsServiceOptions) {
 		if (ObjectUtils.isEmpty(sockJsServiceOptions)) {
 			this.sockJsServiceOptions = new SockJsServiceOptions();
@@ -85,9 +108,18 @@ public class ServerWebSocketContainer extends IntegrationWebSocketContainer impl
 
 	@Override
 	public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
-		WebSocketHandlerRegistration registration = registry.addHandler(this.webSocketHandler, this.paths)
+		WebSocketHandler webSocketHandler = this.webSocketHandler;
+
+		if (this.decoratorFactories != null) {
+			for (WebSocketHandlerDecoratorFactory factory : this.decoratorFactories) {
+				webSocketHandler = factory.decorate(webSocketHandler);
+			}
+		}
+
+		WebSocketHandlerRegistration registration = registry.addHandler(webSocketHandler, this.paths)
 				.setHandshakeHandler(this.handshakeHandler)
 				.addInterceptors(this.interceptors);
+
 		if (this.sockJsServiceOptions != null) {
 			SockJsServiceRegistration sockJsServiceRegistration = registration.withSockJS();
 			if (this.sockJsServiceOptions.webSocketEnabled != null) {

spring-integration-websocket/src/main/java/org/springframework/integration/websocket/config/ServerWebSocketContainerParser.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2014 the original author or authors.
+ * Copyright 2014-2015 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.
@@ -62,6 +62,14 @@ public class ServerWebSocketContainerParser extends AbstractSingleBeanDefinition
 		}
 		builder.addPropertyValue("interceptors", handshakeInterceptorList);
 
+		String decoratorFactories = element.getAttribute("decorator-factories");
+		List<BeanReference> decoratorFactoryList = new ManagedList<BeanReference>();
+		ids = StringUtils.commaDelimitedListToStringArray(decoratorFactories);
+		for (String id : ids) {
+			decoratorFactoryList.add(new RuntimeBeanReference(id));
+		}
+		builder.addPropertyValue("decoratorFactories", decoratorFactoryList);
+
 		Element sockjs = DomUtils.getChildElementByTagName(element, "sockjs");
 
 		if (sockjs != null) {

spring-integration-websocket/src/main/resources/org/springframework/integration/websocket/config/spring-integration-websocket-4.2.xsd

@@ -280,6 +280,17 @@
 					</xsd:documentation>
 				</xsd:annotation>
 			</xsd:attribute>
+			<xsd:attribute name="decorator-factories" type="xsd:string">
+				<xsd:annotation>
+					<xsd:documentation>
+						List of 'org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory' bean references.
+						Configure one or more factories to decorate the handler used to process WebSocket
+						messages. This may be useful for some advanced use cases, for example to allow
+						Spring Security to forcibly close the WebSocket session when the corresponding
+						HTTP session expires.
+					</xsd:documentation>
+				</xsd:annotation>
+			</xsd:attribute>
 			<xsd:attribute name="send-time-limit" default="10000">
 				<xsd:annotation>
 					<xsd:documentation>

spring-integration-websocket/src/test/java/org/springframework/integration/websocket/client/StompIntegrationTests.java

@@ -105,8 +105,6 @@ import org.springframework.web.socket.server.support.DefaultHandshakeHandler;
 @DirtiesContext
 public class StompIntegrationTests {
 
-	private static final SpelExpressionParser PARSER = new SpelExpressionParser();
-
 	@Value("#{server.serverContext}")
 	private ApplicationContext serverContext;
 

spring-integration-websocket/src/test/java/org/springframework/integration/websocket/config/WebSocketParserTests-context.xml

@@ -18,16 +18,17 @@
 									send-buffer-size-limit="100000"
 									send-time-limit="100"
 									handshake-handler="handshakeHandler"
-									handshake-interceptors="handshakeInterceptor">
+									handshake-interceptors="handshakeInterceptor"
+									decorator-factories="decoratorFactory">
 		<int-websocket:sockjs client-library-url="https://foo.sock.js"
-				disconnect-delay="4000"
-				heartbeat-time="30000"
-				message-cache-size="10000"
-				session-cookie-needed="false"
-				stream-bytes-limit="2000"
-				websocket-enabled="false"
-				scheduler="taskScheduler"
-				message-codec="sockJsMessageCodec"/>
+							  disconnect-delay="4000"
+							  heartbeat-time="30000"
+							  message-cache-size="10000"
+							  session-cookie-needed="false"
+							  stream-bytes-limit="2000"
+							  websocket-enabled="false"
+							  scheduler="taskScheduler"
+							  message-codec="sockJsMessageCodec"/>
 	</int-websocket:server-container>
 
 	<bean id="handshakeHandler" class="org.mockito.Mockito" factory-method="mock">
@@ -42,13 +43,16 @@
 		<constructor-arg value="org.springframework.web.socket.sockjs.frame.SockJsMessageCodec"/>
 	</bean>
 
+	<bean id="decoratorFactory"
+		  class="org.springframework.integration.websocket.config.WebSocketParserTests.TestWebSocketHandlerDecoratorFactory" />
+
 	<websocket:message-broker application-destination-prefix="/app">
 		<websocket:stomp-endpoint path="/foo"/>
 		<websocket:simple-broker prefix="/topic"/>
 	</websocket:message-broker>
 
 	<int-websocket:inbound-channel-adapter id="defaultInboundAdapter" container="serverWebSocketContainer"
-			use-broker="true"/>
+										   use-broker="true"/>
 
 	<bean id="webSocketClient" class="org.mockito.Mockito" factory-method="mock">
 		<constructor-arg value="org.springframework.web.socket.client.WebSocketClient"/>
@@ -93,11 +97,11 @@
 	<int-websocket:outbound-channel-adapter id="defaultOutboundAdapter" container="serverWebSocketContainer"/>
 
 	<int-websocket:outbound-channel-adapter id="customOutboundAdapter" container="clientWebSocketContainer"
-										   default-protocol-handler="stompSubProtocolHandler"
-										   protocol-handlers="passThruSubProtocolHandler"
-										   message-converters="simpleMessageConverter,mapMessageConverter"
-										   merge-with-default-converters="true"
-										   channel="clientOutboundChannel"/>
+											default-protocol-handler="stompSubProtocolHandler"
+											protocol-handlers="passThruSubProtocolHandler"
+											message-converters="simpleMessageConverter,mapMessageConverter"
+											merge-with-default-converters="true"
+											channel="clientOutboundChannel"/>
 
 	<int:channel id="clientOutboundChannel"/>
 

spring-integration-websocket/src/test/java/org/springframework/integration/websocket/config/WebSocketParserTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2014 the original author or authors.
+ * Copyright 2014-2015 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.
@@ -19,6 +19,7 @@ package org.springframework.integration.websocket.config;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThat;
@@ -54,8 +55,10 @@ import org.springframework.test.annotation.DirtiesContext;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 import org.springframework.web.servlet.HandlerMapping;
+import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.WebSocketHttpHeaders;
 import org.springframework.web.socket.client.WebSocketClient;
+import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;
 import org.springframework.web.socket.messaging.StompSubProtocolHandler;
 import org.springframework.web.socket.server.HandshakeHandler;
 import org.springframework.web.socket.server.HandshakeInterceptor;
@@ -138,7 +141,11 @@ public class WebSocketParserTests {
 	@Qualifier("customOutboundAdapter.handler")
 	private WebSocketOutboundMessageHandler customOutboundAdapter;
 
+	@Autowired
+	private WebSocketHandlerDecoratorFactory decoratorFactory;
+
 	@Test
+	@SuppressWarnings("unckecked")
 	public void testDefaultInboundChannelAdapterAndServerContainer() {
 		Map<?, ?> urlMap = TestUtils.getPropertyValue(this.handlerMapping, "urlMap", Map.class);
 		assertEquals(1, urlMap.size());
@@ -152,11 +159,19 @@ public class WebSocketParserTests {
 				TestUtils.getPropertyValue(this.serverWebSocketContainer, "handshakeHandler"));
 		HandshakeInterceptor[] interceptors =
 				TestUtils.getPropertyValue(this.serverWebSocketContainer, "interceptors", HandshakeInterceptor[].class);
+		assertNotNull(interceptors);
 		assertEquals(1, interceptors.length);
 		assertSame(this.handshakeInterceptor, interceptors[0]);
 		assertEquals(100, TestUtils.getPropertyValue(this.serverWebSocketContainer, "sendTimeLimit"));
 		assertEquals(100000, TestUtils.getPropertyValue(this.serverWebSocketContainer, "sendBufferSizeLimit"));
 
+		WebSocketHandlerDecoratorFactory[] decoratorFactories =
+				TestUtils.getPropertyValue(this.serverWebSocketContainer, "decoratorFactories",
+						WebSocketHandlerDecoratorFactory[].class);
+		assertNotNull(decoratorFactories);
+		assertEquals(1, decoratorFactories.length);
+		assertSame(this.decoratorFactory, decoratorFactories[0]);
+
 		TransportHandlingSockJsService sockJsService =
 				TestUtils.getPropertyValue(mappedHandler, "sockJsService", TransportHandlingSockJsService.class);
 		assertSame(this.taskScheduler, sockJsService.getTaskScheduler());
@@ -298,4 +313,13 @@ public class WebSocketParserTests {
 		assertTrue(TestUtils.getPropertyValue(this.customOutboundAdapter, "client", Boolean.class));
 	}
 
+	private static class TestWebSocketHandlerDecoratorFactory implements WebSocketHandlerDecoratorFactory {
+
+		@Override
+		public WebSocketHandler decorate(WebSocketHandler handler) {
+			return handler;
+		}
+
+	}
+
 }

spring-integration-websocket/src/test/java/org/springframework/integration/websocket/server/WebSocketServerTests.java

@@ -37,6 +37,11 @@ import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationEvent;
+import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.context.ApplicationEventPublisherAware;
+import org.springframework.context.ApplicationListener;
+import org.springframework.context.PayloadApplicationEvent;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.expression.spel.standard.SpelExpressionParser;
@@ -46,6 +51,7 @@ import org.springframework.integration.channel.DirectChannel;
 import org.springframework.integration.channel.QueueChannel;
 import org.springframework.integration.config.EnableIntegration;
 import org.springframework.integration.core.MessageProducer;
+import org.springframework.integration.event.inbound.ApplicationEventListeningMessageProducer;
 import org.springframework.integration.transformer.ExpressionEvaluatingTransformer;
 import org.springframework.integration.websocket.ClientWebSocketContainer;
 import org.springframework.integration.websocket.IntegrationWebSocketContainer;
@@ -68,11 +74,16 @@ import org.springframework.test.annotation.DirtiesContext;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 import org.springframework.util.MultiValueMap;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.WebSocketMessage;
+import org.springframework.web.socket.WebSocketSession;
 import org.springframework.web.socket.client.WebSocketClient;
 import org.springframework.web.socket.client.standard.StandardWebSocketClient;
 import org.springframework.web.socket.config.annotation.AbstractWebSocketMessageBrokerConfigurer;
 import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
 import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
+import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
+import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;
 import org.springframework.web.socket.messaging.StompSubProtocolHandler;
 import org.springframework.web.socket.messaging.SubProtocolHandler;
 import org.springframework.web.socket.sockjs.client.SockJsClient;
@@ -101,6 +112,9 @@ public class WebSocketServerTests {
 	@Value("#{server.serverContext.getBean('simpleBrokerMessageHandler')}")
 	private SimpleBrokerMessageHandler brokerHandler;
 
+	@Value("#{server.serverContext.getBean('webSocketEvents')}")
+	private PollableChannel webSocketEvents;
+
 	@Test
 	public void testWebSocketOutboundMessageHandler() throws Exception {
 		StompHeaderAccessor headers = StompHeaderAccessor.create(StompCommand.SUBSCRIBE);
@@ -133,6 +147,10 @@ public class WebSocketServerTests {
 		List<String> subscription = subscriptions.values().iterator().next();
 		assertEquals(1, subscription.size());
 		assertEquals("subs1", subscription.get(0));
+
+		Message<?> event = this.webSocketEvents.receive(10000);
+		assertNotNull(event);
+		assertThat(event.getPayload(), instanceOf(WebSocketSession.class));
 	}
 
 	@Test
@@ -224,9 +242,16 @@ public class WebSocketServerTests {
 					.enableSimpleBroker("/queue/", "/topic/");
 		}
 
+		@Bean
+		public WebSocketHandlerDecoratorFactory testWebSocketHandlerDecoratorFactory() {
+			return new TestWebSocketHandlerDecoratorFactory();
+		}
+
 		@Bean
 		public ServerWebSocketContainer serverWebSocketContainer() {
-			return new ServerWebSocketContainer("/ws").withSockJs();
+			return new ServerWebSocketContainer("/ws")
+					.setDecoratorFactories(testWebSocketHandlerDecoratorFactory())
+					.withSockJs();
 		}
 
 		@Bean
@@ -267,6 +292,57 @@ public class WebSocketServerTests {
 					new SubProtocolHandlerRegistry(stompSubProtocolHandler()));
 		}
 
+		@Bean
+		public PollableChannel webSocketEvents() {
+			return new QueueChannel();
+		}
+
+		@Bean
+		@SuppressWarnings("unchecked")
+		public ApplicationListener<ApplicationEvent> webSocketEventListener() {
+			ApplicationEventListeningMessageProducer producer = new ApplicationEventListeningMessageProducer();
+			producer.setEventTypes(PayloadApplicationEvent.class);
+			producer.setExpressionPayload(new SpelExpressionParser().parseExpression("payload"));
+			producer.setOutputChannel(webSocketEvents());
+			return producer;
+		}
+
+	}
+
+	private static class TestWebSocketHandlerDecoratorFactory
+			implements WebSocketHandlerDecoratorFactory, ApplicationEventPublisherAware {
+
+		private ApplicationEventPublisher applicationEventPublisher;
+
+		@Override
+		public WebSocketHandler decorate(WebSocketHandler handler) {
+			return new TestWebSocketHandler(handler);
+		}
+
+		@Override
+		public void setApplicationEventPublisher(ApplicationEventPublisher applicationEventPublisher) {
+			this.applicationEventPublisher = applicationEventPublisher;
+		}
+
+		private class TestWebSocketHandler extends WebSocketHandlerDecorator {
+
+			public TestWebSocketHandler(WebSocketHandler delegate) {
+				super(delegate);
+			}
+
+			@Override
+			public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
+				super.handleMessage(session, message);
+			}
+
+			@Override
+			public void afterConnectionEstablished(WebSocketSession session) throws Exception {
+				super.afterConnectionEstablished(session);
+				applicationEventPublisher.publishEvent(session);
+			}
+
+		}
+
 	}
 
 }