Build a Chat Application Using Spring Boot and WebSocket

In a previous tutorial we created a Spring Boot + WebSocket Application using STOMP and SockJS. But so far, what we already created in previous articles are broadcast applications. But in reality, don't want to broadcast every message to every user. In this tutorial we will learn on how to create a simple chat application, on how to use WebSocket and STOMP to send messages to a single user.

Changes in WebSocketMessageBrokerConfig

Here full code for WebSocketMessageBrokerConfig:

WebSocketMessageBrokerConfig.java
package com.dariawan.websocket.config;

import com.dariawan.websocket.util.UserInterceptor;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketMessageBrokerConfig implements WebSocketMessageBrokerConfigurer {

    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        config.enableSimpleBroker("/topic", "/queue");
        config.setApplicationDestinationPrefixes("/app");
        config.setUserDestinationPrefix("/user");  // this is optional 
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/broadcast");  // this is OK to leave it here
        // custom heartbeat, every 60 sec
        registry.addEndpoint("/broadcast").withSockJS().setHeartbeatTime(60_000);
        registry.addEndpoint("/chat").withSockJS();
    }
}
                    

We enable one more simple broker: /queue. Client will subscribe messages at endpoint, for our example: /queue/messages.

We configure /user, use it as the prefix used to identify user destinations. This is optional, as by default /user/ used by default for this destinations. This user destinations allowing a user to subscribe to queue names unique to their session, and enable others to send messages to those user-specific queues.

We also configure one more endpoint, /chat with SockJS, which enable fallback options for browsers that don’t support WebSocket.

Controllers

In this tutorial, we have two controllers. Our main controller is WebSocketChatController, which will handle all messages.

WebSocketChatController.java
package com.dariawan.websocket.controller;

import com.dariawan.websocket.dto.ChatMessage;
import com.dariawan.websocket.util.ActiveUserManager;
import com.dariawan.websocket.util.ActiveUserChangeListener;
import java.util.Set;
import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
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.SimpMessagingTemplate;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.GetMapping;

@Controller
public class WebSocketChatController implements ActiveUserChangeListener {

    // private final static Logger LOGGER = LoggerFactory.getLogger(WebSocketChatController.class);    
    
    @Autowired
    private SimpMessagingTemplate webSocket;

    @Autowired
    private ActiveUserManager activeUserManager;

    @PostConstruct
    private void init() {
        activeUserManager.registerListener(this);
    }

    @PreDestroy
    private void destroy() {
        activeUserManager.removeListener(this);
    }

    @GetMapping("/sockjs-message")
    public String getWebSocketWithSockJs() {
        return "sockjs-message";
    }

    @MessageMapping("/chat")
    public void send(SimpMessageHeaderAccessor sha, @Payload ChatMessage chatMessage) throws Exception {
        String sender = sha.getUser().getName();
        ChatMessage message = new ChatMessage(chatMessage.getFrom(), chatMessage.getText(), chatMessage.getRecipient());
        if (!sender.equals(chatMessage.getRecipient())) {
            webSocket.convertAndSendToUser(sender, "/queue/messages", message);
        }

        webSocket.convertAndSendToUser(chatMessage.getRecipient(), "/queue/messages", message);
    }

    @Override
    public void notifyActiveUserChange() {
        Set<String> activeUsers = activeUserManager.getAll();
        webSocket.convertAndSend("/topic/active", activeUsers);
    }
}
                    

Function send(SimpMessageHeaderAccessor, @Payload ChatMessage) responsible to check the sender and distribute the copy of message, one to the sender (from) and another to the recipient.

The instance (autowired) SimpMessagingTemplate provides methods for sending messages to a user. The method convertAndSendToUser(...) is used to send a message to the given user.

Class WebSocketChatController also implements interface ActiveUserChangeListener:

ActiveUserChangeListener.java
package com.dariawan.websocket.util;

/**
 * This interface used as an Observer for  ActiveUserManager class
 */
