Consolidate websocket/messaging code

Before this change spring-messaging contained a few WebSocket-related
classes including WebSocket sub-protocol support for STOMP as well
as @EnableWebSocketMessageBroker and related configuration classes.

After this change those classes are located in the spring-websocket
module under org.springframework.web.socket.messaging.

This means the following classes in application configuration must
have their packages updated:

org.springframework.web.socket.messaging.config.EnableWebSocketMessageBroker
org.springframework.web.socket.messaging.config.StompEndpointRegistry
org.springframework.web.socket.messaging.config.WebSocketMessageBrokerConfigurer

MessageBrokerConfigurer has been renamed to MessageBrokerRegistry and
is also located in the above package.

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

@@ -0,0 +1,250 @@
+/*
+ * Copyright 2002-2013 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.messaging;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+import java.security.Principal;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Set;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.simp.SimpMessageType;
+import org.springframework.messaging.simp.handler.UserSessionRegistry;
+import org.springframework.messaging.simp.stomp.*;
+import org.springframework.messaging.support.MessageBuilder;
+import org.springframework.util.Assert;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.TextMessage;
+import org.springframework.web.socket.WebSocketMessage;
+import org.springframework.web.socket.WebSocketSession;
+
+/**
+ * A {@link SubProtocolHandler} for STOMP that supports versions 1.0, 1.1, and 1.2 of the
+ * STOMP specification.
+ *
+ * @author Rossen Stoyanchev
+ * @author Andy Wilkinson
+ * @since 4.0
+ */
+public class StompSubProtocolHandler implements SubProtocolHandler {
+
+	/**
+	 * The name of the header set on the CONNECTED frame indicating the name of the user
+	 * authenticated on the WebSocket session.
+	 */
+	public static final String CONNECTED_USER_HEADER = "user-name";
+
+	private static final Log logger = LogFactory.getLog(StompSubProtocolHandler.class);
+
+
+	private final StompDecoder stompDecoder = new StompDecoder();
+
+	private final StompEncoder stompEncoder = new StompEncoder();
+
+	private UserSessionRegistry userSessionRegistry;
+
+
+	/**
+	 * Provide a registry with which to register active user session ids.
+	 *
+	 * @see {@link org.springframework.messaging.simp.handler.UserDestinationMessageHandler}
+	 */
+	public void setUserSessionRegistry(UserSessionRegistry registry) {
+		this.userSessionRegistry = registry;
+	}
+
+	/**
+	 * @return the configured UserSessionRegistry.
+	 */
+	public UserSessionRegistry getUserSessionRegistry() {
+		return this.userSessionRegistry;
+	}
+
+	@Override
+	public List<String> getSupportedProtocols() {
+		return Arrays.asList("v10.stomp", "v11.stomp", "v12.stomp");
+	}
+
+	/**
+	 * Handle incoming WebSocket messages from clients.
+	 */
+	public void handleMessageFromClient(WebSocketSession session, WebSocketMessage webSocketMessage,
+			MessageChannel outputChannel) {
+
+		Message<?> message;
+		try {
+			Assert.isInstanceOf(TextMessage.class,  webSocketMessage);
+			String payload = ((TextMessage)webSocketMessage).getPayload();
+			ByteBuffer byteBuffer = ByteBuffer.wrap(payload.getBytes(Charset.forName("UTF-8")));
+			message = this.stompDecoder.decode(byteBuffer);
+		}
+		catch (Throwable error) {
+			logger.error("Failed to parse STOMP frame, WebSocket message payload: ", error);
+			sendErrorMessage(session, error);
+			return;
+		}
+
+		try {
+			StompHeaderAccessor headers = StompHeaderAccessor.wrap(message);
+			if (SimpMessageType.HEARTBEAT.equals(headers.getMessageType())) {
+				logger.trace("Received heartbeat from client session=" + session.getId());
+			}
+			else {
+				logger.trace("Received message from client session=" + session.getId());
+			}
+
+			headers.setSessionId(session.getId());
+			headers.setUser(session.getPrincipal());
+
+			message = MessageBuilder.withPayload(message.getPayload()).setHeaders(headers).build();
+			outputChannel.send(message);
+		}
+		catch (Throwable t) {
+			logger.error("Terminating STOMP session due to failure to send message: ", t);
+			sendErrorMessage(session, t);
+		}
+	}
+
+	protected void sendErrorMessage(WebSocketSession session, Throwable error) {
+
+		StompHeaderAccessor headers = StompHeaderAccessor.create(StompCommand.ERROR);
+		headers.setMessage(error.getMessage());
+		Message<byte[]> message = MessageBuilder.withPayload(new byte[0]).setHeaders(headers).build();
+		String payload = new String(this.stompEncoder.encode(message), Charset.forName("UTF-8"));
+		try {
+			session.sendMessage(new TextMessage(payload));
+		}
+		catch (Throwable t) {
+			// ignore
+		}
+	}
+
+	/**
+	 * Handle STOMP messages going back out to WebSocket clients.
+	 */
+	@SuppressWarnings("unchecked")
+	@Override
+	public void handleMessageToClient(WebSocketSession session, Message<?> message) {
+
+		StompHeaderAccessor headers = StompHeaderAccessor.wrap(message);
+
+		if (headers.getMessageType() == SimpMessageType.CONNECT_ACK) {
+			StompHeaderAccessor connectedHeaders = StompHeaderAccessor.create(StompCommand.CONNECTED);
+			connectedHeaders.setVersion(getVersion(headers));
+			connectedHeaders.setHeartbeat(0, 0); // no heart-beat support with simple broker
+			headers = connectedHeaders;
+		}
+		else if (SimpMessageType.MESSAGE.equals(headers.getMessageType())) {
+			headers.updateStompCommandAsServerMessage();
+		}
+
+		if (headers.getCommand() == StompCommand.CONNECTED) {
+			afterStompSessionConnected(headers, session);
+		}
+
+		if (StompCommand.MESSAGE.equals(headers.getCommand()) && (headers.getSubscriptionId() == null)) {
+			logger.error("Ignoring message, no subscriptionId header: " + message);
+			return;
+		}
+
+		if (!(message.getPayload() instanceof byte[])) {
+			logger.error("Ignoring message, expected byte[] content: " + message);
+			return;
+		}
+
+		try {
+			message = MessageBuilder.withPayload(message.getPayload()).setHeaders(headers).build();
+			byte[] bytes = this.stompEncoder.encode((Message<byte[]>) message);
+			session.sendMessage(new TextMessage(new String(bytes, Charset.forName("UTF-8"))));
+		}
+		catch (Throwable t) {
+			sendErrorMessage(session, t);
+		}
+		finally {
+			if (StompCommand.ERROR.equals(headers.getCommand())) {
+				try {
+					session.close(CloseStatus.PROTOCOL_ERROR);
+				}
+				catch (IOException e) {
+					// Ignore
+				}
+			}
+		}
+	}
+
+	private String getVersion(StompHeaderAccessor connectAckHeaders) {
+
+		String name = StompHeaderAccessor.CONNECT_MESSAGE_HEADER;
+		Message<?> connectMessage = (Message<?>) connectAckHeaders.getHeader(name);
+		StompHeaderAccessor connectHeaders = StompHeaderAccessor.wrap(connectMessage);
+		Assert.notNull(connectMessage, "CONNECT_ACK does not contain original CONNECT " + connectAckHeaders);
+
+		Set<String> acceptVersions = connectHeaders.getAcceptVersion();
+		if (acceptVersions.contains("1.2")) {
+			return "1.2";
+		}
+		else if (acceptVersions.contains("1.1")) {
+			return "1.1";
+		}
+		else if (acceptVersions.isEmpty()) {
+			return null;
+		}
+		else {
+			throw new StompConversionException("Unsupported version '" + acceptVersions + "'");
+		}
+	}
+
+	private void afterStompSessionConnected(StompHeaderAccessor headers, WebSocketSession session) {
+		Principal principal = session.getPrincipal();
+		if (principal != null) {
+			headers.setNativeHeader(CONNECTED_USER_HEADER, principal.getName());
+			if (this.userSessionRegistry != null) {
+				this.userSessionRegistry.registerSessionId(principal.getName(), session.getId());
+			}
+		}
+	}
+
+	@Override
+	public String resolveSessionId(Message<?> message) {
+		StompHeaderAccessor headers = StompHeaderAccessor.wrap(message);
+		return headers.getSessionId();
+	}
+
+	@Override
+	public void afterSessionStarted(WebSocketSession session, MessageChannel outputChannel) {
+	}
+
+	@Override
+	public void afterSessionEnded(WebSocketSession session, CloseStatus closeStatus, MessageChannel outputChannel) {
+
+		if ((this.userSessionRegistry != null) && (session.getPrincipal() != null)) {
+			this.userSessionRegistry.unregisterSessionId(session.getPrincipal().getName(), session.getId());
+		}
+
+		StompHeaderAccessor headers = StompHeaderAccessor.create(StompCommand.DISCONNECT);
+		headers.setSessionId(session.getId());
+		Message<?> message = MessageBuilder.withPayload(new byte[0]).setHeaders(headers).build();
+		outputChannel.send(message);
+	}
+
+}

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

