How to retrieve the session ID from Spring Websocket on initial handshake and to be used later on?

Question

How to retrieve the session ID from Spring Websocket on initial handshake and to be used later on?

Goal The Goal is whenever a Websocket client connects to the websocket server. The websocket client will pass in a parameter. It will pass in their tenantID. Upon successfull connection to the websocket server. The websocket server usually generates a session ID. What I hope to achieve is associate this websocket sessionID to that specific tenant parameter. And later on, whenever a websocket server will send an update to the that client/tenant. it will send it to that specific tenant via its sessionID that the websocket server has created.

Here is my websocket server configuration..

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {

    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        config.enableSimpleBroker("/queue");
        config.setApplicationDestinationPrefixes("/app");
    }


    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        //registry.addEndpoint("/client").addInterceptors(new WSHandshakeInterceptor()).withSockJS();
        registry.addEndpoint("/agent").addInterceptors(new WSHandshakeInterceptor()).withSockJS();
    }
}

WSHandhsakeInterceptor.java

public class WSHandshakeInterceptor implements HandshakeInterceptor {
    @Override
    public boolean beforeHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Map<String, Object> map) throws Exception {
        return true;
    }


    @Override
    public void afterHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Exception e) {
        HttpHeaders header = serverHttpRequest.getHeaders();
        String client = header.get("client-id").get(0);
        String sessionId = null;
        if (serverHttpRequest instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) serverHttpRequest;
            HttpSession session = servletRequest.getServletRequest().getSession();
            sessionId = session.getId();
            System.out.println("Session ID + "+sessionId);
            System.out.println("CLIENT ID "+client);
            ClassChangeNotificationServiceImpl.clientSessionMap.put(client, sessionId);
        }
    }
}

Here is how my websocket client connects to the server.

WebSocketClient simpleWebSocketClient = new StandardWebSocketClient();
        List<Transport> transports = new ArrayList<>(1);
        transports.add(new WebSocketTransport(simpleWebSocketClient));
        SockJsClient sockJsClient = new SockJsClient(transports);
        stompClient = new WebSocketStompClient(sockJsClient);
        stompClient.setMessageConverter(new MappingJackson2MessageConverter());
        String url = "ws://localhost:8081/agent";
        sessionHandler = new MyStompSessionHandler();
        try {
            WebSocketHttpHeaders wsHTTPheaders = new WebSocketHttpHeaders();
            wsHTTPheaders.add("CLIENT-ID","XXXTESTCLIENTXXX");
            session = stompClient.connect(url, wsHTTPheaders,sessionHandler).get();
        } catch (InterruptedException e) {
            e.printStackTrace();
        } catch (ExecutionException e) {
            e.printStackTrace();
        }

MyStompSessionHandler.java

import org.springframework.messaging.simp.stomp.StompFrameHandler;
import org.springframework.messaging.simp.stomp.StompHeaders;
import org.springframework.messaging.simp.stomp.StompSession;
import org.springframework.messaging.simp.stomp.StompSessionHandlerAdapter;

import java.lang.reflect.Type;
import java.util.List;
import java.util.Map;

public class  MyStompSessionHandler extends StompSessionHandlerAdapter {

    private void subscribeTopic(String topic, StompSession session) {
        session.subscribe(topic, new StompFrameHandler() {
            @Override
            public Type getPayloadType(StompHeaders headers) {

                System.out.println("HANDLER");
                return String.class;
            }

            @Override
            public void handleFrame(StompHeaders headers, Object payload) {
                System.out.println("TEST HEKHEk");
            }
        });

    }

    @Override
    public Type getPayloadType(StompHeaders headers) {
        return String.class;
    }

    /**
     * This implementation is empty.
     */
    @Override
    public void handleFrame(StompHeaders headers, Object payload) {
        String resp = (String) payload;
        System.out.println("Received responses from websocket server: "+ resp);
    }

    @Override
    public void afterConnected(StompSession session, StompHeaders connectedHeaders) {
        subscribeTopic("/user/queue/response", session);

        System.out.println("CONNECTEd");
        for(Map.Entry<String, List<String>> entry: connectedHeaders.entrySet()) {
            System.out.println(entry.getKey()+"S");
        }
    }
}

Now For the fun part. or what I've been trying to achieve. is I can send an update to the websocket client via this code

this.template.convertAndSendToUser(client, "/queue/reply", message);

Where client is the sessionID that spring has generated. and it will send a message to that specific client only

Question How do I retrieve the sessionID after a successful handshake?? can I use this sessionID to send specific updates to that client. Here's my code..

Answer 1

Just add this method anywhere in your code to get the session ID after handshake:

@EventListener
    private void handleSessionConnected(SessionConnectEvent event) {
        String sessionId = SimpAttributesContextHolder.currentAttributes().getSessionId();
    }