public interface ActiveUserChangeListener {

    /**
     * call when Observable's internal state is changed.
     */
    void notifyActiveUserChange();
}
                    

This interface used as an Observer to track active users change. To store the changes, we are create class ActiveUserManager:

ActiveUserManager.java
package com.dariawan.websocket.util;

import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import org.springframework.stereotype.Component;

@Component
public class ActiveUserManager {

    private final Map<String, Object> map;

    private final List<ActiveUserChangeListener> listeners;

    private final ThreadPoolExecutor notifyPool;

    private ActiveUserManager() {
        map = new ConcurrentHashMap<>();
        listeners = new CopyOnWriteArrayList<>();
        notifyPool = new ThreadPoolExecutor(1, 5, 10, TimeUnit.SECONDS, new ArrayBlockingQueue<>(100));
    }

    public void add(String userName, String remoteAddress) {
        map.put(userName, remoteAddress);
        notifyListeners();
    }

    /**
     * Removes username from the concurrentHashMap.
     *
     * @param username - to be removed
     */
    public void remove(String username) {
        map.remove(username);
        notifyListeners();
    }

    /**
     * Get all active user
     *
     * @return - Set of active users.
     */
    public Set<String> getAll() {
        return map.keySet();
    }

    /**
     * Get a set of all active user except username that passed in the parameter
     *
     * @param username - current username
     * @return - set of users except passed username
     */
    public Set<String> getActiveUsersExceptCurrentUser(String username) {
        Set<String> users = new HashSet<>(map.keySet());
        users.remove(username);
        return users;
    }
    
    /**
     * To get notified when active users changed
     *
     * @param listener - object that implements ActiveUserChangeListener
     */
    public void registerListener(ActiveUserChangeListener listener) {
        listeners.add(listener);
    }

    /**
     * Stop receiving notification.
     *
     * @param listener - object that implements ActiveUserChangeListener
     */
    public void removeListener(ActiveUserChangeListener listener) {
        listeners.remove(listener);
    }

    private void notifyListeners() {
        notifyPool.submit(() -> listeners.forEach(ActiveUserChangeListener::notifyActiveUserChange));
    }
}
                    

ActiveUserManager is maintained by another controller, which is a REST controller: WebSocketConnectionRestController. This REST controller responsible to handle user's connection, and add/remove it to/from ActiveUserManager.

WebSocketConnectionRestController.java
package com.dariawan.websocket.controller;

import com.dariawan.websocket.util.ActiveUserManager;
import java.util.Set;
import javax.servlet.http.HttpServletRequest;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.ModelAttribute;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RestController;

@RestController
public class WebSocketConnectionRestController {
    
    @Autowired
    private ActiveUserManager activeSessionManager;
    
    @PostMapping("/rest/user-connect")
    public String userConnect(HttpServletRequest request,
            @ModelAttribute("username") String userName) {
        String remoteAddr = "";
        if (request != null) {
            remoteAddr = request.getHeader("Remote_Addr");
            if (StringUtils.isEmpty(remoteAddr)) {
                remoteAddr = request.getHeader("X-FORWARDED-FOR");
                if (remoteAddr == null || "".equals(remoteAddr)) {
                    remoteAddr = request.getRemoteAddr();
                }
            }
        }
        
        activeSessionManager.add(userName, remoteAddr);
        return remoteAddr;
    }
    
    @PostMapping("/rest/user-disconnect")
    public String userDisconnect(@ModelAttribute("username") String userName) {
        activeSessionManager.remove(userName);
        return "disconnected";
    }
    
    @GetMapping("/rest/active-users-except/{userName}")
    public Set<String> getActiveUsersExceptCurrentUser(@PathVariable String userName) {
        return activeSessionManager.getActiveUsersExceptCurrentUser(userName);
    }
}
                    

There are three functions in this class:

  • userConnect(...): add new user to the pool in ActiveUserManager
  • userDisconnect(...): remove user from the pool in ActiveUserManager
  • getActiveUsersExceptCurrentUser(...): list all active user, except the one who request it (in short: get another users active)

