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

Allow athentication at the STOMP level

Browse Source 
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
This commit is contained in:
Rossen Stoyanchev
2016-12-09 17:39:15 -05:00
parent b14d189db9
commit e16d753439
3 changed files with 209 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -105,6 +105,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;
@@ -272,11 +274,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();
@@ -286,7 +287,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())) {
@@ -297,15 +299,23 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
SimpAttributesContextHolder.setAttributesFromMessage(message);
boolean sent = outputChannel.send(message);
if (sent && this.eventPublisher != null) {
if (StompCommand.CONNECT.equals(headerAccessor.getCommand())) {
publishEvent(new SessionConnectEvent(this, message, user));
if (sent) {
if (isConnect) {
Principal user = headerAccessor.getUser();
if (user != null && user != session.getPrincipal()) {
this.stompAuthentications.put(session.getId(), user);
}
}
else if (StompCommand.SUBSCRIBE.equals(headerAccessor.getCommand())) {
publishEvent(new SessionSubscribeEvent(this, message, user));
}
else if (StompCommand.UNSUBSCRIBE.equals(headerAccessor.getCommand())) {
publishEvent(new SessionUnsubscribeEvent(this, message, 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, getUser(session)));
}
else if (StompCommand.UNSUBSCRIBE.equals(headerAccessor.getCommand())) {
publishEvent(new SessionUnsubscribeEvent(this, message, getUser(session)));
}
}
}
}
@@ -323,6 +333,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();
}
@SuppressWarnings("deprecation")
private void handleError(WebSocketSession session, Throwable ex, Message<byte[]> clientMessage) {
if (getErrorHandler() == null) {
@@ -425,7 +440,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 {
@@ -566,7 +581,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());
@@ -613,7 +628,7 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
public void afterSessionEnded(WebSocketSession session, CloseStatus closeStatus, MessageChannel outputChannel) {
this.decoders.remove(session.getId());
Principal principal = session.getPrincipal();
Principal principal = getUser(session);
if (principal != null && this.userSessionRegistry != null) {
String userName = getSessionRegistryUserName(principal);
this.userSessionRegistry.unregisterSessionId(userName, session.getId());
@@ -624,12 +639,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();
}
@@ -642,7 +658,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());
}

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

@@ -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;
@@ -68,7 +71,7 @@ import static org.mockito.Mockito.*;
*/
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 +213,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 +359,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 +533,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;
}
}
}