Add WebSocket scope

This change adds support for a custom "websocket" scope.

WebSocket-scoped beans may be injected into controllers with message
handling methods as well as channel interceptor registered on the
"inboundClientChannel".

Issue: SPR-11305

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

@@ -17,8 +17,12 @@
 package org.springframework.web.socket.config;
 
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 
+import org.springframework.beans.factory.config.CustomScopeConfigurer;
+import org.springframework.messaging.simp.SimpSessionScope;
 import org.w3c.dom.Element;
 
 import org.springframework.beans.MutablePropertyValues;
@@ -167,6 +171,11 @@ class MessageBrokerBeanDefinitionParser implements BeanDefinitionParser {
 		registerUserDestinationMessageHandler(clientInChannel, clientOutChannel, brokerChannel,
 				userDestinationResolver, parserCxt, source);
 
+		Map<String, Object> scopeMap = Collections.<String, Object>singletonMap("websocket", new SimpSessionScope());
+		RootBeanDefinition scopeConfigurerDef = new RootBeanDefinition(CustomScopeConfigurer.class);
+		scopeConfigurerDef.getPropertyValues().add("scopes", scopeMap);
+		registerBeanDefByName("webSocketScopeConfigurer", scopeConfigurerDef, parserCxt, source);
+
 		parserCxt.popAndRegisterContainingComponent();
 
 		return null;

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

@@ -58,6 +58,6 @@ import org.springframework.context.annotation.Import;
 @Retention(RetentionPolicy.RUNTIME)
 @Target(ElementType.TYPE)
 @Documented
-@Import(DelegatingWebSocketConfiguration.class)
+@Import({DelegatingWebSocketConfiguration.class, WebSocketScopeConfiguration.class})
 public @interface EnableWebSocket {
 }

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

@@ -16,7 +16,9 @@
 
 package org.springframework.web.socket.config.annotation;
 
+import org.springframework.beans.factory.config.CustomScopeConfigurer;
 import org.springframework.context.annotation.Bean;
+import org.springframework.messaging.simp.SimpSessionScope;
 import org.springframework.messaging.simp.config.AbstractMessageBrokerConfiguration;
 import org.springframework.messaging.simp.user.UserSessionRegistry;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
@@ -24,6 +26,8 @@ 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.
@@ -75,6 +79,8 @@ public abstract class WebSocketMessageBrokerConfigurationSupport extends Abstrac
 	protected void configureWebSocketTransport(WebSocketTransportRegistration registry) {
 	}
 
+	protected abstract void registerStompEndpoints(StompEndpointRegistry registry);
+
 	/**
 	 * The default TaskScheduler to use if none is configured via
 	 * {@link SockJsServiceRegistration#setTaskScheduler(org.springframework.scheduling.TaskScheduler)}, i.e.
@@ -100,6 +106,11 @@ public abstract class WebSocketMessageBrokerConfigurationSupport extends Abstrac
 		return scheduler;
 	}
 
-	protected abstract void registerStompEndpoints(StompEndpointRegistry registry);
+	@Bean
+	public static CustomScopeConfigurer webSocketScopeConfigurer() {
+		CustomScopeConfigurer configurer = new CustomScopeConfigurer();
+		configurer.setScopes(Collections.<String, Object>singletonMap("websocket", new SimpSessionScope()));
+		return configurer;
+	}
 
 }

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

@@ -0,0 +1,39 @@
+/*
+ * 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.
+ * 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 org.springframework.web.socket.config.annotation;
+
+import org.springframework.beans.factory.config.CustomScopeConfigurer;
+import org.springframework.context.annotation.Bean;
+import org.springframework.messaging.simp.SimpSessionScope;
+
+import java.util.Collections;
+
+/**
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.1
+ */
+public class WebSocketScopeConfiguration {
+
+	@Bean
+	public CustomScopeConfigurer webSocketScopeConfigurer() {
+		CustomScopeConfigurer configurer = new CustomScopeConfigurer();
+		configurer.setScopes(Collections.<String, Object>singletonMap("websocket", new SimpSessionScope()));
+		return configurer;
+	}
+
+}

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

@@ -33,6 +33,8 @@ import org.springframework.context.ApplicationEventPublisher;
 import org.springframework.context.ApplicationEventPublisherAware;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.simp.SimpAttributes;
+import org.springframework.messaging.simp.SimpAttributesContextHolder;
 import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
 import org.springframework.messaging.simp.SimpMessageType;
 import org.springframework.messaging.simp.stomp.BufferingStompDecoder;
@@ -221,7 +223,13 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 					publishEvent(new SessionConnectEvent(this, message));
 				}
 
-				outputChannel.send(message);
+				try {
+					SimpAttributesContextHolder.setAttributesFromMessage(message);
+					outputChannel.send(message);
+				}
+				finally {
+					SimpAttributesContextHolder.resetAttributes();
+				}
 			}
 			catch (Throwable ex) {
 				logger.error("Terminating STOMP session due to failure to send message", ex);
@@ -420,22 +428,33 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 			this.userSessionRegistry.unregisterSessionId(userName, session.getId());
 		}
 
-		if (logger.isDebugEnabled()) {
-			logger.debug("WebSocket session ended, sending DISCONNECT message to broker");
+		if (this.eventPublisher != null) {
+			publishEvent(new SessionDisconnectEvent(this, session.getId(), closeStatus));
 		}
 
+		Message<?> message = createDisconnectMessage(session);
+		SimpAttributes simpAttributes = SimpAttributes.fromMessage(message);
+		try {
+			if (logger.isDebugEnabled()) {
+				logger.debug("WebSocket session ended, sending DISCONNECT message to broker");
+			}
+			SimpAttributesContextHolder.setAttributes(simpAttributes);
+			outputChannel.send(message);
+		}
+		finally {
+			SimpAttributesContextHolder.resetAttributes();
+			simpAttributes.sessionCompleted();
+		}
+	}
+
+	private Message<?> createDisconnectMessage(WebSocketSession session) {
 		StompHeaderAccessor headerAccessor = StompHeaderAccessor.create(StompCommand.DISCONNECT);
 		if (getHeaderInitializer() != null) {
 			getHeaderInitializer().initHeaders(headerAccessor);
 		}
 		headerAccessor.setSessionId(session.getId());
-		Message<?> message = MessageBuilder.createMessage(EMPTY_PAYLOAD, headerAccessor.getMessageHeaders());
-
-		if (this.eventPublisher != null) {
-			publishEvent(new SessionDisconnectEvent(this, session.getId(), closeStatus));
-		}
-
-		outputChannel.send(message);
+		headerAccessor.setSessionAttributes(session.getAttributes());
+		return MessageBuilder.createMessage(EMPTY_PAYLOAD, headerAccessor.getMessageHeaders());
 	}
 
 }