@@ -0,0 +1,94 @@
+/*
+ * Copyright 2002-2013 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.messaging;
+
+import java.util.List;
+
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.WebSocketMessage;
+import org.springframework.web.socket.WebSocketSession;
+
+
+/**
+ * A contract for handling WebSocket messages as part of a higher level protocol, referred
+ * to as "sub-protocol" in the WebSocket RFC specification. Handles both
+ * {@link WebSocketMessage}s from a client as well as {@link Message}s to a client.
+ * <p>
+ * Implementations of this interface can be configured on a
+ * {@link SubProtocolWebSocketHandler} which selects a sub-protocol handler to delegate
+ * messages to based on the sub-protocol requested by the client through the
+ * {@code Sec-WebSocket-Protocol} request header.
+ *
+ * @author Andy Wilkinson
+ * @author Rossen Stoyanchev
+ *
+ * @since 4.0
+ */
+public interface SubProtocolHandler {
+
+	/**
+	 * Return the list of sub-protocols supported by this handler, never {@code null}.
+	 */
+	List<String> getSupportedProtocols();
+
+	/**
+	 * Handle the given {@link WebSocketMessage} received from a client.
+	 *
+	 * @param session the client session
+	 * @param message the client message
+	 * @param outputChannel an output channel to send messages to
+	 */
+	void handleMessageFromClient(WebSocketSession session, WebSocketMessage message,
+			MessageChannel outputChannel) throws Exception;
+
+	/**
+	 * Handle the given {@link Message} to the client associated with the given WebSocket
+	 * session.
+	 *
+	 * @param session the client session
+	 * @param message the client message
+	 */
+	void handleMessageToClient(WebSocketSession session, Message<?> message) throws Exception;
+
+	/**
+	 * Resolve the session id from the given message or return {@code null}.
+	 *
+	 * @param message the message to resolve the session id from
+	 */
+	String resolveSessionId(Message<?> message);
+
+	/**
+	 * Invoked after a {@link WebSocketSession} has started.
+	 *
+	 * @param session the client session
+	 * @param outputChannel a channel
+	 */
+	void afterSessionStarted(WebSocketSession session, MessageChannel outputChannel) throws Exception;
+
+	/**
+	 * Invoked after a {@link WebSocketSession} has ended.
+	 *
+	 * @param session the client session
+	 * @param closeStatus the reason why the session was closed
+	 * @param outputChannel a channel
+	 */
+	void afterSessionEnded(WebSocketSession session, CloseStatus closeStatus,
+			MessageChannel outputChannel) throws Exception;
+
+}

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

