Add Java config support for WebSocket and STOMP

Issue: SPR-10835

Author: rstoyanchev

build.gradle

@@ -317,6 +317,7 @@ project("spring-messaging") {
 		compile(project(":spring-core"))
 		compile(project(":spring-context"))
 		optional(project(":spring-websocket"))
+		optional(project(":spring-webmvc"))
 		optional("com.fasterxml.jackson.core:jackson-databind:2.2.0")
 		optional("org.projectreactor:reactor-core:1.0.0.M2")
 		optional("org.projectreactor:reactor-tcp:1.0.0.M2")

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

@@ -17,8 +17,10 @@
 package org.springframework.messaging.handler.websocket;
 
 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;
 
@@ -79,21 +81,25 @@ public SubProtocolWebSocketHandler(MessageChannel outputChannel) {
 	public void setProtocolHandlers(List<SubProtocolHandler> protocolHandlers) {
 		this.protocolHandlers.clear();
 		for (SubProtocolHandler handler: protocolHandlers) {
-			List<String> protocols = handler.getSupportedProtocols();
-			if (CollectionUtils.isEmpty(protocols)) {
-				logger.warn("No sub-protocols, ignoring handler " + handler);
-				continue;
-			}
-			for (String protocol: protocols) {
-				SubProtocolHandler replaced = this.protocolHandlers.put(protocol, handler);
-				if (replaced != null) {
-					throw new IllegalStateException("Failed to map handler " + handler
-							+ " to protocol '" + protocol + "', it is already mapped to handler " + replaced);
-				}
-			}
+			addProtocolHandler(handler);
 		}
-		if ((this.protocolHandlers.size() == 1) &&(this.defaultProtocolHandler == null)) {
-			this.defaultProtocolHandler = this.protocolHandlers.values().iterator().next();
+	}
+
+	/**
+	 * 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);
+			}
 		}
 	}
 
@@ -128,10 +134,10 @@ public SubProtocolHandler getDefaultProtocolHandler() {
 	@Override
 	public void afterConnectionEstablished(WebSocketSession session) throws Exception {
 		this.sessions.put(session.getId(), session);
-		getProtocolHandler(session).afterSessionStarted(session, this.outputChannel);
+		findProtocolHandler(session).afterSessionStarted(session, this.outputChannel);
 	}
 
-	protected final SubProtocolHandler getProtocolHandler(WebSocketSession session) {
+	protected final SubProtocolHandler findProtocolHandler(WebSocketSession session) {
 		SubProtocolHandler handler;
 		String protocol = session.getAcceptedProtocol();
 		if (!StringUtils.isEmpty(protocol)) {
@@ -140,16 +146,26 @@ protected final SubProtocolHandler getProtocolHandler(WebSocketSession session)
 					"No handler for sub-protocol '" + protocol + "', handlers=" + this.protocolHandlers);
 		}
 		else {
-			handler = this.defaultProtocolHandler;
-			Assert.state(handler != null,
-					"No sub-protocol was requested and a default sub-protocol handler was not configured");
+			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 {
-		getProtocolHandler(session).handleMessageFromClient(session, message, this.outputChannel);
+		findProtocolHandler(session).handleMessageFromClient(session, message, this.outputChannel);
 	}
 
 	@Override
@@ -168,7 +184,7 @@ public void handleMessage(Message<?> message) throws MessagingException {
 		}
 
 		try {
-			getProtocolHandler(session).handleMessageToClient(session, message);
+			findProtocolHandler(session).handleMessageToClient(session, message);
 		}
 		catch (Exception e) {
 			logger.error("Failed to send message to client " + message, e);
@@ -198,7 +214,7 @@ public void handleTransportError(WebSocketSession session, Throwable exception)
 	@Override
 	public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
 		this.sessions.remove(session.getId());
-		getProtocolHandler(session).afterSessionEnded(session, closeStatus, this.outputChannel);
+		findProtocolHandler(session).afterSessionEnded(session, closeStatus, this.outputChannel);
 	}
 
 	@Override

spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractBrokerRegistration.java

@@ -0,0 +1,60 @@
+/*
+ * 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.messaging.simp.config;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.simp.handler.AbstractBrokerMessageHandler;
+
+import reactor.util.Assert;
+
+
+/**
+ * Base class for message broker registration classes.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public abstract class AbstractBrokerRegistration {
+
+	private final MessageChannel webSocketReplyChannel;
+
+	private final String[] destinationPrefixes;
+
+
+	public AbstractBrokerRegistration(MessageChannel webSocketReplyChannel, String[] destinationPrefixes) {
+		Assert.notNull(webSocketReplyChannel, "");
+		this.webSocketReplyChannel = webSocketReplyChannel;
+		this.destinationPrefixes = destinationPrefixes;
+	}
+
+
+	protected MessageChannel getWebSocketReplyChannel() {
+		return this.webSocketReplyChannel;
+	}
+
+	protected Collection<String> getDestinationPrefixes() {
+		return (this.destinationPrefixes != null)
+				? Arrays.<String>asList(this.destinationPrefixes) : Collections.<String>emptyList();
+	}
+
+	protected abstract AbstractBrokerMessageHandler getMessageHandler();
+
+}

spring-messaging/src/main/java/org/springframework/messaging/simp/config/DelegatingWebSocketMessageBrokerConfiguration.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.messaging.simp.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.util.CollectionUtils;
+
+
+/**
+ * A {@link WebSocketMessageBrokerConfiguration} 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(MessageBrokerConfigurer configurer) {
+		for (WebSocketMessageBrokerConfigurer c : this.configurers) {
+			c.configureMessageBroker(configurer);
+		}
+	}
+
+}

spring-messaging/src/main/java/org/springframework/messaging/simp/config/EnableWebSocketMessageBroker.java

@@ -0,0 +1,65 @@
+/*
+ * 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.messaging.simp.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(MessageBrokerConfigurer configurer) {
+ * 		configurer.enableStompBrokerRelay("/queue/", "/topic/");
+ * 		configurer.setAnnotationMethodDestinationPrefixes("/app/");
+ * 	}
+ * }
+ * </pre>
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.TYPE)
+@Documented
+@Import(DelegatingWebSocketMessageBrokerConfiguration.class)
+public @interface EnableWebSocketMessageBroker {
+}

spring-messaging/src/main/java/org/springframework/messaging/simp/config/MessageBrokerConfigurer.java

@@ -0,0 +1,84 @@
+/*
+ * 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.messaging.simp.config;
+
+import java.util.Arrays;
+import java.util.Collection;
+
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.simp.handler.AbstractBrokerMessageHandler;
+
+import reactor.util.Assert;
+
+
+/**
+ * A helper class for configuring message broker options.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public class MessageBrokerConfigurer {
+
+	private final MessageChannel webSocketReplyChannel;
+
+	private SimpleBrokerRegistration simpleBroker;
+
+	private StompBrokerRelayRegistration stompRelay;
+
+	private String[] annotationMethodDestinationPrefixes;
+
+
+	public MessageBrokerConfigurer(MessageChannel webSocketReplyChannel) {
+		Assert.notNull(webSocketReplyChannel);
+		this.webSocketReplyChannel = webSocketReplyChannel;
+	}
+
+	public SimpleBrokerRegistration enableSimpleBroker(String... destinationPrefixes) {
+		this.simpleBroker = new SimpleBrokerRegistration(this.webSocketReplyChannel, destinationPrefixes);
+		return this.simpleBroker;
+	}
+
+	public StompBrokerRelayRegistration enableStompBrokerRelay(String... destinationPrefixes) {
+		this.stompRelay = new StompBrokerRelayRegistration(this.webSocketReplyChannel, destinationPrefixes);
+		return this.stompRelay;
+	}
+
+	public MessageBrokerConfigurer setAnnotationMethodDestinationPrefixes(String... destinationPrefixes) {
+		this.annotationMethodDestinationPrefixes = destinationPrefixes;
+		return this;
+	}
+
+	protected AbstractBrokerMessageHandler getSimpleBroker() {
+		initSimpleBrokerIfNecessary();
+		return (this.simpleBroker != null) ? this.simpleBroker.getMessageHandler() : null;
+	}
+
+	protected void initSimpleBrokerIfNecessary() {
+		if ((this.simpleBroker == null) && (this.stompRelay == null)) {
+			this.simpleBroker = new SimpleBrokerRegistration(this.webSocketReplyChannel, null);
+		}
+	}
+
+	protected AbstractBrokerMessageHandler getStompBrokerRelay() {
+		return (this.stompRelay != null) ? this.stompRelay.getMessageHandler() : null;
+	}
+
+	protected Collection<String> getAnnotationMethodDestinationPrefixes() {
+		return (this.annotationMethodDestinationPrefixes != null)
+				? Arrays.asList(this.annotationMethodDestinationPrefixes) : null;
+	}
+}