Disconnect Websocket on server side based on a token passed from UI

Question

My aim is to connect browser clients having proper headers with the server. I pass these headers from StompClient.

My UI code in which i passed token in the header is

function connect() {
    var socket = new SockJS('/websocket/api/add');
    stompClient = Stomp.over(socket);
    stompClient.connect({"token" : "12345"}, function(frame) {
        setConnected(true);
        console.log('Connected: ' + frame);                
    });
}

In backend i am able to read the headers in the preSend() method of ChannelInterceptorAdapter

@Override
public Message<?> preSend(Message<?> message, MessageChannel channel) {
    MessageHeaders headers = message.getHeaders();
    System.out.println("preSend : HEADERS : {}" + headers);
    return super.preSend(message, channel);
}

But here i am not able to close the wesocket session. How can we do that?

Also i was able to close the websocket session but i couldn't receive the headers in afterConnectionEstablished() method of WebSocketHandlerDecorator

public void configureWebSocketTransport(final WebSocketTransportRegistration registration) {
    registration.addDecoratorFactory(new WebSocketHandlerDecoratorFactory() {
        @Override
        public WebSocketHandler decorate(final WebSocketHandler handler) {
            return new WebSocketHandlerDecorator(handler) {
                @Override
                public void afterConnectionEstablished(final WebSocketSession session) throws Exception {   

                    session.close(CloseStatus.NOT_ACCEPTABLE);
                    super.afterConnectionEstablished(session);
                }
            };
        }
    });
    super.configureWebSocketTransport(registration);
}

Can someone guide me how can i close the websocketsession based on the header we pass from UI at server side?

Answer 1

You can try sending the client's token to the server through a message when connection established, then let the server save that session into a map, whose key is the corresponding token.

So when you want to close a session by its token, you can query for the session from that map using the token.

---

Sample Code:

Save the session with its token:

@Override
public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
    String messageToString = message.getPayload().toString();
    if (messageToString.startsWith("token=")) {
        tokenToSessionMapping.put(messageToString.substring("token=".length()));
    }
    // Other handling message code...
}

Close the session by token:

WebSocketSession sessionByToken = tokenToSessionMapping.get(token);
if (sessionByToken != null && sessionByToken.isOpen()) {
    sessionByToken.close(CloseStatus.NOT_ACCEPTABLE);
}

And other things to notice:

1. Since the tokenToSessionMapping is static and is shared among sessions. You should use a thread-safe implementation such as ConcurrentHashMap.

2. When the session is closed, you'd better remove the corresponding entry from the map tokenToSessionMapping. Otherwise the map size will just keep growing. You can do this by the override method afterConnectionClosed().

   @Override
   public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
       Log.info("Socket session closed: {}", status.toString());
       String foundKey = null;
       for (Map.Entry<String, String> entry : tokenToSessionMapping.entrySet()) {
           if (Objects.equals(entry.getValue(), session)) {
               foundKey = entry.getKey();
           }
       }
       if (foundKey != null) {
           tokenToSessionMapping.remove(foundKey);
       }
   }