@@ -0,0 +1,234 @@
+/*
+ * Copyright 2002-2013 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.messaging;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeMap;
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.MessageHandler;
+import org.springframework.messaging.MessagingException;
+import org.springframework.util.Assert;
+import org.springframework.util.CollectionUtils;
+import org.springframework.util.StringUtils;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.WebSocketMessage;
+import org.springframework.web.socket.WebSocketSession;
+
+
+/**
+ * An implementation of {@link WebSocketHandler} that delegates incoming WebSocket
+ * messages to a {@link SubProtocolHandler} along with a {@link MessageChannel} to
+ * which the sub-protocol handler can send messages from WebSocket clients to
+ * the application.
+ * <p>
+ * Also an implementation of {@link MessageHandler} that finds the WebSocket
+ * session associated with the {@link Message} and passes it, along with the message,
+ * to the sub-protocol handler to send messages from the application back to the
+ * client.
+ *
+ * @author Rossen Stoyanchev
+ * @author Andy Wilkinson
+ *
+ * @since 4.0
+ */
+public class SubProtocolWebSocketHandler implements WebSocketHandler, MessageHandler {
+
+	private final Log logger = LogFactory.getLog(SubProtocolWebSocketHandler.class);
+
+	private final MessageChannel clientOutboundChannel;
+
+	private final Map<String, SubProtocolHandler> protocolHandlers =
+			new TreeMap<String, SubProtocolHandler>(String.CASE_INSENSITIVE_ORDER);
+
+	private SubProtocolHandler defaultProtocolHandler;
+
+	private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<String, WebSocketSession>();
+
+
+	public SubProtocolWebSocketHandler(MessageChannel clientOutboundChannel) {
+		Assert.notNull(clientOutboundChannel, "clientOutboundChannel is required");
+		this.clientOutboundChannel = clientOutboundChannel;
+	}
+
+
+	/**
+	 * Configure one or more handlers to use depending on the sub-protocol requested by
+	 * the client in the WebSocket handshake request.
+	 *
+	 * @param protocolHandlers the sub-protocol handlers to use
+	 */
+	public void setProtocolHandlers(List<SubProtocolHandler> protocolHandlers) {
+		this.protocolHandlers.clear();
+		for (SubProtocolHandler handler: protocolHandlers) {
+			addProtocolHandler(handler);
+		}
+	}
+
+	/**
+	 * Register a sub-protocol handler.
+	 */
+	public void addProtocolHandler(SubProtocolHandler handler) {
+		List<String> protocols = handler.getSupportedProtocols();
+		if (CollectionUtils.isEmpty(protocols)) {
+			logger.warn("No sub-protocols, ignoring handler " + handler);
+			return;
+		}
+		for (String protocol: protocols) {
+			SubProtocolHandler replaced = this.protocolHandlers.put(protocol, handler);
+			if ((replaced != null) && (replaced != handler) ) {
+				throw new IllegalStateException("Failed to map handler " + handler
+						+ " to protocol '" + protocol + "', it is already mapped to handler " + replaced);
+			}
+		}
+	}
+
+	/**
+	 * @return the configured sub-protocol handlers
+	 */
+	public Map<String, SubProtocolHandler> getProtocolHandlers() {
+		return this.protocolHandlers;
+	}
+
+	/**
+	 * Set the {@link SubProtocolHandler} to use when the client did not request a
+	 * sub-protocol.
+	 *
+	 * @param defaultProtocolHandler the default handler
+	 */
+	public void setDefaultProtocolHandler(SubProtocolHandler defaultProtocolHandler) {
+		this.defaultProtocolHandler = defaultProtocolHandler;
+		if (this.protocolHandlers.isEmpty()) {
+			setProtocolHandlers(Arrays.asList(defaultProtocolHandler));
+		}
+	}
+
+	/**
+	 * @return the default sub-protocol handler to use
+	 */
+	public SubProtocolHandler getDefaultProtocolHandler() {
+		return this.defaultProtocolHandler;
+	}
+
+	/**
+	 * Return all supported protocols.
+	 */
+	public Set<String> getSupportedProtocols() {
+		return this.protocolHandlers.keySet();
+	}
+
+	@Override
+	public void afterConnectionEstablished(WebSocketSession session) throws Exception {
+		this.sessions.put(session.getId(), session);
+		findProtocolHandler(session).afterSessionStarted(session, this.clientOutboundChannel);
+	}
+
+	protected final SubProtocolHandler findProtocolHandler(WebSocketSession session) {
+		SubProtocolHandler handler;
+		String protocol = session.getAcceptedProtocol();
+		if (!StringUtils.isEmpty(protocol)) {
+			handler = this.protocolHandlers.get(protocol);
+			Assert.state(handler != null,
+					"No handler for sub-protocol '" + protocol + "', handlers=" + this.protocolHandlers);
+		}
+		else {
+			if (this.defaultProtocolHandler != null) {
+				handler = this.defaultProtocolHandler;
+			}
+			else {
+				Set<SubProtocolHandler> handlers = new HashSet<SubProtocolHandler>(this.protocolHandlers.values());
+				if (handlers.size() == 1) {
+					handler = handlers.iterator().next();
+				}
+				else {
+					throw new IllegalStateException(
+							"No sub-protocol was requested and a default sub-protocol handler was not configured");
+				}
+			}
+		}
+		return handler;
+	}
+
+	@Override
+	public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
+		findProtocolHandler(session).handleMessageFromClient(session, message, this.clientOutboundChannel);
+	}
+
+	@Override
+	public void handleMessage(Message<?> message) throws MessagingException {
+
+		String sessionId = resolveSessionId(message);
+		if (sessionId == null) {
+			logger.error("sessionId not found in message " + message);
+			return;
+		}
+
+		WebSocketSession session = this.sessions.get(sessionId);
+		if (session == null) {
+			logger.error("Session not found for session with id " + sessionId);
+			return;
+		}
+
+		try {
+			findProtocolHandler(session).handleMessageToClient(session, message);
+		}
+		catch (Exception e) {
+			logger.error("Failed to send message to client " + message, e);
+		}
+	}
+
+	private String resolveSessionId(Message<?> message) {
+		for (SubProtocolHandler handler : this.protocolHandlers.values()) {
+			String sessionId = handler.resolveSessionId(message);
+			if (sessionId != null) {
+				return sessionId;
+			}
+		}
+		if (this.defaultProtocolHandler != null) {
+			String sessionId = this.defaultProtocolHandler.resolveSessionId(message);
+			if (sessionId != null) {
+				return sessionId;
+			}
+		}
+		return null;
+	}
+
+	@Override
+	public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
+	}
+
+	@Override
+	public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
+		this.sessions.remove(session.getId());
+		findProtocolHandler(session).afterSessionEnded(session, closeStatus, this.clientOutboundChannel);
+	}
+
+	@Override
+	public boolean supportsPartialMessages() {
+		return false;
+	}
+
+}

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

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2002-2013 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.messaging.config;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.messaging.simp.config.MessageBrokerRegistry;
+import org.springframework.util.CollectionUtils;
+
+
+/**
+ * A {@link WebSocketMessageBrokerConfigurationSupport} extension that detects beans of type
+ * {@link WebSocketMessageBrokerConfigurer}
+ * and delegates to all of them allowing callback style customization of the
+ * configuration provided in {@link WebSocketMessageBrokerConfigurationSupport}.
+ *
+ * <p>This class is typically imported via {@link EnableWebSocketMessageBroker}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+@Configuration
+public class DelegatingWebSocketMessageBrokerConfiguration extends WebSocketMessageBrokerConfigurationSupport {
+
+	private List<WebSocketMessageBrokerConfigurer> configurers = new ArrayList<WebSocketMessageBrokerConfigurer>();
+
+
+	@Autowired(required=false)
+	public void setConfigurers(List<WebSocketMessageBrokerConfigurer> configurers) {
+		if (CollectionUtils.isEmpty(configurers)) {
+			return;
+		}
+		this.configurers.addAll(configurers);
+	}
+
+
+	@Override
+	protected void registerStompEndpoints(StompEndpointRegistry registry) {
+		for (WebSocketMessageBrokerConfigurer c : this.configurers) {
+			c.registerStompEndpoints(registry);
+		}
+	}
+
+	@Override
+	protected void configureMessageBroker(MessageBrokerRegistry registry) {
+		for (WebSocketMessageBrokerConfigurer c : this.configurers) {
+			c.configureMessageBroker(registry);
+		}
+	}
+
+}

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

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2002-2013 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.messaging.config;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+import org.springframework.context.annotation.Import;
+
+
+/**
+ * Add this annotation to an {@code @Configuration} class to enable broker-backed
+ * messaging over WebSocket using a higher-level messaging sub-protocol.
+ *
+ * <pre class="code">
+ * &#064;Configuration
+ * &#064;EnableWebSocketMessageBroker
+ * public class MyWebSocketConfig {
+ *
+ * }
+ * </pre>
+ * <p>
+ * Customize the imported configuration by implementing the
+ * {@link WebSocketMessageBrokerConfigurer} interface:
+ *
+ * <pre class="code">
+ * &#064;Configuration
+ * &#064;EnableWebSocketMessageBroker
+ * public class MyConfiguration implements implements WebSocketMessageBrokerConfigurer {
+ *
+ * 	&#064;Override
+ * 	public void registerStompEndpoints(StompEndpointRegistry registry) {
+ * 		registry.addEndpoint("/portfolio").withSockJS();
+ * 	}
+ *
+ * 	&#064;Bean
+ * 	public void configureMessageBroker(MessageBrokerRegistry registry) {
+ * 		registry.enableStompBrokerRelay("/queue/", "/topic/");
+ * 		registry.setApplicationDestinationPrefixes("/app/");
+ * 	}
+ * }
+ * </pre>
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.TYPE)
+@Documented
+@Import(DelegatingWebSocketMessageBrokerConfiguration.class)
+public @interface EnableWebSocketMessageBroker {
+
+}

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2002-2013 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.messaging.config;
+
+/**
+ * A contract for registering STOMP over WebSocket endpoints.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public interface StompEndpointRegistry {
+
+	/**
+	 * Register a STOMP over WebSocket endpoint at the given mapping path.
+	 */
+	StompWebSocketEndpointRegistration addEndpoint(String... paths);
+
+}

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

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2002-2013 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.messaging.config;
+
+import org.springframework.web.socket.server.HandshakeHandler;
+import org.springframework.web.socket.server.config.SockJsServiceRegistration;
+
+/**
+ * A contract for configuring a STOMP over WebSocket endpoint.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public interface StompWebSocketEndpointRegistration {
+
+	/**
+	 * Enable SockJS fallback options.
+	 */
+	SockJsServiceRegistration withSockJS();
+
+	/**
+	 * Configure the HandshakeHandler to use.
+	 */
+	StompWebSocketEndpointRegistration setHandshakeHandler(HandshakeHandler handshakeHandler);
+
+}

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

