使用 Springboot websocket 实现聊天室

文章首发于个人博客，欢迎访问关注：https://www.lin2j.tech

通过使用spring集成的websocket和原生H5,实现一个聊天室,以此加深对websocket的了解.
文末会附带代码, 部分代码有注释

• 出现原因
  • 弥补 HTTP协议的不足,使用HTTP协议,服务端无法对客户端进行主动推送,只能依靠长连接,或者轮询来进行获取实时消息
• 简单原理:
  • 客户端先借用HTTP协议, 在第一次握手的时候,告诉服务端,接下来我要把请求升级为Websocket, 然后服务端同意就返回true, 否则返回false, 返回true后将建立 TCP 连接
  • 保持住 TCP 连接,接下来的消息就通过该 TCP 连接进行传输.
  • 传输数据时,只传送真正的数据(不必向HTTP一样,除了真正的数据外,还有HTTP Header)
  • 因为数据在网络链路的传输要经过若干个中间节点,有些节点会将一段时间内没有发送信息的连接断开,因此我们可以用 Ping/Pong Frame 包(一种特殊的数据包)来维持连接.
  • 这是一条全双工的通信连接, 此时, 客户端和服务端是平等的, 因此服务端可以主动发送信息给客户端
• 优点
  • 效率高
  • 全双工
• 应用场景
  • 多人在线游戏
  • WebIM
  • 多人在线编辑
  • 实时消息推送

详细原理在这里

示例源码 GitHub

• 简单介绍接下来使用的类,类名是我定义的, 可以先看代码,有不清楚的地方可以看这里有没有解释 😃
  • WebSocketConfiguration: 配置类, 该类实现了 WebSocketConfigurer 接口, 需要实现其 registerWebSocketHandlers 方法, 用于配置webSocket的入口(就是链接地址), 处理器, 拦截器.
  • WebSocketInterceptor: 拦截器,该类实现了 HandShakeIntercepetor 接口, 并实现了其中的beforeHandshake 方法, 每个socket要建立之前,先拦截该请求, 解析客户端的请求的地址
  • ChatRoomWebSocketHandler: 处理器, 该类实现了 WebSocketHandler, 并重写了其中的方法, 这个类的功能主要是将客户端的 id 和对应的会话用 Map 保存起来, 然后在有 连接、接收、发送, 断开连接时, 对 map进行相应的操作
  • 涉及到的注解
    • @Configuration: 标记 WebSocketConfiguration 类
    • @EnableWebSocket: 开启spring的websocket功能
    • @Service: 这里是用来把处理器标记成一个服务层组件来使用

-上代码
WebSocketInterceptor.java

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;

import java.util.Map;

 // 拦截器,在建立连接之前,拦截请求
public class WebSocketInterceptor implements HandshakeInterceptor {

    /**
     * 在webSocket连接建立之前(tcp三次握手),拦截连接请求,将请求中的部分信息存储起来,存储在 map 中
     * 继续握手返回true, 中断握手返回false.
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse
            , WebSocketHandler webSocketHandler, Map<String, Object> map) throws Exception {
        if(serverHttpRequest instanceof ServletServerHttpRequest){
            ServletServerHttpRequest request = (ServletServerHttpRequest)serverHttpRequest;
            // 将 url 分割成两部分, userId= 前和 userId= 后
            // 比如 http://localhost:8080/chat/userId=1234
            // 被分割成 http://localhost:8080/chat 和 1234, 1234就是我们需要的id
            String userId = request.getURI().toString().split("userId=")[1];
            map.put("WEBSOCKET_USER_ID", userId);
        }
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse
            , WebSocketHandler webSocketHandler, Exception e) {
        // do nothing
    }
}

ChatRoomWebSocketHandler.java

import org.springframework.stereotype.Service;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

// 这里将处理器作为一个服务层组件来使用
@Service
public class ChatRoomWebSocketHandler implements WebSocketHandler {
    /**
     * 用户广播, 使用一个 Map 来存储用户id和websocket会话的映射
     */
    public static final ConcurrentHashMap<String, WebSocketSession> USER_SESSIONS = new ConcurrentHashMap<>();


    @Override
    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws Exception {
        // 取得此次连接的用户的id
        String userId = webSocketSession.getUri().toString().split("userId=")[1];
        // 将此用户的会话保存起来
        USER_SESSIONS.put(userId, webSocketSession);
    }

    @Override
    public void handleMessage(WebSocketSession webSocketSession, WebSocketMessage<?> webSocketMessage) throws Exception {
        // 使用 sendMessageToUser 和 sendMessageToAllUsers 来代替这个方法
    }

    @Override
    public void handleTransportError(WebSocketSession webSocketSession, Throwable throwable) throws Exception {
        if(webSocketSession.isOpen()){
            // 移除会话
            USER_SESSIONS.remove(getClientId(webSocketSession));
            webSocketSession.close();
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession webSocketSession, CloseStatus closeStatus) throws Exception {
        // 移除会话
        USER_SESSIONS.remove(getClientId(webSocketSession));
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 获取用户的id
     */
    private String getClientId(WebSocketSession webSocketSession){
        if(webSocketSession != null){
            // 此处 .getAttributes() 返回的是一个 Map ,
            // 就是在 WebSocketInterceptor 的 beforeHandshake 方法中保存了 userId 的那个 Map
            return (String)webSocketSession.getAttributes().get("WEBSOCKET_USER_ID");
        }
        return null;
    }

    /**
     * 发送给某个在线用户
     * @param userId
     * @param message
     */
    public boolean sendMessageToUser(String userId, TextMessage message){
        if(userId == null || "".equals(userId)){
            return false;
        }

        WebSocketSession session;
        if((session = USER_SESSIONS.get(userId)) == null){
            return false;
        }

        if(!session.isOpen()){
            return false;
        }

        try{
            session.sendMessage(message);
        }catch (IOException e){
            return false;
        }

        return true;
    }

    /**
     * 给全体在线用户发送信息
     * @param message
     * @return
     */
    public boolean sendMessageTOAllUsers(TextMessage message){
        AtomicBoolean success = new AtomicBoolean(true);
        // 使用 lambda 表达式
        USER_SESSIONS.values().forEach((session)-> {
            try {
                session.sendMessage(message);
            } catch (IOException e) {
                success.set(false);
            }
        });

        return success.get();
    }
}

WebSocketConfiguration.java

import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
//开启注解接收和发送消息
@EnableWebSocket
public class WebSocketConfiguration implements WebSocketConfigurer {
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // addHandler 为每个请求分配一个处理器
        // "/chat-room/{userId}" 为服务器终端url地址
        // 增加拦截器 map
        // 设置所有的域都可以访问
        registry.addHandler(new ChatRoomWebSocketHandler(), "/chat-room/{userId}")
                .addInterceptors(new WebSocketInterceptor())
                .setAllowedOrigins("*");
    }
}

ChatController.java

import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestMethod;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.socket.TextMessage;

@RestController
@RequestMapping("/chat")
public class ChatController {

    /**
     * 注入
     */
    final ChatRoomWebSocketHandler handler;

    public ChatController(ChatRoomWebSocketHandler handler) {
        this.handler = handler;
    }

    /**
     * 两个用户之间发送信息, 虽然这里用 GET 不是很合适
     * @param sendId 发送人
     * @param receiveId 接收人
     * @param message 信息内容
     */
    @RequestMapping(value = "/{sendId}/to/{receiveId}", method = RequestMethod.GET)
    public void  sendToUser(@PathVariable("sendId") String sendId
            , @PathVariable("receiveId") String receiveId, String message){
        if((!"".equals(receiveId) && (!"".equals(sendId)))){
            // 组织语言
            handler.sendMessageToUser(receiveId, new TextMessage(message));
        }
    }

    @RequestMapping(value = "/{sendId}/toAll", method = RequestMethod.GET)
    public void sendToAllUser(@PathVariable("sendId") String sendId, String message){
        if(!"".equals(sendId)){
            handler.sendMessageTOAllUsers(new TextMessage(message));
        }
    }
}

前端页面

<!DOCTYPE html>
<html>
<head>
    <title>chat.html</title>
    <meta charset="UTF-8">
</head>

<body>
    <input type="text"  id="userName" placeholder="请先输入用户名" />
    <button onclick="connectWebSocket()">加入房间</button>
    <button onclick="closeWebSocket()">退出群聊</button><br>

    <input id="all" type="text" placeholder="群发"/> <button onclick="sendAll()">发送</button><br>

    <p>群聊房间</p>
    <textarea id="room" cols="40" rows="10" readonly="readonly"></textarea>

<script type="text/javascript" src="../js/jquery-3.3.1.js"></script>
<script type="text/javascript">
    var websocket=null;

    //关闭浏览器时
    window.onunload = function() {
        //关闭连接
        closeWebSocket();
    }

    //建立WebSocket连接
    function connectWebSocket(){

        var userId = document.getElementById('userName').value;
        // 不可更改用户名
        document.getElementById('userName').setAttribute("readOnly", true);

        //建立webSocket连接
        websocket = new WebSocket("ws://localhost:8080/chat-room/userId=" + userId);

        //打开webSokcet连接
        websocket.onopen = function () {
            var content = userId + '加入群聊';
            $.get('http://localhost:8080/chat/' + userId + '/toAll?message=' + content);
        }

        //关闭webSocket连接
        websocket.onclose = function () {
            //关闭连接
            console.log("onclose");
        }

        //接收信息
        websocket.onmessage = function (msg) {
            document.getElementById('room').append(msg.data + '\n');
        }
    }

    function sendAll(){
        // 聊天内容
        var content = document.getElementById('all').value;
        // 发送人
        var userId = document.getElementById('userName').value;
        $.get('http://localhost:8080/chat/' + userId + '/toAll?message=' + userId + ': ' + content);

    }

    //关闭连接
    function closeWebSocket(){
        var userId = document.getElementById('userName').value;
        $.get('http://localhost:8080/chat/' + userId + '/toAll?message=' + '退出群聊');
        websocket.send(userId + '退出群聊');
        if(websocket != null) {
            websocket.close();
        }
        alert("退出成功")
    }

</script>
</body>
</html>

运行截图

如果要实现私聊的功能, 应该为这两个用户之间另外建立一条连接