These REST web-services will later consumed by JavaScript client.

Now, the best part of this article... It'll be broader and deeper topic if I cover how to secure WebSockets implementation in this article. So, I'll not cover that in this article. I also not include Spring Security in my example, which lead me into a specific issue:

Function send(...) in WebSocketChatController using Spring provided API convertAndSendToUser(username, destination, message) to send to specific. This API accepts a String username, so with that we will be able to send messages to specific user subscribed to a topic — if we have those user's unique username. But the question is: where does this username come from? No, this is not the username you're specified when joining the topic. If your username is "Alicia", doesn't mean that system unique username also "ALicia"

The unique username is part of a java.security.Principal interface. Each StompHeaderAccessor or WebSocketSession object has instance of this principal and you can get the user name from it. In my example, I will use the username choose by user as system unique username. To do that, first we need to have User class which implement Principal interface:

User.java
package com.dariawan.websocket.util;

import java.security.Principal;

public class User implements Principal {

    String name;

    public User(String name) {
        this.name = name;
    }

    @Override
    public String getName() {
        return name;
    }
}
                    

Then I have another class, UserInterceptor. In this class, I override preSend(...). For each StompCommand.CONNECT, I set User's name is the name that user choose (and key in in the browser), and putting it into StompHeaderAccessor.

UserInterceptor.java
package com.dariawan.websocket.util;

import java.util.LinkedList;
import java.util.Map;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;

public class UserInterceptor implements ChannelInterceptor {

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {

        StompHeaderAccessor accessor
                = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);

        if (StompCommand.CONNECT.equals(accessor.getCommand())) {
            Object raw = message
                    .getHeaders()
                    .get(SimpMessageHeaderAccessor.NATIVE_HEADERS);

            if (raw instanceof Map) {
                Object name = ((Map) raw).get("username");

                if (name instanceof LinkedList) {
                    accessor.setUser(new User(((LinkedList) name).get(0).toString()));
                }
            }
        }
        return message;
    }
}
                    

Next, we need to register this interceptor. Change and add/override configureClientInboundChannel(...) in WebSocketMessageBrokerConfig as below:

import com.dariawan.websocket.util.UserInterceptor;
...
import org.springframework.messaging.simp.config.ChannelRegistration;
...

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketMessageBrokerConfig implements WebSocketMessageBrokerConfigurer {

    ...

    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(new UserInterceptor());
    }
}
                    

That's all, and we all set for server side. Yes, is not a full proof solution, where you need to think what if the username is not unique, and another user want to join using a username that already exists in ActiveUserManager? Well, we not building a proven enterprise solution, and that's not our focus.

For the client layout, here sockjs-message.html which will be invoked by Thymeleaf as in WebSocketChatController