@@ -0,0 +1,128 @@
+/*
+ * Copyright 2002-2013 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.messaging.config;
+
+import java.util.*;
+
+import org.springframework.messaging.simp.handler.UserSessionRegistry;
+import org.springframework.web.socket.messaging.StompSubProtocolHandler;
+import org.springframework.scheduling.TaskScheduler;
+import org.springframework.util.Assert;
+import org.springframework.util.MultiValueMap;
+import org.springframework.web.HttpRequestHandler;
+import org.springframework.web.servlet.handler.AbstractHandlerMapping;
+import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
+import org.springframework.web.socket.support.WebSocketHandlerDecorator;
+
+
+/**
+ * A registry for STOMP over WebSocket endpoints that maps the endpoints with a
+ * {@link SimpleUrlHandlerMapping} for use in Spring MVC.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public class WebMvcStompEndpointRegistry implements StompEndpointRegistry {
+
+	private final WebSocketHandler webSocketHandler;
+
+	private final SubProtocolWebSocketHandler subProtocolWebSocketHandler;
+
+	private final StompSubProtocolHandler stompHandler;
+
+	private final List<WebMvcStompWebSocketEndpointRegistration> registrations =
+			new ArrayList<WebMvcStompWebSocketEndpointRegistration>();
+
+	private final TaskScheduler sockJsScheduler;
+
+	private int order = 1;
+
+
+	public WebMvcStompEndpointRegistry(WebSocketHandler webSocketHandler,
+			UserSessionRegistry userSessionRegistry, TaskScheduler defaultSockJsTaskScheduler) {
+
+		Assert.notNull(webSocketHandler);
+		Assert.notNull(userSessionRegistry);
+
+		this.webSocketHandler = webSocketHandler;
+		this.subProtocolWebSocketHandler = unwrapSubProtocolWebSocketHandler(webSocketHandler);
+		this.stompHandler = new StompSubProtocolHandler();
+		this.stompHandler.setUserSessionRegistry(userSessionRegistry);
+		this.sockJsScheduler = defaultSockJsTaskScheduler;
+	}
+
+	private static SubProtocolWebSocketHandler unwrapSubProtocolWebSocketHandler(WebSocketHandler webSocketHandler) {
+
+		WebSocketHandler actual = (webSocketHandler instanceof WebSocketHandlerDecorator) ?
+				((WebSocketHandlerDecorator) webSocketHandler).getLastHandler() : webSocketHandler;
+
+		Assert.isInstanceOf(SubProtocolWebSocketHandler.class, actual,
+						"No SubProtocolWebSocketHandler found: " + webSocketHandler);
+
+		return (SubProtocolWebSocketHandler) actual;
+	}
+
+
+	@Override
+	public StompWebSocketEndpointRegistration addEndpoint(String... paths) {
+
+		this.subProtocolWebSocketHandler.addProtocolHandler(this.stompHandler);
+		Set<String> subProtocols = this.subProtocolWebSocketHandler.getSupportedProtocols();
+
+		WebMvcStompWebSocketEndpointRegistration registration = new WebMvcStompWebSocketEndpointRegistration(
+				paths, this.webSocketHandler, subProtocols, this.sockJsScheduler);
+		this.registrations.add(registration);
+
+		return registration;
+	}
+
+	/**
+	 * Set the order for the resulting {@link SimpleUrlHandlerMapping} relative to
+	 * other handler mappings configured in Spring MVC.
+	 * <p>
+	 * The default value is 1.
+	 */
+	public void setOrder(int order) {
+		this.order = order;
+	}
+
+	public int getOrder() {
+		return this.order;
+	}
+
+	/**
+	 * Returns a handler mapping with the mapped ViewControllers; or {@code null} in case of no registrations.
+	 */
+	protected AbstractHandlerMapping getHandlerMapping() {
+		Map<String, Object> urlMap = new LinkedHashMap<String, Object>();
+		for (WebMvcStompWebSocketEndpointRegistration registration : this.registrations) {
+			MultiValueMap<HttpRequestHandler, String> mappings = registration.getMappings();
+			for (HttpRequestHandler httpHandler : mappings.keySet()) {
+				for (String pattern : mappings.get(httpHandler)) {
+					urlMap.put(pattern, httpHandler);
+				}
+			}
+		}
+		SimpleUrlHandlerMapping hm = new SimpleUrlHandlerMapping();
+		hm.setUrlMap(urlMap);
+		hm.setOrder(this.order);
+		return hm;
+	}
+
+}

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

