Maintain Spring Session during WebSocket Connection

In this post, we will be discussing about how to maintain spring session during websocket connection through HandshakeInterceptor. Doing so we can track user session during every websocket request and can utilize this session to track client activities from the server or can be used to provide an extra level of security even after the server is connected through websocket protocol. Now let us look into the code how can we maintain spring session during websocket connection in Spring boot.

The question comes why we sometimes require to maintain session during a websocket connection. Actually, we can't send any HTTP headers using stompClient, because STOMP is layer over websockets, and only when websockets handshake happen we have possibility to send custom headers. Hence, if we require to add any extra validations or security checks over the data sent from the client to the server then this session could be useful.

Here, we won't be discussing about the authentication and authorization for websocket connections because with Spring 4 WebSockets reuse the same authentication information that is found in the HTTP request when the WebSocket connection was made. To ensure a user has authenticated to your WebSocket application, all that is necessary is to ensure that you setup Spring Security to authenticate your HTTP based web application.

To add authentication to your websocket connection, you can visit this nice explanation on spring official doc.

How Session is Maintained

As we know, before making a websocket connection to a server, client makes a handshake request as an upgrade request and this request is a HTTP request. Hence, to maintain a websocket session, we require to intercept this HTTP request and keep the session id somewhere from where it can be accessed everytime whenever a websocket request is made. Here we will be using STOMP header attributes to track the session.

Environment Setup

Maven Dependencies

There is no any extra maven dependency is required for this case that we used in our previous post of Spring Websocket Integration Example without STOMP. Hence let us ignore it for time being.

Server Side

Spring Boot Websocket Configuration

First of all let us configure our message broker. In WebSocketConfig.java we have defined our STOMP endpoint and message broker. The important thing to notice here is the configuration for handshakeInterceptor.So, whenever any websocket handshake request is received by the server,this class will come into picture.

The endpoint /greeting is the endpoint that client will be using to start with the handshake and that configuration is done below.

During this handshake itself, we can put any extra attribute in the session object that we require for any future use case.

package com.developertack;

import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.web.socket.config.annotation.AbstractWebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig extends AbstractWebSocketMessageBrokerConfigurer {
	
	 @Override
	 public void configureMessageBroker(MessageBrokerRegistry config) {
	  config.enableSimpleBroker("/topic/", "/queue/");
	  config.setApplicationDestinationPrefixes("/app");
	 }
	 
	 @Override
	 public void registerStompEndpoints(StompEndpointRegistry registry) {
	  registry.addEndpoint("/greeting").addInterceptors(new HttpHandshakeInterceptor());
	 }
	
}

 Other Interesting Posts

Websocket spring Boot Integration Example
Spring Boot Websocket without STOMP
Spring Security Password Encoding using Bcrypt Encoder
Spring Boot Security Redirect after Login with complete JavaConfig
Spring Security Hibernate Example with complete JavaConfig
Securing REST API with Spring Security Basic Authentication
Spring Security with Spring MVC Example Using Spring Boot
Spring JMS Activemq Integration with Spring Boot

Spring Boot Websocket Interceptor

Now let us define the httpInterceptor. We have overriden beforeHandshake() to intercept the handshake request and set our custom attribute in the session. We are setting our custom attribute sessionId with actual sessionId and this is the same key which we will be using later to track the session from our controller class.

You can use this method to store any other attribute of your wish.

package com.developertack;

import java.util.Map;

import javax.servlet.http.HttpSession;

import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

public class HttpHandshakeInterceptor implements HandshakeInterceptor {

	@Override
	public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
			Map attributes) throws Exception {
		if (request instanceof ServletServerHttpRequest) {
			ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
			HttpSession session = servletRequest.getServletRequest().getSession();
			attributes.put("sessionId", session.getId());
		}
		return true;
	}

	public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
			Exception ex) {
	}
}

Websocket Subscription Event Listener

Following is the listener class which will be executed whenever an event is raised when a new WebSocket client using a Simple Messaging Protocol (STOMP) sends a subscription request for any queue or topic. We are also validating the sessionId which we created during handshake.

package com.developertack;

import org.springframework.context.ApplicationListener;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.messaging.SessionSubscribeEvent;

@Component
public class SubscribeEventListener implements ApplicationListener {

	@Override
	public void onApplicationEvent(SessionSubscribeEvent sessionSubscribeEvent) {
		StompHeaderAccessor headerAccessor = StompHeaderAccessor.wrap(sessionSubscribeEvent.getMessage());
		System.out.println(headerAccessor.getSessionAttributes().get("sessionId").toString());
	}
}

Websocket Request Mapping

