Spring boot项目整合WebSocket方法

下面是关于Spring boot整合WebSocket的完整攻略。

环境准备

JDK 1.8及以上版本
Maven 3.2及以上版本
Spring Boot 2.x版本
IDE：Eclipse、IDEA

依赖配置

在Spring Boot项目中开启WebSocket功能需要引入相关依赖，添加以下依赖到项目的pom.xml文件中：

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-websocket</artifactId>
</dependency>

WebSocket配置

自定义WebSocket处理器

自定义一个Websocket处理器来处理WebSocket连接和消息：

@Component
public class MyWebSocketHandler extends TextWebSocketHandler {

    private static final Logger LOGGER = LoggerFactory.getLogger(MyWebSocketHandler.class);

    // 保存在线用户列表
    private static final Map<String, WebSocketSession> USERS = new ConcurrentHashMap<>();

    // 添加一个新用户到在线列表
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String username = (String) session.getAttributes().get("username");
        LOGGER.info("新用户 {} 加入聊天室", username);
        USERS.put(username, session);
    }

    // 处理用户发送的消息
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        LOGGER.info("收到用户 {} 发送的消息：{}", session.getAttributes().get("username"), message.getPayload());
        sendMessageToAllUsers(session, message.getPayload()); // 群发消息
    }

    // 从在线列表中删除用户并断开WebSocket连接
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String username = (String) session.getAttributes().get("username");
        LOGGER.info("用户 {} 离开聊天室，断开WebSocket连接", username);
        USERS.remove(username);
    }

    // 群发消息
    private void sendMessageToAllUsers(WebSocketSession fromUser, String message) throws IOException {
        TextMessage textMessage = new TextMessage(message);
        for (Map.Entry<String, WebSocketSession> entry : USERS.entrySet()) {
            WebSocketSession session = entry.getValue();
            if (!session.isOpen()) {
                continue;
            }
            if (session == fromUser) {
                continue;
            }
            session.sendMessage(textMessage);
        }
    }
}

WebSocket配置类

定义一个WebSocket配置类来设置WebSocket的相关信息：

@Configuration
@EnableWebSocket
public class MyWebSocketConfig implements WebSocketConfigurer {

    @Autowired
    private MyWebSocketHandler myWebSocketHandler;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(myWebSocketHandler, "/ws").setAllowedOrigins("*"); // 配置WebSocket端点，允许跨域访问
    }
}

Controller层

在Controller层中添加一个方法来接收WebSocket连接，并将用户信息存储在WebSocketSession中：

@Controller
public class UserController {

    @GetMapping("/login")
    public String login(String username, HttpServletRequest request) {
        request.getSession().setAttribute("username", username); // 将username保存到session中
        return "chat";
    }

    @Autowired
    private SimpMessagingTemplate messagingTemplate;

    // 处理WebSocket连接请求
    @MessageMapping("/send")
    public void handleWebSocketMessage(Principal principal, ChatMessage message) throws Exception {
        String fromUser = principal.getName(); // 获取发送者的用户名
        String toUser = message.getToUser();
        if (StringUtils.isBlank(toUser)) {
            messagingTemplate.convertAndSend("/topic/public", new ChatMessage(fromUser, message.getContent())); // 广播消息
        } else {
            messagingTemplate.convertAndSendToUser(toUser, "/queue/message", new ChatMessage(fromUser, message.getContent())); // 指定用户发送消息
        }
    }
}

页面展示

使用JSP页面来展示聊天页面，用户在登录页面输入用户名后，跳转到聊天室页面：

<%@ page contentType="text/html;charset=UTF-8" language="java" %>
<html>
<head>
    <title>Spring Boot WebSocket Demo</title>
    <style>
        .messageForm {
            display: flex;
            margin-bottom: 10px;
        }

        .messageInput {
            flex-grow: 1;
            margin-right: 5px;
        }

        .messageButton {
            width: 80px;
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>Spring Boot WebSocket Demo</h1>
        <hr/>
        <h2>Welcome <%=session.getAttribute("username")%></h2>
        <div id="chatBox">
        </div>
        <div class="messageForm">
            <input type="text" id="messageInput" class="messageInput" placeholder="请输入你的消息...">
            <button type="button" class="messageButton" onclick="sendMessage()">发送</button>
        </div>
    </div>
    <script src="/webjars/jquery/3.3.1/jquery.min.js"></script>
    <script src="/webjars/sockjs-client/1.3.0/sockjs.min.js"></script>
    <script src="/webjars/stomp-websocket/2.3.3/stomp.min.js"></script>
    <script>
        // 连接WebSocket服务器
        var stompClient = Stomp.over(new SockJS('/ws'));
        stompClient.connect({}, function (frame) {
            console.log('Connected: ' + frame);
            stompClient.subscribe('/user/queue/message', function (response) {
                showMessage(JSON.parse(response.body));
            });
            stompClient.subscribe('/topic/public', function (response) {
                showMessage(JSON.parse(response.body));
            });
        });

        // 显示消息
        function showMessage(message) {
            var $chatBox = $('#chatBox');
            var $messageDiv = $('<div></div>');
            $messageDiv.html(message.fromUser + ' : ' + message.content);
            $chatBox.append($messageDiv);
        }

        // 发送消息
        function sendMessage() {
            var $messageInput = $('#messageInput');
            var message = $messageInput.val();
            $messageInput.val('');
            stompClient.send("/app/send", {}, JSON.stringify({'content': message}));
        }
    </script>
</body>
</html>