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