Now let us define the controller which will intercept the message coming from the stomp client and validate the session using custom attribute sessionId which we set during handshake.This attribute can be accessed from the headerAccessor.

package com.developertack.controller;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.handler.annotation.Payload;
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
import org.springframework.messaging.simp.SimpMessageSendingOperations;
import org.springframework.stereotype.Controller;

@Controller
public class WebSocketController {
	
	@Autowired
	private SimpMessageSendingOperations messagingTemplate;
	
	@MessageMapping("/message")
    public void processMessageFromClient(@Payload String message, SimpMessageHeaderAccessor  headerAccessor) throws Exception {
		String sessionId = headerAccessor.getSessionAttributes().get("sessionId").toString();
		System.out.println(sessionId);
		headerAccessor.setSessionId(sessionId);
		messagingTemplate.convertAndSend("/topic/reply", new Gson().fromJson(message, Map.class).get("name"));
      
    }

}

Client Side Implementation

At the client side, we have a .html file to trigger the websocket connection over a button click. The text entered in the text box will be sent to server over websocket protocol and the response from the server will be diaplayed in the Greetings div element.

<!DOCTYPE html>
<html>
<head>
    <title>Hello WebSocket</title>
    <link href="/bootstrap.min.css" rel="stylesheet">
    <link href="/main.css" rel="stylesheet">
    <script src="/jquery-1.10.2.min.js"></script>
    <script src="/app.js"></script>
    <script src="/stomp.js"></script>
    
</head>
<body>
<noscript><h2 style="color: #ff0000">Seems your browser doesn't support Javascript! Websocket relies on Javascript being
    enabled. Please enable
    Javascript and reload this page!</h2></noscript>
<div id="main-content" class="container">
    <div class="row">
        <div class="col-md-6">
            <form class="form-inline">
                <div class="form-group">
                    <label for="connect">WebSocket connection:</label>
                    <button id="connect" class="btn btn-default" type="submit">Connect</button>
                    <button id="disconnect" class="btn btn-default" type="submit" disabled="disabled">Disconnect
                    </button>
                </div>
            </form>
        </div>
        <div class="col-md-6">
            <form class="form-inline">
                <div class="form-group">
                    <label for="name">What is your name?</label>
                    <input type="text" id="name" class="form-control" placeholder="Your name here...">
                </div>
                <button id="send" class="btn btn-default" type="submit">Send</button>
            </form>
        </div>
    </div>
    <div class="row">
        <div class="col-md-12">
            <table id="conversation" class="table table-striped">
                <thead>
                <tr>
                    <th>Greetings</th>
                </tr>
                </thead>
                <tbody id="greetings">
                </tbody>
            </table>
        </div>
    </div>
 
</div>
</body>
</html>

Now let us define our app.js. In app.js, we have connect() method that is triggered on the button click. The connection is requested at /greeting endpoint. This is the same endpoint that we configured as STOMP endpoint in our server side websocket configuration.

We also have the configuration to subscribe at /topic/reply and whenever a new message is recieved, the showGreeting() method will be invoked.

var ws;
function setConnected(connected) {
    $("#connect").prop("disabled", connected);
    $("#disconnect").prop("disabled", !connected);
    if (connected) {
        $("#conversation").show();
    }
    else {
        $("#conversation").hide();
    }
    $("#greetings").html("");
}

function connect() {
	//connect to stomp where stomp endpoint is exposed
	var socket = new WebSocket("ws://localhost:8080/greeting");
	ws = Stomp.over(socket);

	//usernaeme and password
	ws.connect({}, function(frame){
	
	  ws.subscribe("/topic/reply", function(msg) {
	    alert(msg);
	    showGreeting(msg.body);
	  });

	 setConnected(true);
	});
}

function disconnect() {
    if (ws != null) {
        ws.close();
    }
    setConnected(false);
    console.log("Disconnected");
}

function sendName() {
	var data = JSON.stringify({'name': $("#name").val()})
	 ws.send("/app/message", {},data);
}

function showGreeting(message) {
    $("#greetings").append(" " + message + "");
}

$(function () {
    $("form").on('submit', function (e) {
        e.preventDefault();
    });
    $( "#connect" ).click(function() { connect(); });
    $( "#disconnect" ).click(function() { disconnect(); });
    $( "#send" ).click(function() { sendName(); });
});

Run Application

1. Run Application.java as a java application.

2. Open the browser and hit the URL - http:localhost:8080.

3. Click on connect button to make the socket connection.

4. Enter your name and click send.Following is the result.

On the server side you can check the console and find the sessionIds are same.

Conclusion

I hope this article served you that you were looking for. If you have anything that you want to add or share then please share it below in the comment section. The source code can be downloaded from te github here.