Why is SimpUserRegistry not working properly on EC2 Instance

I was able to track the issue after some debugging by adding a few logger statements in the StompSubProtocolHandler. After finding the cause, conclusion was that a channel-interceptor is not a correct place to authenticate an user. At least for my use-case.

Following are some of the code snippets from StompSubProtocolHandler -

The handleMessageFromClient method adds the user to the stompAuthentications map and publishes a SessionConnectEvent event -

public void handleMessageFromClient(WebSocketSession session, WebSocketMessage<?> webSocketMessage, MessageChannel outputChannel) {
    //...
    SimpAttributesContextHolder.setAttributesFromMessage(message);
    boolean sent = outputChannel.send(message);

    if (sent) {
        if (isConnect) {
            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)));
            }
    //...

And the handleMessageToClient retrieves the user from the stompAuthentications map and publishes a SessionConnectedEvent -

public void handleMessageToClient(WebSocketSession session, Message<?> message) {
    //...
    SimpAttributes simpAttributes = new SimpAttributes(session.getId(), session.getAttributes());
    SimpAttributesContextHolder.setAttributes(simpAttributes);
    Principal user = getUser(session);
    publishEvent(new SessionConnectedEvent(this, (Message<byte[]>) message, user));
    //...

getUser method used by above methods -

private Principal getUser(WebSocketSession session) {
    Principal user = this.stompAuthentications.get(session.getId());
    return user != null ? user : session.getPrincipal();
}

Now, the problem occurs when the handleMessageToClient snippet executes before the handleMessageFromClient snippet. In this case, user is never added to the DefaultSimpUserRegistry, as it only checks the SessionConnectedEvent.

Below is the event listener snippet from DefaultSimpUserRegistry -

public void onApplicationEvent(ApplicationEvent event) {
    //...
    else if (event instanceof SessionConnectedEvent) {
        Principal user = subProtocolEvent.getUser();
        if (user == null) {
            return;
        }
    //...

Solution

The solution is to extend DefaultHandshakeHandler and override determineUser method, which is based on this answer. But, as I am using SockJS, this requires the client to send auth-token as a query parameter. And the reason for the query parameter requirement is discussed here.