Spring/spring-framework
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Allow plugging in a WebSocketHandlerDecorator

Browse Source 
The WebSocketMessageBroker config now allows wrapping the
SubProtocolWebSocketHandler to enable advanced use cases that may
require access to the underlying WebSocketSession.

Issue: SPR-12314
This commit is contained in:
Rossen Stoyanchev
2014-10-13 09:35:51 -04:00
parent f0323be786
commit 97596fb9f6
8 changed files with 283 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -62,7 +62,10 @@ import org.springframework.web.context.support.GenericWebApplicationContext;
import org.springframework.web.servlet.HandlerMapping;
import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TestWebSocketSession;
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.SubProtocolWebSocketHandler;
import org.springframework.web.socket.server.HandshakeHandler;
@@ -93,7 +96,7 @@ public class MessageBrokerBeanDefinitionParserTests {
@Test
public void simpleBroker() {
public void simpleBroker() throws Exception {
loadBeanDefinitions("websocket-config-broker-simple.xml");
HandlerMapping hm = this.appContext.getBean(HandlerMapping.class);
@@ -113,6 +116,10 @@ public class MessageBrokerBeanDefinitionParserTests {
List<HandshakeInterceptor> interceptors = wsHttpRequestHandler.getHandshakeInterceptors();
assertThat(interceptors, contains(instanceOf(FooTestInterceptor.class), instanceOf(BarTestInterceptor.class)));
WebSocketSession session = new TestWebSocketSession("id");
wsHttpRequestHandler.getWebSocketHandler().afterConnectionEstablished(session);
assertEquals(true, session.getAttributes().get("decorated"));
WebSocketHandler wsHandler = unwrapWebSocketHandler(wsHttpRequestHandler.getWebSocketHandler());
assertNotNull(wsHandler);
assertThat(wsHandler, Matchers.instanceOf(SubProtocolWebSocketHandler.class));
@@ -429,7 +436,6 @@ class CustomArgumentResolver implements HandlerMethodArgumentResolver {
public Object resolveArgument(MethodParameter parameter, Message<?> message) throws Exception {
return null;
}
}
class CustomReturnValueHandler implements HandlerMethodReturnValueHandler {
@@ -443,4 +449,18 @@ class CustomReturnValueHandler implements HandlerMethodReturnValueHandler {
public void handleReturnValue(Object returnValue, MethodParameter returnType, Message<?> message) throws Exception {
}
}
class TestWebSocketHandlerDecoratorFactory implements WebSocketHandlerDecoratorFactory {
@Override
public WebSocketHandler decorate(WebSocketHandler handler) {
return new WebSocketHandlerDecorator(handler) {
@Override
public void afterConnectionEstablished(WebSocketSession session) throws Exception {
session.getAttributes().put("decorated", true);
super.afterConnectionEstablished(session);
}
};
}
}

87
spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupportTests.java
View File

@@ -22,9 +22,9 @@ import java.util.Map;
import java.util.Set;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import org.junit.Before;
import org.junit.Test;
import org.springframework.context.ApplicationContext;
import org.springframework.context.annotation.AnnotationConfigApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
@@ -43,13 +43,17 @@ import org.springframework.stereotype.Controller;
import org.springframework.web.servlet.HandlerMapping;
import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.config.WebSocketMessageBrokerStats;
import org.springframework.web.socket.handler.TestWebSocketSession;
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.StompTextMessageBuilder;
import org.springframework.web.socket.messaging.SubProtocolHandler;
import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
import org.springframework.web.socket.server.support.WebSocketHttpRequestHandler;
import static org.junit.Assert.*;
import static org.junit.Assert.assertEquals;
@@ -62,19 +66,11 @@ import static org.junit.Assert.assertEquals;
*/
public class WebSocketMessageBrokerConfigurationSupportTests {
private AnnotationConfigApplicationContext config;
@Before
public void setupOnce() {
this.config = new AnnotationConfigApplicationContext();
this.config.register(TestWebSocketMessageBrokerConfiguration.class, TestSimpleMessageBrokerConfig.class);
this.config.refresh();
}
@Test
public void handlerMapping() {
SimpleUrlHandlerMapping hm = (SimpleUrlHandlerMapping) this.config.getBean(HandlerMapping.class);
ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
SimpleUrlHandlerMapping hm = (SimpleUrlHandlerMapping) config.getBean(HandlerMapping.class);
assertEquals(1, hm.getOrder());
Map<String, Object> handlerMap = hm.getHandlerMap();
@@ -84,8 +80,9 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
@Test
public void clientInboundChannelSendMessage() throws Exception {
TestChannel channel = this.config.getBean("clientInboundChannel", TestChannel.class);
SubProtocolWebSocketHandler webSocketHandler = this.config.getBean(SubProtocolWebSocketHandler.class);
ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
TestChannel channel = config.getBean("clientInboundChannel", TestChannel.class);
SubProtocolWebSocketHandler webSocketHandler = config.getBean(SubProtocolWebSocketHandler.class);
WebSocketSession session = new TestWebSocketSession("s1");
webSocketHandler.afterConnectionEstablished(session);
@@ -102,7 +99,8 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
@Test
public void clientOutboundChannelChannel() {
TestChannel channel = this.config.getBean("clientOutboundChannel", TestChannel.class);
ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
TestChannel channel = config.getBean("clientOutboundChannel", TestChannel.class);
Set<MessageHandler> handlers = channel.getSubscribers();
assertEquals(1, handlers.size());
@@ -111,8 +109,9 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
@Test
public void webSocketTransportOptions() {
ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
SubProtocolWebSocketHandler subProtocolWebSocketHandler =
this.config.getBean("subProtocolWebSocketHandler", SubProtocolWebSocketHandler.class);
config.getBean("subProtocolWebSocketHandler", SubProtocolWebSocketHandler.class);
assertEquals(1024 * 1024, subProtocolWebSocketHandler.getSendBufferSizeLimit());
assertEquals(25 * 1000, subProtocolWebSocketHandler.getSendTimeLimit());
@@ -126,8 +125,9 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
@Test
public void messageBrokerSockJsTaskScheduler() {
ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
ThreadPoolTaskScheduler taskScheduler =
this.config.getBean("messageBrokerSockJsTaskScheduler", ThreadPoolTaskScheduler.class);
config.getBean("messageBrokerSockJsTaskScheduler", ThreadPoolTaskScheduler.class);
ScheduledThreadPoolExecutor executor = taskScheduler.getScheduledThreadPoolExecutor();
assertEquals(Runtime.getRuntime().availableProcessors(), executor.getCorePoolSize());
@@ -136,8 +136,9 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
@Test
public void webSocketMessageBrokerStats() {
ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
String name = "webSocketMessageBrokerStats";
WebSocketMessageBrokerStats stats = this.config.getBean(name, WebSocketMessageBrokerStats.class);
WebSocketMessageBrokerStats stats = config.getBean(name, WebSocketMessageBrokerStats.class);
String actual = stats.toString();
String expected = "WebSocketSession\\[0 current WS\\(0\\)-HttpStream\\(0\\)-HttpPoll\\(0\\), " +
"0 total, 0 closed abnormally \\(0 connect failure, 0 send limit, 0 transport error\\)\\], " +
@@ -150,6 +151,29 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
assertTrue("\nExpected: " + expected.replace("\\", "") + "\n Actual: " + actual, actual.matches(expected));
}
@Test
public void webSocketHandlerDecorator() throws Exception {
ApplicationContext config = createConfig(WebSocketHandlerDecoratorConfig.class);
WebSocketHandler handler = config.getBean(SubProtocolWebSocketHandler.class);
assertNotNull(handler);
SimpleUrlHandlerMapping mapping = (SimpleUrlHandlerMapping) config.getBean("stompWebSocketHandlerMapping");
WebSocketHttpRequestHandler httpHandler = (WebSocketHttpRequestHandler) mapping.getHandlerMap().get("/test");
handler = httpHandler.getWebSocketHandler();
WebSocketSession session = new TestWebSocketSession("id");
handler.afterConnectionEstablished(session);
assertEquals(true, session.getAttributes().get("decorated"));
}
private ApplicationContext createConfig(Class<?>... configClasses) {
AnnotationConfigApplicationContext config = new AnnotationConfigApplicationContext();
config.register(configClasses);
config.refresh();
return config;
}
@Controller
static class TestController {
@@ -167,7 +191,7 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
}
@Configuration
static class TestSimpleMessageBrokerConfig extends AbstractWebSocketMessageBrokerConfigurer {
static class TestConfigurer extends AbstractWebSocketMessageBrokerConfigurer {
@Bean
public TestController subscriptionController() {
@@ -188,7 +212,7 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
}
@Configuration
static class TestWebSocketMessageBrokerConfiguration extends DelegatingWebSocketMessageBrokerConfiguration {
static class TestChannelConfig extends DelegatingWebSocketMessageBrokerConfiguration {
@Override
@Bean
@@ -208,6 +232,31 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
}
}
@Configuration
static class WebSocketHandlerDecoratorConfig extends WebSocketMessageBrokerConfigurationSupport {
@Override
protected void registerStompEndpoints(StompEndpointRegistry registry) {
registry.addEndpoint("/test");
}
@Override
protected void configureWebSocketTransport(WebSocketTransportRegistration registry) {
registry.addDecoratorFactory(new WebSocketHandlerDecoratorFactory() {
@Override
public WebSocketHandlerDecorator decorate(WebSocketHandler handler) {
return new WebSocketHandlerDecorator(handler) {
@Override
public void afterConnectionEstablished(WebSocketSession session) throws Exception {
session.getAttributes().put("decorated", true);
super.afterConnectionEstablished(session);
}
};
}
});
}
}
private static class TestChannel extends ExecutorSubscribableChannel {
private final List<Message<?>> messages = new ArrayList<>();