@@ -0,0 +1,137 @@
+/*
+ * Copyright 2002-2013 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.messaging.config;
+
+import java.util.Set;
+
+import org.springframework.scheduling.TaskScheduler;
+import org.springframework.util.Assert;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
+import org.springframework.util.ObjectUtils;
+import org.springframework.web.HttpRequestHandler;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.server.DefaultHandshakeHandler;
+import org.springframework.web.socket.server.HandshakeHandler;
+import org.springframework.web.socket.server.config.SockJsServiceRegistration;
+import org.springframework.web.socket.server.support.WebSocketHttpRequestHandler;
+import org.springframework.web.socket.sockjs.SockJsHttpRequestHandler;
+import org.springframework.web.socket.sockjs.SockJsService;
+import org.springframework.web.socket.sockjs.transport.handler.WebSocketTransportHandler;
+
+
+/**
+ * An abstract base class class for configuring STOMP over WebSocket/SockJS endpoints.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public class WebMvcStompWebSocketEndpointRegistration implements StompWebSocketEndpointRegistration {
+
+	private final String[] paths;
+
+	private final WebSocketHandler webSocketHandler;
+
+	private final String[] subProtocols;
+
+	private final TaskScheduler sockJsTaskScheduler;
+
+	private HandshakeHandler handshakeHandler;
+
+	private StompSockJsServiceRegistration registration;
+
+
+	public WebMvcStompWebSocketEndpointRegistration(String[] paths, WebSocketHandler webSocketHandler,
+			Set<String> subProtocols, TaskScheduler sockJsTaskScheduler) {
+
+		Assert.notEmpty(paths, "No paths specified");
+		Assert.notNull(webSocketHandler, "'webSocketHandler' is required");
+		Assert.notNull(subProtocols, "'subProtocols' is required");
+
+		this.paths = paths;
+		this.webSocketHandler = webSocketHandler;
+		this.subProtocols = subProtocols.toArray(new String[subProtocols.size()]);
+		this.sockJsTaskScheduler = sockJsTaskScheduler;
+
+		this.handshakeHandler = new DefaultHandshakeHandler();
+		updateHandshakeHandler();
+	}
+
+	private void updateHandshakeHandler() {
+		if (handshakeHandler instanceof DefaultHandshakeHandler) {
+			DefaultHandshakeHandler defaultHandshakeHandler = (DefaultHandshakeHandler) handshakeHandler;
+			if (ObjectUtils.isEmpty(defaultHandshakeHandler.getSupportedProtocols())) {
+				defaultHandshakeHandler.setSupportedProtocols(this.subProtocols);
+			}
+		}
+	}
+
+	/**
+	 * Provide a custom or pre-configured {@link HandshakeHandler}.
+	 */
+	@Override
+	public StompWebSocketEndpointRegistration setHandshakeHandler(HandshakeHandler handshakeHandler) {
+		Assert.notNull(handshakeHandler, "'handshakeHandler' must not be null");
+		this.handshakeHandler = handshakeHandler;
+		updateHandshakeHandler();
+		return this;
+	}
+
+	/**
+	 * Enable SockJS fallback options.
+	 */
+	@Override
+	public SockJsServiceRegistration withSockJS() {
+		this.registration = new StompSockJsServiceRegistration(this.sockJsTaskScheduler);
+		WebSocketTransportHandler transportHandler = new WebSocketTransportHandler(this.handshakeHandler);
+		this.registration.setTransportHandlerOverrides(transportHandler);
+		return this.registration;
+	}
+
+	protected final MultiValueMap<HttpRequestHandler, String> getMappings() {
+		MultiValueMap<HttpRequestHandler, String> mappings = new LinkedMultiValueMap<HttpRequestHandler, String>();
+		if (this.registration != null) {
+			SockJsService sockJsService = this.registration.getSockJsService();
+			for (String path : this.paths) {
+				String pattern = path.endsWith("/") ? path + "**" : path + "/**";
+				SockJsHttpRequestHandler handler = new SockJsHttpRequestHandler(sockJsService, this.webSocketHandler);
+				mappings.add(handler, pattern);
+			}
+		}
+		else {
+			for (String path : this.paths) {
+				WebSocketHttpRequestHandler handler =
+						new WebSocketHttpRequestHandler(this.webSocketHandler, this.handshakeHandler);
+				mappings.add(handler, path);
+			}
+		}
+		return mappings;
+	}
+
+
+	private static class StompSockJsServiceRegistration extends SockJsServiceRegistration {
+
+		public StompSockJsServiceRegistration(TaskScheduler defaultTaskScheduler) {
+			super(defaultTaskScheduler);
+		}
+
+		protected SockJsService getSockJsService() {
+			return super.getSockJsService();
+		}
+	}
+
+}

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

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2002-2013 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.messaging.config;
+
+import org.springframework.context.annotation.Bean;
+import org.springframework.messaging.simp.config.AbstractMessageBrokerConfiguration;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
+import org.springframework.web.servlet.HandlerMapping;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.server.config.SockJsServiceRegistration;
+import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
+
+
+/**
+ * Extends {@link AbstractMessageBrokerConfiguration} and adds configuration for
+ * receiving and responding to STOMP messages from WebSocket clients.
+ * <p>
+ * Typically used in conjunction with
+ * {@link EnableWebSocketMessageBroker @EnableWebSocketMessageBroker} but can
+ * also be extended directly.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public abstract class WebSocketMessageBrokerConfigurationSupport extends AbstractMessageBrokerConfiguration {
+
+
+	protected WebSocketMessageBrokerConfigurationSupport() {
+	}
+
+	@Bean
+	public HandlerMapping stompWebSocketHandlerMapping() {
+
+		WebMvcStompEndpointRegistry registry = new WebMvcStompEndpointRegistry(
+				subProtocolWebSocketHandler(), userSessionRegistry(), messageBrokerSockJsTaskScheduler());
+
+		registerStompEndpoints(registry);
+		return registry.getHandlerMapping();
+	}
+
+	@Bean
+	public WebSocketHandler subProtocolWebSocketHandler() {
+		SubProtocolWebSocketHandler handler = new SubProtocolWebSocketHandler(clientInboundChannel());
+		clientOutboundChannel().subscribe(handler);
+		return handler;
+	}
+
+	/**
+	 * The default TaskScheduler to use if none is configured via
+	 * {@link SockJsServiceRegistration#setTaskScheduler(org.springframework.scheduling.TaskScheduler)}, i.e.
+	 * <pre class="code">
+	 * &#064;Configuration
+	 * &#064;EnableWebSocketMessageBroker
+	 * public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
+	 *
+	 *   public void registerStompEndpoints(StompEndpointRegistry registry) {
+	 *     registry.addEndpoint("/stomp").withSockJS().setTaskScheduler(myScheduler());
+	 *   }
+	 *
+	 *   // ...
+	 *
+	 * }
+	 * </pre>
+	 */
+	@Bean
+	public ThreadPoolTaskScheduler messageBrokerSockJsTaskScheduler() {
+		ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
+		scheduler.setThreadNamePrefix("MessageBrokerSockJS-");
+		return scheduler;
+	}
+
+	protected abstract void registerStompEndpoints(StompEndpointRegistry registry);
+
+}

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

@@ -0,0 +1,43 @@
+/*
+ * Copyright 2002-2013 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.messaging.config;
+
+
+import org.springframework.messaging.simp.config.MessageBrokerRegistry;
+
+/**
+ * Defines methods for configuring message handling with simple messaging
+ * protocols (e.g. STOMP) from WebSocket clients. Typically used to customize
+ * the configuration provided via
+ * {@link EnableWebSocketMessageBroker @EnableWebSocketMessageBroker}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public interface WebSocketMessageBrokerConfigurer {
+
+	/**
+	 * Configure STOMP over WebSocket endpoints.
+	 */
+	void registerStompEndpoints(StompEndpointRegistry registry);
+
+	/**
+	 * Configure message broker options.
+	 */
+	void configureMessageBroker(MessageBrokerRegistry registry);
+
+}

spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHttpRequestHandler.java

@@ -84,6 +84,13 @@ public class WebSocketHttpRequestHandler implements HttpRequestHandler {
 		return this.wsHandler;
 	}
 
+	/**
+	 * Return the HandshakeHandler.
+	 */
+	public HandshakeHandler getHandshakeHandler() {
+		return this.handshakeHandler;
+	}
+
 	/**
 	 * Configure one or more WebSocket handshake request interceptors.
 	 */