Allow configuring custom argument types

The WebSocket messaging namespace now exposes configuration options for
custom argument resolvers and return value handlers.

Issue: SPR-12217

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

@@ -401,9 +401,30 @@ class MessageBrokerBeanDefinitionParser implements BeanDefinitionParser {
 			String pathMatcherRef = messageBrokerElement.getAttribute("path-matcher");
 			beanDef.getPropertyValues().add("pathMatcher", new RuntimeBeanReference(pathMatcherRef));
 		}
+
+		Element resolversElement = DomUtils.getChildElementByTagName(messageBrokerElement, "argument-resolvers");
+		if (resolversElement != null) {
+			values.add("customArgumentResolvers", extractBeanSubElements(resolversElement, context));
+		}
+
+		Element handlersElement = DomUtils.getChildElementByTagName(messageBrokerElement, "return-value-handlers");
+		if (handlersElement != null) {
+			values.add("customReturnValueHandlers", extractBeanSubElements(handlersElement, context));
+		}
+
 		registerBeanDef(beanDef, context, source);
 	}
 
+	private ManagedList<Object> extractBeanSubElements(Element parentElement, ParserContext parserContext) {
+		ManagedList<Object> list = new ManagedList<Object>();
+		list.setSource(parserContext.extractSource(parentElement));
+		for (Element beanElement : DomUtils.getChildElementsByTagName(parentElement, "bean", "ref")) {
+			Object object = parserContext.getDelegate().parsePropertySubElement(beanElement, null);
+			list.add(object);
+		}
+		return list;
+	}
+
 	private RuntimeBeanReference registerUserDestinationResolver(Element brokerElem,
 			RuntimeBeanReference userSessionRegistry, ParserContext context, Object source) {
 

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

@@ -17,6 +17,8 @@
 package org.springframework.web.socket.config.annotation;
 
 import org.springframework.messaging.converter.MessageConverter;
+import org.springframework.messaging.handler.invocation.HandlerMethodArgumentResolver;
+import org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler;
 import org.springframework.messaging.simp.config.ChannelRegistration;
 import org.springframework.messaging.simp.config.MessageBrokerRegistry;
 
@@ -49,6 +51,14 @@ public abstract class AbstractWebSocketMessageBrokerConfigurer implements WebSoc
 		return true;
 	}
 
+	@Override
+	public void addArgumentResolvers(List<HandlerMethodArgumentResolver> argumentResolvers) {
+	}
+
+	@Override
+	public void addReturnValueHandlers(List<HandlerMethodReturnValueHandler> returnValueHandlers) {
+	}
+
 	@Override
 	public void configureMessageBroker(MessageBrokerRegistry registry) {
 	}

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

@@ -22,6 +22,8 @@ import java.util.List;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.messaging.converter.MessageConverter;
+import org.springframework.messaging.handler.invocation.HandlerMethodArgumentResolver;
+import org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler;
 import org.springframework.messaging.simp.config.ChannelRegistration;
 import org.springframework.messaging.simp.config.MessageBrokerRegistry;
 import org.springframework.util.CollectionUtils;
@@ -79,6 +81,20 @@ public class DelegatingWebSocketMessageBrokerConfiguration extends WebSocketMess
 		}
 	}
 
