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 d4411f4cc6
commit 2191d80a31
3 changed files with 222 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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 (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)));
}
}
}
}
@@ -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());
}

79
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;
@@ -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.junit.Assert.*;
import static org.hamcrest.Matchers.is;
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;
}
}
}