diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractMessageBrokerConfiguration.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractMessageBrokerConfiguration.java index e22e76d0f8..fa475b6811 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractMessageBrokerConfiguration.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractMessageBrokerConfiguration.java @@ -384,12 +384,14 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC } @Bean - @SuppressWarnings("deprecation") public SimpUserRegistry userRegistry() { return (getBrokerRegistry().getUserRegistryBroadcast() != null ? new MultiServerUserRegistry(createLocalUserRegistry()) : createLocalUserRegistry()); } + /** + * Create the user registry that provides access to the local users. + */ protected abstract SimpUserRegistry createLocalUserRegistry(); /** diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistry.java b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistry.java index 74038ae137..8239c70566 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistry.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistry.java @@ -60,8 +60,10 @@ public class WebMvcStompEndpointRegistry implements StompEndpointRegistry { private UrlPathHelper urlPathHelper; + @SuppressWarnings("deprecation") public WebMvcStompEndpointRegistry(WebSocketHandler webSocketHandler, WebSocketTransportRegistration transportRegistration, + org.springframework.messaging.simp.user.UserSessionRegistry userSessionRegistry, TaskScheduler defaultSockJsTaskScheduler) { Assert.notNull(webSocketHandler, "'webSocketHandler' is required "); @@ -78,11 +80,13 @@ public class WebMvcStompEndpointRegistry implements StompEndpointRegistry { } this.stompHandler = new StompSubProtocolHandler(); + this.stompHandler.setUserSessionRegistry(userSessionRegistry); if (transportRegistration.getMessageSizeLimit() != null) { this.stompHandler.setMessageSizeLimit(transportRegistration.getMessageSizeLimit()); } + this.sockJsScheduler = defaultSockJsTaskScheduler; } diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupport.java b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupport.java index 2421d5f560..417ad5d642 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupport.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupport.java @@ -27,7 +27,6 @@ import org.springframework.messaging.simp.config.AbstractMessageBrokerConfigurat import org.springframework.messaging.simp.stomp.StompBrokerRelayMessageHandler; import org.springframework.messaging.simp.user.SimpUserRegistry; import org.springframework.messaging.simp.user.UserSessionRegistryAdapter; -import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; import org.springframework.web.servlet.HandlerMapping; import org.springframework.web.socket.WebSocketHandler; import org.springframework.web.socket.config.WebSocketMessageBrokerStats; @@ -59,12 +58,22 @@ public abstract class WebSocketMessageBrokerConfigurationSupport extends Abstrac clientOutboundChannel(), brokerMessagingTemplate()); } + @Override + @SuppressWarnings("deprecation") + protected SimpUserRegistry createLocalUserRegistry() { + org.springframework.messaging.simp.user.UserSessionRegistry sessionRegistry = userSessionRegistry(); + if (sessionRegistry != null) { + return new UserSessionRegistryAdapter(sessionRegistry); + } + return new DefaultSimpUserRegistry(); + } + @Bean + @SuppressWarnings("deprecation") public HandlerMapping stompWebSocketHandlerMapping() { WebSocketHandler handler = decorateWebSocketHandler(subProtocolWebSocketHandler()); - WebSocketTransportRegistration transport = getTransportRegistration(); - ThreadPoolTaskScheduler scheduler = messageBrokerTaskScheduler(); - WebMvcStompEndpointRegistry registry = new WebMvcStompEndpointRegistry(handler, transport, scheduler); + WebMvcStompEndpointRegistry registry = new WebMvcStompEndpointRegistry(handler, + getTransportRegistration(), userSessionRegistry(), messageBrokerTaskScheduler()); registry.setApplicationContext(getApplicationContext()); registerStompEndpoints(registry); return registry.getHandlerMapping(); @@ -93,19 +102,6 @@ public abstract class WebSocketMessageBrokerConfigurationSupport extends Abstrac protected void configureWebSocketTransport(WebSocketTransportRegistration registry) { } - @Override - @SuppressWarnings("deprecation") - protected SimpUserRegistry createLocalUserRegistry() { - org.springframework.messaging.simp.user.UserSessionRegistry sessionRegistry = userSessionRegistry(); - if (sessionRegistry == null) { - return new DefaultSimpUserRegistry(); - } - else { - return (userSessionRegistry() instanceof SimpUserRegistry ? - (SimpUserRegistry) userSessionRegistry() : new UserSessionRegistryAdapter(sessionRegistry)); - } - } - protected abstract void registerStompEndpoints(StompEndpointRegistry registry); @Bean diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/DefaultSimpUserRegistry.java b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/DefaultSimpUserRegistry.java index 2921979a79..e25a9b8d52 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/DefaultSimpUserRegistry.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/DefaultSimpUserRegistry.java @@ -37,7 +37,9 @@ import org.springframework.messaging.support.MessageHeaderAccessor; import org.springframework.util.Assert; /** - * Default, mutable, thread-safe implementation of {@link SimpUserRegistry}. + * Default, mutable, thread-safe implementation of {@link SimpUserRegistry} that + * listens ApplicationContext events of type {@link AbstractSubProtocolEvent} to + * keep track of user presence and subscription information. * * @author Rossen Stoyanchev * @since 4.2 diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java index e9baeca42b..b8ee1d47b7 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java @@ -43,6 +43,7 @@ import org.springframework.messaging.simp.stomp.StompCommand; import org.springframework.messaging.simp.stomp.StompDecoder; import org.springframework.messaging.simp.stomp.StompEncoder; import org.springframework.messaging.simp.stomp.StompHeaderAccessor; +import org.springframework.messaging.simp.user.DestinationUserNameProvider; import org.springframework.messaging.support.AbstractMessageChannel; import org.springframework.messaging.support.ChannelInterceptor; import org.springframework.messaging.support.ImmutableMessageChannelInterceptor; @@ -93,6 +94,9 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE private int messageSizeLimit = 64 * 1024; + @SuppressWarnings("deprecation") + private org.springframework.messaging.simp.user.UserSessionRegistry userSessionRegistry; + private final StompEncoder stompEncoder = new StompEncoder(); private final StompDecoder stompDecoder = new StompDecoder(); @@ -149,6 +153,28 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE return this.messageSizeLimit; } + + /** + * Provide a registry with which to register active user session ids. + * @see org.springframework.messaging.simp.user.UserDestinationMessageHandler + * @deprecated as of 4.2 in favor of {@link DefaultSimpUserRegistry} which relies + * on the ApplicationContext events published by this class and is created via + * {@link org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurationSupport#createLocalUserRegistry + * WebSocketMessageBrokerConfigurationSupport.createLocalUserRegistry} + */ + @Deprecated + public void setUserSessionRegistry(org.springframework.messaging.simp.user.UserSessionRegistry registry) { + this.userSessionRegistry = registry; + } + + /** + * @deprecated as of 4.2 + */ + @Deprecated + public org.springframework.messaging.simp.user.UserSessionRegistry getUserSessionRegistry() { + return this.userSessionRegistry; + } + /** * Configure a {@link MessageHeaderInitializer} to apply to the headers of all * messages created from decoded STOMP frames and other messages sent to the @@ -506,6 +532,10 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE if (principal != null) { accessor = toMutableAccessor(accessor, message); accessor.setNativeHeader(CONNECTED_USER_HEADER, principal.getName()); + if (this.userSessionRegistry != null) { + String userName = getSessionRegistryUserName(principal); + this.userSessionRegistry.registerSessionId(userName, session.getId()); + } } long[] heartbeat = accessor.getHeartbeat(); if (heartbeat[1] > 0) { @@ -517,6 +547,14 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE return accessor; } + private String getSessionRegistryUserName(Principal principal) { + String userName = principal.getName(); + if (principal instanceof DestinationUserNameProvider) { + userName = ((DestinationUserNameProvider) principal).getDestinationUserName(); + } + return userName; + } + @Override public String resolveSessionId(Message message) { return SimpMessageHeaderAccessor.getSessionId(message.getHeaders()); @@ -533,6 +571,11 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE @Override public void afterSessionEnded(WebSocketSession session, CloseStatus closeStatus, MessageChannel outputChannel) { this.decoders.remove(session.getId()); + Principal principal = session.getPrincipal(); + if (principal != null && this.userSessionRegistry != null) { + String userName = getSessionRegistryUserName(principal); + this.userSessionRegistry.unregisterSessionId(userName, session.getId()); + } Message message = createDisconnectMessage(session); SimpAttributes simpAttributes = SimpAttributes.fromMessage(message); try { diff --git a/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistryTests.java b/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistryTests.java index 8651638405..25fa934447 100644 --- a/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistryTests.java +++ b/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistryTests.java @@ -52,7 +52,7 @@ public class WebMvcStompEndpointRegistryTests { WebSocketTransportRegistration transport = new WebSocketTransportRegistration(); TaskScheduler scheduler = Mockito.mock(TaskScheduler.class); - this.endpointRegistry = new WebMvcStompEndpointRegistry(this.webSocketHandler, transport, scheduler); + this.endpointRegistry = new WebMvcStompEndpointRegistry(this.webSocketHandler, transport, null, scheduler); }