+	@Override
+	protected void addArgumentResolvers(List<HandlerMethodArgumentResolver> argumentResolvers) {
+		for (WebSocketMessageBrokerConfigurer c : this.configurers) {
+			c.addArgumentResolvers(argumentResolvers);
+		}
+	}
+
+	@Override
+	protected void addReturnValueHandlers(List<HandlerMethodReturnValueHandler> returnValueHandlers) {
+		for (WebSocketMessageBrokerConfigurer c : this.configurers) {
+			c.addReturnValueHandlers(returnValueHandlers);
+		}
+	}
+
 	@Override
 	protected boolean configureMessageConverters(List<MessageConverter> messageConverters) {
 		boolean registerDefaults = true;

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

@@ -18,6 +18,8 @@ package org.springframework.web.socket.config.annotation;
 
 
 import org.springframework.messaging.converter.MessageConverter;
+import org.springframework.messaging.handler.invocation.HandlerMethodArgumentResolver;
+import org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler;
 import org.springframework.messaging.simp.config.ChannelRegistration;
 import org.springframework.messaging.simp.config.MessageBrokerRegistry;
 
@@ -62,6 +64,26 @@ public interface WebSocketMessageBrokerConfigurer {
 	 */
 	void configureClientOutboundChannel(ChannelRegistration registration);
 
+	/**
+	 * Add resolvers to support custom controller method argument types.
+	 * <p>This does not override the built-in support for resolving handler
+	 * method arguments. To customize the built-in support for argument
+	 * resolution, configure {@code SimpAnnotationMethodMessageHandler} directly.
+	 * @param argumentResolvers initially an empty list
+	 * @since 4.1.1
+	 */
+	void addArgumentResolvers(List<HandlerMethodArgumentResolver> argumentResolvers);
+
+	/**
+	 * Add handlers to support custom controller method return value types.
+	 * <p>Using this option does not override the built-in support for handling
+	 * return values. To customize the built-in support for handling return
+	 * values, configure  {@code SimpAnnotationMethodMessageHandler} directly.
+	 * @param returnValueHandlers initially an empty list
+	 * @since 4.1.1
+	 */
+	void addReturnValueHandlers(List<HandlerMethodReturnValueHandler> returnValueHandlers);
+
 	/**
 	 * Configure the message converters to use when extracting the payload of
 	 * messages in annotated methods and when sending messages (e.g. through the

spring-websocket/src/main/resources/org/springframework/web/socket/config/spring-websocket-4.1.xsd

@@ -615,6 +615,70 @@
 					<xsd:element name="simple-broker" type="simple-broker"/>
 					<xsd:element name="stomp-broker-relay" type="stomp-broker-relay"/>
 				</xsd:choice>
+				<xsd:element name="argument-resolvers" minOccurs="0">
+					<xsd:annotation>
+						<xsd:documentation><![CDATA[
+	Configures HandlerMethodArgumentResolver types to support custom controller method argument types.
+	Using this option does not override the built-in support for resolving handler method arguments.
+	To customize the built-in support for argument resolution configure SimpAnnotationMethodMessageHandler directly.
+						]]></xsd:documentation>
+					</xsd:annotation>
+					<xsd:complexType>
+						<xsd:choice minOccurs="1" maxOccurs="unbounded">
+							<xsd:element ref="beans:bean" minOccurs="0" maxOccurs="unbounded">
+								<xsd:annotation>
+									<xsd:documentation><![CDATA[
+	The HandlerMethodArgumentResolver (or WebArgumentResolver for backwards compatibility) bean definition.
+									]]></xsd:documentation>
+								</xsd:annotation>
+							</xsd:element>
+							<xsd:element ref="beans:ref" minOccurs="0" maxOccurs="unbounded">
+								<xsd:annotation>
+									<xsd:documentation><![CDATA[
+	A reference to a HandlerMethodArgumentResolver bean definition.
+									]]></xsd:documentation>
+									<xsd:appinfo>
+										<tool:annotation kind="ref">
+											<tool:expected-type type="java:org.springframework.messaging.handler.invocation.HandlerMethodArgumentResolver" />
+										</tool:annotation>
+									</xsd:appinfo>
+								</xsd:annotation>
+							</xsd:element>
+						</xsd:choice>
+					</xsd:complexType>
+				</xsd:element>
+				<xsd:element name="return-value-handlers" minOccurs="0">
+					<xsd:annotation>
+						<xsd:documentation><![CDATA[
+	Configures HandlerMethodReturnValueHandler types to support custom controller method return value handling.
+	Using this option does not override the built-in support for handling return values.
+	To customize the built-in support for handling return values configure SimpAnnotationMethodMessageHandler directly.
+						]]></xsd:documentation>
+					</xsd:annotation>
+					<xsd:complexType>
+						<xsd:choice minOccurs="1" maxOccurs="unbounded">
+							<xsd:element ref="beans:bean" minOccurs="0" maxOccurs="unbounded">
+								<xsd:annotation>
+									<xsd:documentation><![CDATA[
+	The HandlerMethodReturnValueHandler bean definition.
+									]]></xsd:documentation>
+								</xsd:annotation>
+							</xsd:element>
+							<xsd:element ref="beans:ref" minOccurs="0" maxOccurs="unbounded">
+								<xsd:annotation>
+									<xsd:documentation><![CDATA[
+	A reference to a HandlerMethodReturnValueHandler bean definition.
+									]]></xsd:documentation>
+									<xsd:appinfo>
+										<tool:annotation kind="ref">
+											<tool:expected-type type="java:org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler" />
+										</tool:annotation>
+									</xsd:appinfo>
+								</xsd:annotation>
+							</xsd:element>
+						</xsd:choice>
+					</xsd:complexType>
+				</xsd:element>
 				<xsd:element name="message-converters" minOccurs="0">
 					<xsd:annotation>
 						<xsd:documentation><![CDATA[

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

@@ -32,7 +32,9 @@ import org.springframework.beans.DirectFieldAccessor;
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
 import org.springframework.beans.factory.config.CustomScopeConfigurer;
 import org.springframework.beans.factory.xml.XmlBeanDefinitionReader;
+import org.springframework.core.MethodParameter;
 import org.springframework.core.io.ClassPathResource;
+import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageHandler;
 import org.springframework.messaging.converter.ByteArrayMessageConverter;
 import org.springframework.messaging.converter.CompositeMessageConverter;
@@ -41,6 +43,8 @@ import org.springframework.messaging.converter.DefaultContentTypeResolver;
 import org.springframework.messaging.converter.MappingJackson2MessageConverter;
 import org.springframework.messaging.converter.MessageConverter;
 import org.springframework.messaging.converter.StringMessageConverter;
+import org.springframework.messaging.handler.invocation.HandlerMethodArgumentResolver;
+import org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler;
 import org.springframework.messaging.simp.SimpMessagingTemplate;
 import org.springframework.messaging.simp.annotation.support.SimpAnnotationMethodMessageHandler;
 import org.springframework.messaging.simp.broker.SimpleBrokerMessageHandler;
@@ -336,6 +340,23 @@ public class MessageBrokerBeanDefinitionParserTests {
 		assertFalse(this.appContext.containsBean("brokerChannelExecutor"));
 	}
 
+	@Test
+	public void customArgumentAndReturnValueTypes() {
+		loadBeanDefinitions("websocket-config-broker-custom-argument-and-return-value-types.xml");
+
+		SimpAnnotationMethodMessageHandler handler = this.appContext.getBean(SimpAnnotationMethodMessageHandler.class);
+
+		List<HandlerMethodArgumentResolver> customResolvers = handler.getCustomArgumentResolvers();
+		assertEquals(2, customResolvers.size());
+		assertTrue(handler.getArgumentResolvers().contains(customResolvers.get(0)));
+		assertTrue(handler.getArgumentResolvers().contains(customResolvers.get(1)));
+
+		List<HandlerMethodReturnValueHandler> customHandlers = handler.getCustomReturnValueHandlers();
+		assertEquals(2, customHandlers.size());
+		assertTrue(handler.getReturnValueHandlers().contains(customHandlers.get(0)));
+		assertTrue(handler.getReturnValueHandlers().contains(customHandlers.get(1)));
+	}
+
 	@Test
 	public void messageConverters() {
 		loadBeanDefinitions("websocket-config-broker-converters.xml");
@@ -396,3 +417,30 @@ public class MessageBrokerBeanDefinitionParserTests {
 	}
 
 }
+
+class CustomArgumentResolver implements HandlerMethodArgumentResolver {
+
+	@Override
+	public boolean supportsParameter(MethodParameter parameter) {
+		return false;
+	}
+
+	@Override
+	public Object resolveArgument(MethodParameter parameter, Message<?> message) throws Exception {
+		return null;
+	}
+
+}
+
+class CustomReturnValueHandler implements HandlerMethodReturnValueHandler {
+
+	@Override
+	public boolean supportsReturnType(MethodParameter returnType) {
+		return false;
+	}
+
+	@Override
+	public void handleReturnValue(Object returnValue, MethodParameter returnType, Message<?> message) throws Exception {
+
+	}
+}

spring-websocket/src/test/resources/org/springframework/web/socket/config/websocket-config-broker-custom-argument-and-return-value-types.xml

@@ -0,0 +1,30 @@
+<beans xmlns="http://www.springframework.org/schema/beans"
+	   xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+	   xmlns:websocket="http://www.springframework.org/schema/websocket"
+	   xsi:schemaLocation="http://www.springframework.org/schema/beans http://www.springframework.org/schema/beans/spring-beans.xsd
+		http://www.springframework.org/schema/websocket http://www.springframework.org/schema/websocket/spring-websocket.xsd">
+
+	<websocket:message-broker application-destination-prefix="/app" user-destination-prefix="/personal">
+		<websocket:stomp-endpoint path="/foo,/bar">
+			<websocket:handshake-handler ref="myHandler"/>
+		</websocket:stomp-endpoint>
+		<websocket:simple-broker prefix="/topic"/>
+		<websocket:argument-resolvers>
+			<bean class="org.springframework.web.socket.config.CustomArgumentResolver" />
+			<ref bean="myArgumentResolver" />
+		</websocket:argument-resolvers>
+		<websocket:return-value-handlers>
+			<bean class="org.springframework.web.socket.config.CustomReturnValueHandler" />
+			<ref bean="myReturnValueHandler" />
+		</websocket:return-value-handlers>
+	</websocket:message-broker>
+
+	<bean id="myHandler" class="org.springframework.web.socket.config.TestHandshakeHandler"/>
+
+	<bean id="myInterceptor" class="org.springframework.web.socket.config.TestChannelInterceptor"/>
+
+	<bean id="myArgumentResolver" class="org.springframework.web.socket.config.CustomArgumentResolver" />
+
+	<bean id="myReturnValueHandler" class="org.springframework.web.socket.config.CustomReturnValueHandler" />
+
+</beans>