sockjs-message.html
<!DOCTYPE HTML>
<html xmlns:th="http://www.thymeleaf.org">
    <head>
        <title>WebSocket With STOMP & SockJS Broadcast Example</title>
        <th:block th:include="fragments/common.html :: headerfiles"></th:block>
    </head>
    <body>
        <div class="container">
            <div class="py-5 text-center">
                <a href="/"><h2>WebSocket</h2></a>
                <p class="lead">WebSocket Chat - with STOMP & SockJS.</p>
            </div>
            <div class="row">
                <div class="col-md-6">
                    <div class="mb-3">
                        <div class="input-group">
                            <input type="text" id="from" class="form-control" placeholder="Choose a nickname"/>
                            <div class="btn-group">
                                <button type="button" id="connect" class="btn btn-sm btn-outline-secondary" onclick="connect()">Connect</button>
                                <button type="button" id="disconnect" class="btn btn-sm btn-outline-secondary" onclick="disconnect()" disabled>Disconnect</button>
                            </div>                        
                        </div>
                    </div>
                    <div class="mb-3" id="users" style="display: none;">
                        <span id="active-users-span"></span>
                        <ul id="active-users" class="list-group list-group-horizontal-sm">
                            <!--/*--> 
                            <!--/*/
                            <div th:with="condition=${#lists.size(activeUsers)}" th:remove="tag">
                                <p th:if="${condition}"><i>No active users found.</i></p>
                                <p th:unless="${condition}" class="text-muted">click on user to begin chat</p>
                            </div>
                            <li th:each="username ${activeUsers}" class="list-group-item">
                                <a class="active-user" href="javascript:void(0)" onclick="setSelectedUser('${username}')">${username}</a>
                            </li>
                            /*/-->
                            <!--*/-->                            
                        </ul>
                    </div>
                    <div id="divSelectedUser" class="mb-3" style="display: none;">
                        <span id="selectedUser" class="badge badge-secondary"></span> Selected
                    </div>
                    <div class="mb-3">                        
                        <div class="input-group" id="sendmessage" style="display: none;">
                            <input type="text" id="message" class="form-control" placeholder="Message">
                            <div class="input-group-append">
                                <button id="send" class="btn btn-primary" onclick="send()">Send</button>
                            </div>
                        </div>
                    </div>                    
                </div>
                <div class="col-md-6">
                    <div id="content"></div>
                    <div>
                        <span class="float-right">
                            <button id="clear" class="btn btn-primary" onclick="clearMessages()" style="display: none;">Clear</button>
                        </span>
                    </div>
                    <div id="response"></div>                        
                </div>
            </div>
        </div>

        <footer th:insert="fragments/common.html :: footer"></footer>
        
        <script th:src="@{/webjars/sockjs-client/1.1.2/sockjs.js}" type="text/javascript"></script>
        <script th:src="@{/webjars/stomp-websocket/2.3.3-1/stomp.js}" type="text/javascript"></script>
        <script type="text/javascript">
            var stompClient = null;
            var selectedUser = null;
            var userName = $("#from").val();
            
            function setConnected(connected) {
                $("#from").prop("disabled", connected);
                $("#connect").prop("disabled", connected);
                $("#disconnect").prop("disabled", !connected);
                if (connected) {
                    $("#users").show();
                    $("#sendmessage").show();
                } else {
                    $("#users").hide();
                    $("#sendmessage").hide();
                }
            }
            
            function connect() {
                userName = $("#from").val();
                if (userName == null || userName === "") {
                    alert('Please input a nickname!');
                    return;
                }
                $.post('/rest/user-connect',
                    { username: userName }, 
                    function(remoteAddr, status, xhr) {
                        var socket = new SockJS('/chat');
                        stompClient = Stomp.over(socket);
                        stompClient.connect({username: userName}, function () {
                            stompClient.subscribe('/topic/broadcast', function (output) {
                                showMessage(createTextNode(JSON.parse(output.body)));
                            });
                            
                            stompClient.subscribe('/topic/active', function () {
                                updateUsers(userName);
                            });
                            
                            stompClient.subscribe('/user/queue/messages', function (output) {
                                showMessage(createTextNode(JSON.parse(output.body)));
                            });

                            sendConnection(' connected to server');
                            setConnected(true);
                        }, function (err) {
                            alert('error' + err);
                        });  

                    }).done(function() { 
                        // alert('Request done!'); 
                    }).fail(function(jqxhr, settings, ex) {
                        console.log('failed, ' + ex); 
                    }
                );                              
            }

            function disconnect() {
                if (stompClient != null) {                    
                    $.post('/rest/user-disconnect',
                        { username: userName }, 
                        function() {
                            sendConnection(' disconnected from server'); 
                    
                            stompClient.disconnect(function() {
                                console.log('disconnected...');
                                setConnected(false);
                            });

                        }).done(function() { 
                            // alert('Request done!'); 
                        }).fail(function(jqxhr, settings, ex) {
                            console.log('failed, ' + ex); 
                        }
                    );                                      
                }                
            }
            
            function sendConnection(message) {
                var text = userName + message;
                sendBroadcast({'from': 'server', 'text': text});
                
                // for first time or last time, list active users:
                updateUsers(userName);
            }
            
            function sendBroadcast(json) {
                stompClient.send("/app/broadcast", {}, JSON.stringify(json));
            }
           
            function send() {
                var text = $("#message").val();
                if (selectedUser == null) {
                    alert('Please select a user.');
                    return;
                }
                stompClient.send("/app/chat", {'sender': userName},
                        JSON.stringify({'from': userName, 'text': text, 'recipient': selectedUser}));                
                $("#message").val("");
            }

            function createTextNode(messageObj) {
                var classAlert = 'alert-info';
                var fromTo = messageObj.from;
                var addTo =  fromTo;
                
                if (userName == messageObj.from) {
                    fromTo = messageObj.recipient;
                    addTo =  'to: ' + fromTo;
                }
                
                if (userName != messageObj.from && messageObj.from != "server") {
                    classAlert = "alert-warning";
                }
                
                if (messageObj.from != "server") {
                    addTo = '<a href="javascript:void(0)" onclick="setSelectedUser(\'' + fromTo + '\')">' + addTo + '</a>'
                }
                return '<div class="row alert ' + classAlert + '"><div class="col-md-8">' +
                        messageObj.text +
                        '</div><div class="col-md-4 text-right"><small>[<b>' +
                        addTo +
                        '</b> ' +
                        messageObj.time + 
                        ']</small>' +
                        '</div></div>';
            }
            
            function showMessage(message) {
                $("#content").html($("#content").html() + message);
                $("#clear").show();
            }
            
            function clearMessages() {
                $("#content").html("");
                $("#clear").hide();
            }
            
            function setSelectedUser(username) {
                selectedUser = username;
                $("#selectedUser").html(selectedUser);
                if ($("#selectedUser").html() == "") {
                    $("#divSelectedUser").hide();
                }
                else {
                    $("#divSelectedUser").show();
                }
            }
            
            function updateUsers(userName) {
                // console.log('List of users : ' + userList);               
                var activeUserSpan = $("#active-users-span");
                var activeUserUL = $("#active-users");

                var index;
                activeUserUL.html('');
                
                var url = '/rest/active-users-except/' + userName;
                $.ajax({
                    type: 'GET',
                    url: url,
                    // data: data, 
                    dataType: 'json', // ** ensure you add this line **
                    success: function(userList) {
                        if (userList.length == 0) {
                            activeUserSpan.html('<p><i>No active users found.</i></p>');
                        }
                        else {
                            activeUserSpan.html('<p class="text-muted">click on user to begin chat</p>');
                            
                            for (index = 0; index < userList.length; ++index) {
                                if (userList[index] != userName) {
                                    activeUserUL.html(activeUserUL.html() + createUserNode(userList[index], index));                                    
                                }
                            }
                            /*
                            $.each(userList, function(index, item) {
                                //now you can access properties using dot notation
                            }); 
                            */
                        }
                    },
                    error: function(XMLHttpRequest, textStatus, errorThrown) {
                        alert("error occurred");
                    }
                });
            }
            
            function createUserNode(username, index) {
                return '<li class="list-group-item">' +
                       '<a class="active-user" href="javascript:void(0)" onclick="setSelectedUser(\'' + username + '\')">' + username + '</a>' +
                       '</li>';
            }
        </script>
    </body>
</html>
                    

And that's all. Here the result in the browser, when you navigate to http://localhost:8080/sockjs-message:

1st User's Screen=

1st User's Screen

2nd User's Screen=

2nd User's Screen

3rd User's Screen=

3rd User's Screen

Conclusion

As I mentioned in above, there are lot of room for improvement to get full blown solution of Spring Boot and WebSocket application. As example, you can use Spring Security to secure your WebSockets implementation. And so far, we're only using simple message broker. It's very simple to switch to full-featured message broker like ActiveMQ or RabbitMQ.

References: