Allow athentication at the STOMP level

This commit makes it possible for a ChannelInterceptor to override the
user header in a Spring Message that contains a STOMP CONNECT frame.

After the message is sent, the updated user header is observed and
saved to be associated with session thereafter.

Issue: SPR-14690

spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java

@@ -101,6 +101,8 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 
 	private MessageHeaderInitializer headerInitializer;
 
+	private final Map<String, Principal> stompAuthentications = new ConcurrentHashMap<String, Principal>();
+
 	private Boolean immutableMessageInterceptorPresent;
 
 	private ApplicationEventPublisher eventPublisher;
@@ -247,11 +249,10 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 			try {
 				StompHeaderAccessor headerAccessor =
 						MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
-				Principal user = session.getPrincipal();
 
 				headerAccessor.setSessionId(session.getId());
 				headerAccessor.setSessionAttributes(session.getAttributes());
-				headerAccessor.setUser(user);
+				headerAccessor.setUser(getUser(session));
 				headerAccessor.setHeader(SimpMessageHeaderAccessor.HEART_BEAT_HEADER, headerAccessor.getHeartbeat());
 				if (!detectImmutableMessageInterceptor(outputChannel)) {
 					headerAccessor.setImmutable();
@@ -261,7 +262,8 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 					logger.trace("From client: " + headerAccessor.getShortLogMessage(message.getPayload()));
 				}
 
-				if (StompCommand.CONNECT.equals(headerAccessor.getCommand())) {
+				boolean isConnect = StompCommand.CONNECT.equals(headerAccessor.getCommand());
+				if (isConnect) {
 					this.stats.incrementConnectCount();
 				}
 				else if (StompCommand.DISCONNECT.equals(headerAccessor.getCommand())) {
@@ -272,15 +274,23 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 					SimpAttributesContextHolder.setAttributesFromMessage(message);
 					boolean sent = outputChannel.send(message);
 
-					if (sent && this.eventPublisher != null) {
+					if (sent) {
-						if (StompCommand.CONNECT.equals(headerAccessor.getCommand())) {
+						if (isConnect) {
-							publishEvent(new SessionConnectEvent(this, message, user));
+							Principal user = headerAccessor.getUser();
+							if (user != null && user != session.getPrincipal()) {
+								this.stompAuthentications.put(session.getId(), user);
+							}
+						}
+						if (this.eventPublisher != null) {
+							if (isConnect) {
+								publishEvent(new SessionConnectEvent(this, message, getUser(session)));
 							}
 							else if (StompCommand.SUBSCRIBE.equals(headerAccessor.getCommand())) {
-							publishEvent(new SessionSubscribeEvent(this, message, user));
+								publishEvent(new SessionSubscribeEvent(this, message, getUser(session)));
 							}
 							else if (StompCommand.UNSUBSCRIBE.equals(headerAccessor.getCommand())) {
-							publishEvent(new SessionUnsubscribeEvent(this, message, user));
+								publishEvent(new SessionUnsubscribeEvent(this, message, getUser(session)));
+							}
 						}
 					}
 				}
@@ -298,6 +308,11 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 		}
 	}
 
+	private Principal getUser(WebSocketSession session) {
+		Principal user = this.stompAuthentications.get(session.getId());
+		return user != null ? user : session.getPrincipal();
+	}
+
 	private void handleError(WebSocketSession session, Throwable ex, Message<byte[]> clientMessage) {
 		if (getErrorHandler() == null) {
 			sendErrorMessage(session, ex);
@@ -395,7 +410,7 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 				try {
 					SimpAttributes simpAttributes = new SimpAttributes(session.getId(), session.getAttributes());
 					SimpAttributesContextHolder.setAttributes(simpAttributes);
-					Principal user = session.getPrincipal();
+					Principal user = getUser(session);
 					publishEvent(new SessionConnectedEvent(this, (Message<byte[]>) message, user));
 				}
 				finally {
@@ -535,7 +550,7 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 	private StompHeaderAccessor afterStompSessionConnected(Message<?> message, StompHeaderAccessor accessor,
 			WebSocketSession session) {
 
-		Principal principal = session.getPrincipal();
+		Principal principal = getUser(session);
 		if (principal != null) {
 			accessor = toMutableAccessor(accessor, message);
 			accessor.setNativeHeader(CONNECTED_USER_HEADER, principal.getName());
@@ -574,12 +589,13 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 		try {
 			SimpAttributesContextHolder.setAttributes(simpAttributes);
 			if (this.eventPublisher != null) {
-				Principal user = session.getPrincipal();
+				Principal user = getUser(session);
 				publishEvent(new SessionDisconnectEvent(this, message, session.getId(), closeStatus, user));
 			}
 			outputChannel.send(message);
 		}
 		finally {
+			this.stompAuthentications.remove(session.getId());
 			SimpAttributesContextHolder.resetAttributes();
 			simpAttributes.sessionCompleted();
 		}
@@ -592,7 +608,7 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 		}
 		headerAccessor.setSessionId(session.getId());
 		headerAccessor.setSessionAttributes(session.getAttributes());
-		headerAccessor.setUser(session.getPrincipal());
+		headerAccessor.setUser(getUser(session));
 		return MessageBuilder.createMessage(EMPTY_PAYLOAD, headerAccessor.getMessageHeaders());
 	}
 

spring-websocket/src/test/java/org/springframework/web/socket/messaging/StompSubProtocolHandlerTests.java

@@ -17,6 +17,7 @@
 package org.springframework.web.socket.messaging;
 
 import java.io.IOException;
+import java.security.Principal;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -34,6 +35,8 @@ import org.springframework.context.ApplicationEventPublisher;
 import org.springframework.context.PayloadApplicationEvent;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.MessageHandler;
+import org.springframework.messaging.MessagingException;
 import org.springframework.messaging.simp.SimpAttributes;
 import org.springframework.messaging.simp.SimpAttributesContextHolder;
 import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
@@ -56,19 +59,29 @@ import org.springframework.web.socket.WebSocketMessage;
 import org.springframework.web.socket.handler.TestWebSocketSession;
 import org.springframework.web.socket.sockjs.transport.SockJsSession;
 
-import static org.hamcrest.Matchers.*;
+import static org.hamcrest.Matchers.is;
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.any;
-import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.mockito.Mockito.when;
 
 /**
  * Test fixture for {@link StompSubProtocolHandler} tests.
- *
  * @author Rossen Stoyanchev
  */
 public class StompSubProtocolHandlerTests {
 
-	public static final byte[] EMPTY_PAYLOAD = new byte[0];
+	private static final byte[] EMPTY_PAYLOAD = new byte[0];
 
 	private StompSubProtocolHandler protocolHandler;
 
@@ -210,22 +223,26 @@ public class StompSubProtocolHandlerTests {
 	public void handleMessageToClientWithHeartbeatSuppressingSockJsHeartbeat() throws IOException {
 
 		SockJsSession sockJsSession = Mockito.mock(SockJsSession.class);
+		when(sockJsSession.getId()).thenReturn("s1");
 		StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.CONNECTED);
 		accessor.setHeartbeat(0, 10);
 		Message<byte[]> message = MessageBuilder.createMessage(EMPTY_PAYLOAD, accessor.getMessageHeaders());
 		this.protocolHandler.handleMessageToClient(sockJsSession, message);
 
+		verify(sockJsSession).getId();
 		verify(sockJsSession).getPrincipal();
 		verify(sockJsSession).disableHeartbeat();
 		verify(sockJsSession).sendMessage(any(WebSocketMessage.class));
 		verifyNoMoreInteractions(sockJsSession);
 
 		sockJsSession = Mockito.mock(SockJsSession.class);
+		when(sockJsSession.getId()).thenReturn("s1");
 		accessor = StompHeaderAccessor.create(StompCommand.CONNECTED);
 		accessor.setHeartbeat(0, 0);
 		message = MessageBuilder.createMessage(EMPTY_PAYLOAD, accessor.getMessageHeaders());
 		this.protocolHandler.handleMessageToClient(sockJsSession, message);
 
+		verify(sockJsSession).getId();
 		verify(sockJsSession).getPrincipal();
 		verify(sockJsSession).sendMessage(any(WebSocketMessage.class));
 		verifyNoMoreInteractions(sockJsSession);
@@ -352,6 +369,28 @@ public class StompSubProtocolHandlerTests {
 		assertFalse(mutable.get());
 	}
 
+	@Test // SPR-14690
+	public void handleMessageFromClientWithTokenAuthentication() {
+		ExecutorSubscribableChannel channel = new ExecutorSubscribableChannel();
+		channel.addInterceptor(new AuthenticationInterceptor("__pete__@gmail.com"));
+		channel.addInterceptor(new ImmutableMessageChannelInterceptor());
+
+		TestMessageHandler messageHandler = new TestMessageHandler();
+		channel.subscribe(messageHandler);
+
+		StompSubProtocolHandler handler = new StompSubProtocolHandler();
+		handler.afterSessionStarted(this.session, channel);
+
+		TextMessage wsMessage = StompTextMessageBuilder.create(StompCommand.CONNECT).build();
+		handler.handleMessageFromClient(this.session, wsMessage, channel);
+
+		assertEquals(1, messageHandler.getMessages().size());
+		Message<?> message = messageHandler.getMessages().get(0);
+		Principal user = SimpMessageHeaderAccessor.getUser(message.getHeaders());
+		assertNotNull(user);
+		assertEquals("__pete__@gmail.com", user.getName());
+	}
+
 	@Test
 	public void handleMessageFromClientWithInvalidStompCommand() {
 
@@ -504,4 +543,34 @@ public class StompSubProtocolHandlerTests {
 		}
 	}
 
+	private static class TestMessageHandler implements MessageHandler {
+
+		private final List<Message> messages = new ArrayList<>();
+
+		public List<Message> getMessages() {
+			return this.messages;
+		}
+
+		@Override
+		public void handleMessage(Message<?> message) throws MessagingException {
+			this.messages.add(message);
+		}
+	}
+
+	private static class AuthenticationInterceptor extends ChannelInterceptorAdapter {
+
+		private final String name;
+
+
+		public AuthenticationInterceptor(String name) {
+			this.name = name;
+		}
+
+		@Override
+		public Message<?> preSend(Message<?> message, MessageChannel channel) {
+			TestPrincipal user = new TestPrincipal(name);
+			MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class).setUser(user);
+			return message;
+		}
+	}
 }

src/asciidoc/web-websocket.adoc

@@ -1647,35 +1647,127 @@ If the application prefix is set to "/app" then the foo method is effectively ma
 [[websocket-stomp-authentication]]
 === Authentication
 
-In a WebSocket-style application it is often useful to know who sent a message.
+Every STOMP over WebSocket messaging session begins with an HTTP request --
-Therefore some form of authentication is needed to establish the user identity
+that can be a request to upgrade to WebSockets (i.e. a WebSocket handshake)
-and associate it with the current session.
+or in the case of SockJS fallbacks a series of SockJS HTTP transport requests.
 
-Existing Web applications already use HTTP based authentication.
+Web applications already have authentication and authorization in place to
-For example Spring Security can secure the HTTP URLs of the application as usual.
+secure HTTP requests. Typically a user is authenticated via Spring Security
-Since a WebSocket session begins with an HTTP handshake, that means URLs mapped
+using some mechanism such as a login page, HTTP basic authentication, or other.
-to STOMP/WebSocket are already automatically protected and require authentication.
+The security context for the authenticated user is saved in the HTTP session
-Moreover the page that opens the WebSocket connection is itself likely protected
+and is associated with subsequent requests in the same cookie-based session.
-and so by the time of the actual handshake, the user should have been authenticated.
 
-When a WebSocket handshake is made and a new WebSocket session is created,
+Therefore for a WebSocket handshake, or for SockJS HTTP transport requests,
-Spring's WebSocket support automatically propagates the `java.security.Principal`
+typically there will already be an authenticated user accessible via
-from the HTTP request to the WebSocket session. After that every message flowing
+`HttpServletRequest#getUserPrincipal()`. Spring automatically associates that user
-through the application on that WebSocket session is enriched with
+with a WebSocket or SockJS session created for them and subsequently with all
-the user information. It's present in the message as a header.
+STOMP messages transported over that session through a user header.
-Controller methods can access the current user by adding a method argument of
-type `javax.security.Principal`.
 
-Note that even though the STOMP `CONNECT` frame has "login" and "passcode" headers
+In short there is nothing special a typical web application needs to do above
-that can be used for authentication, Spring's STOMP WebSocket support ignores them
+and beyond what it already does for security. The user is authenticated at
-and currently expects users to have been authenticated already via HTTP.
+the HTTP request level with a security context maintained through a cookie-based
+HTTP session which is then associated with WebSocket or SockJS sessions created
+for that user and results in a user header stamped on every `Message` flowing
+through the application.
+
+Note that the STOMP protocol does have a "login" and "passcode" headers
+on the `CONNECT` frame. Those were originally designed for and are still needed
+for example for STOMP over TCP. However for STOMP over WebSocket by default
+Spring ignores authorization headers at the STOMP protocol level and assumes
+the user is already authenticated at the HTTP transport level and expects that
+the WebSocket or SockJS session contain the authenticated user.
+
+[NOTE]
+====
+Spring Security provides
+https://docs.spring.io/spring-security/site/docs/current/reference/htmlsingle/#websocket[WebSocket sub-protocol authorization]
+that uses a `ChannelInterceptor` to authorize messages based on the user header in them.
+Also Spring Session provides a
+http://docs.spring.io/spring-session/docs/current/reference/html5/#websocket[WebSocket integration]
+that ensures the user HTTP session does not expire when the WebSocket session is still active.
+====
+
+
+
+[[websocket-stomp-authentication-token-based]]
+=== Token-based Authentication
+
+https://github.com/spring-projects/spring-security-oauth[Spring Security OAuth]
+provides support for token based security including JSON Web Token (JWT).
+This can be used as the authentication mechanism in Web applications
+including STOMP over WebSocket interactions just as described in the previous
+section, i.e. maintaining identity through a cookie-based session.
+
+At the same time cookie-based sessions are not always the best fit for example
+in applications that don't wish to maintain a server-side session at all or in
+mobile applications where it's common to use headers for authentication.
+
+The https://tools.ietf.org/html/rfc6455#section-10.5[WebSocket protocol RFC 6455]
+"doesn't prescribe any particular way that servers can authenticate clients during
+the WebSocket handshake." In practice however browser clients can only use standard
+authentication headers (i.e. basic HTTP authentication) or cookies and cannot for example
+provide custom headers. Likewise the SockJS JavaScript client does not provide
+a way to send HTTP headers with SockJS transport requests, see
+https://github.com/sockjs/sockjs-client/issues/196[sockjs-client issue 196].
+Instead it does allow sending query parameters that can be used to send a token
+but that has its own drawbacks, for example as the token may be inadvertently
+logged with the URL in server logs.
+
+[NOTE]
+====
+The above limitations are for browser-based clients and do not apply to the
+Spring Java-based STOMP client which does support sending headers with both
+WebSocket and SockJS requests.
+====
+
+Therefore applications that wish to avoid the use of cookies may not have any good
+alternatives for authentication at the HTTP protocol level. Instead of using cookies
+they may prefer to authenticate with headers at the STOMP messaging protocol level
+There are 2 simple steps to doing that:
+
+1. Use the STOMP client to pass authentication header(s) at connect time.
+2. Process the authentication header(s) with a `ChannelInterceptor`.
+
+Below is the example server-side configuration to register a custom authentication
+interceptor. Note that an interceptor only needs to authenticate and set
+the user header on the CONNECT `Message`. Spring will note and save the authenticated
+user and associate it with subsequent STOMP messages on the same session:
+
+[source,java,indent=0]
+[subs="verbatim,quotes"]
+----
+  @Configuration
+  @EnableWebSocketMessageBroker
+  public class MyConfig extends AbstractWebSocketMessageBrokerConfigurer {
+
+    @Override
+    public void configureClientInboundChannel(ChannelRegistration registration) {
+      registration.setInterceptors(new ChannelInterceptorAdapter() {
+
+          @Override
+          public Message<?> preSend(Message<?> message, MessageChannel channel) {
+
+              StompHeaderAccessor accessor =
+                  MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
+
+              if (StompCommand.CONNECT.equals(accessor.getCommand())) {
+                  Principal user = ... ; // access authentication header(s)
+                  accessor.setUser(user);
+              }
+
+              return message;
+          }
+      });
+    }
+  }
+----
+
+Also note that when using Spring Security's authorization for messages, at present
+you will need to ensure that the authentication `ChannelInterceptor` config is ordered
+ahead of Spring Security's. This is best done by declaring the custom interceptor in
+its own sub-class of `AbstractWebSocketMessageBrokerConfigurer` marked with
+`@Order(Ordered.HIGHEST_PRECEDENCE + 99)`.
 
-In some cases it may be useful to assign an identity to a WebSocket session even
-when the user has not been formally authenticated. For example, a mobile app might
-assign some identity to anonymous users, perhaps based on geographical location.
-The do that currently, an application can sub-class `DefaultHandshakeHandler`
-and override the `determineUser` method. The custom handshake handler can then
-be plugged in (see examples in <<websocket-server-deployment>>).