Fix issue with clearing subscriptions on disconnect

Author: rstoyanchev

spring-messaging/src/main/java/org/springframework/messaging/simp/handler/SimpleBrokerMessageHandler.java

@@ -56,7 +56,8 @@ public void setSubscriptionRegistry(SubscriptionRegistry subscriptionRegistry) {
 
 	@Override
 	protected Collection<SimpMessageType> getSupportedMessageTypes() {
-		return Arrays.asList(SimpMessageType.MESSAGE, SimpMessageType.SUBSCRIBE, SimpMessageType.UNSUBSCRIBE);
+		return Arrays.asList(SimpMessageType.MESSAGE, SimpMessageType.SUBSCRIBE,
+				SimpMessageType.UNSUBSCRIBE, SimpMessageType.DISCONNECT);
 	}
 
 	@Override

spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompWebSocketHandler.java

@@ -47,20 +47,20 @@ public class StompWebSocketHandler extends TextWebSocketHandlerAdapter implement
 
 	private static Log logger = LogFactory.getLog(StompWebSocketHandler.class);
 
-	private MessageChannel outputChannel;
+	private MessageChannel clientInputChannel;
 
 	private final StompMessageConverter stompMessageConverter = new StompMessageConverter();
 
 	private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<String, WebSocketSession>();
 
 
 	/**
-	 * @param outputChannel the channel to which incoming STOMP/WebSocket messages should
+	 * @param clientInputChannel the channel to which incoming STOMP/WebSocket messages should
 	 *        be sent to
 	 */
-	public StompWebSocketHandler(MessageChannel outputChannel) {
-		Assert.notNull(outputChannel, "clientInputChannel is required");
-		this.outputChannel = outputChannel;
+	public StompWebSocketHandler(MessageChannel clientInputChannel) {
+		Assert.notNull(clientInputChannel, "clientInputChannel is required");
+		this.clientInputChannel = clientInputChannel;
 	}
 
 
@@ -70,7 +70,7 @@ public StompMessageConverter getStompMessageConverter() {
 
 	@Override
 	public void afterConnectionEstablished(WebSocketSession session) throws Exception {
-		Assert.notNull(this.outputChannel, "No output channel for STOMP messages.");
+		Assert.notNull(this.clientInputChannel, "No output channel for STOMP messages.");
 		this.sessions.put(session.getId(), session);
 	}
 
@@ -102,7 +102,7 @@ protected void handleTextMessage(WebSocketSession session, TextMessage textMessa
 				}
 
 				message = MessageBuilder.fromMessage(message).copyHeaders(headers.toMap()).build();
-				this.outputChannel.send(message);
+				this.clientInputChannel.send(message);
 			}
 			catch (Throwable t) {
 				logger.error("Terminating STOMP session due to failure to send message: ", t);
@@ -170,7 +170,7 @@ public void afterConnectionClosed(WebSocketSession session, CloseStatus status)
 		SimpMessageHeaderAccessor headers = SimpMessageHeaderAccessor.create(SimpMessageType.DISCONNECT);
 		headers.setSessionId(session.getId());
 		Message<?> message = MessageBuilder.withPayload(new byte[0]).copyHeaders(headers.toMap()).build();
-		this.outputChannel.send(message);
+		this.clientInputChannel.send(message);
 	}
 
 	/**

spring-messaging/src/test/java/org/springframework/messaging/simp/handler/SimpleBrokerWebMessageHandlerTests.java

@@ -28,7 +28,6 @@
 import org.springframework.messaging.MessageChannel;
 import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
 import org.springframework.messaging.simp.SimpMessageType;
-import org.springframework.messaging.simp.handler.SimpleBrokerMessageHandler;
 import org.springframework.messaging.support.MessageBuilder;
 
 import static org.junit.Assert.*;
@@ -60,7 +59,8 @@ public void setup() {
 
 	@Test
 	public void getSupportedMessageTypes() {
-		assertEquals(Arrays.asList(SimpMessageType.MESSAGE, SimpMessageType.SUBSCRIBE, SimpMessageType.UNSUBSCRIBE),
+		assertEquals(Arrays.asList(SimpMessageType.MESSAGE, SimpMessageType.SUBSCRIBE,
+				SimpMessageType.UNSUBSCRIBE, SimpMessageType.DISCONNECT),
 				this.messageHandler.getSupportedMessageTypes());
 	}