Spring WebSocket 实现方式

Spring Boot 提供了对 WebSocket 的全面支持，通过封装了如 TextWebSocketHandler、HandshakeInterceptor、WebSocketConfigurer 等接口，开发者可以便捷地实现实时的 WebSocket 通信。相较于传统的 HTTP 请求-响应模式，WebSocket 是一种长连接，可以支持服务器与客户端之间的全双工通信。这种方式特别适用于实时性要求较高的应用场景，比如：在线聊天、实时推送等。

本实现通过五个主要部分展示了如何配置和使用 WebSocket：

1. 依赖引入：引入 Spring WebSocket 相关依赖以开启其自动配置功能。
2. WebSocket 处理器：通过继承 TextWebSocketHandler，处理连接的建立、消息接收、连接关闭等 WebSocket 事件。
3. 会话管理：通过自定义的 WsSessionManager 实现会话的存储、查找和移除，确保 WebSocket 连接的有效管理。
4. 握手拦截器：通过实现 HandshakeInterceptor，在 WebSocket 握手前进行身份验证或其他预处理操作。
5. WebSocket 全局配置：通过 WebSocketConfigurer，进行 WebSocket 处理器和拦截器的全局注册，并提供跨域支持。

实现步骤：

1. 添加 WebSocket 依赖

首先，需要在 Spring Boot 项目中引入 WebSocket 依赖。Spring Boot 提供了 spring-boot-starter-websocket 依赖，自动集成了 WebSocket 的核心功能。

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

引入该依赖后，Spring Boot 自动加载 WebSocket 的相关配置，使开发者可以通过简单的配置和扩展实现 WebSocket 服务。

2. WebSocket 处理器 HttpAuthHandler

HttpAuthHandler 是一个继承自 TextWebSocketHandler 的类，用于处理 WebSocket 连接的建立、消息接收和连接关闭事件。通过重写核心的三个方法，能够处理各种 WebSocket 生命周期中的事件。

package cn.coder4j.study.example.websocket.handler;

import cn.coder4j.study.example.websocket.config.WsSessionManager;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.time.LocalDateTime;

/**
 * WebSocket 消息处理器，继承 TextWebSocketHandler 实现文本消息的处理。
 *
 * 通过重写 afterConnectionEstablished、handleTextMessage 和 afterConnectionClosed 方法，
 * 实现 WebSocket 连接建立、消息处理和连接关闭时的逻辑。
 */
@Component
public class HttpAuthHandler extends TextWebSocketHandler {

    /**
     * 当 WebSocket 连接建立成功时触发，类似于原生注解中的 @OnOpen。
     *
     * @param session 当前 WebSocket 会话
     * @throws Exception 处理异常
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Object token = session.getAttributes().get("token");
        if (token != null) {
            // 将用户连接存入在线用户缓存
            WsSessionManager.add(token.toString(), session);
        } else {
            throw new RuntimeException("用户登录已失效!");
        }
    }

    /**
     * 当 WebSocket 接收到消息时触发，类似于原生注解中的 @OnMessage。
     *
     * @param session 当前 WebSocket 会话
     * @param message 接收到的文本消息
     * @throws Exception 处理异常
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 获取客户端发送的消息内容
        String payload = message.getPayload();
        Object token = session.getAttributes().get("token");
        System.out.println("服务器接收到 " + token + " 发送的消息：" + payload);
        // 向客户端发送响应消息
        session.sendMessage(new TextMessage("服务器回复给 " + token + " 的消息：" + payload + "，时间：" + LocalDateTime.now().toString()));
    }

    /**
     * 当 WebSocket 连接关闭时触发，类似于原生注解中的 @OnClose。
     *
     * @param session 当前 WebSocket 会话
     * @param status  连接关闭状态
     * @throws Exception 处理异常
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Object token = session.getAttributes().get("token");
        if (token != null) {
            // 用户退出，移除会话缓存
            WsSessionManager.remove(token.toString());
        }
    }
}

关键方法解释：

• afterConnectionEstablished: 在 WebSocket 连接成功后触发。通常用于初始化 WebSocket 会话、进行身份验证和存储会话信息。
• handleTextMessage: 在接收到文本消息时触发。处理客户端发来的消息，并返回响应。
• afterConnectionClosed: 在连接关闭时触发，清理 WebSocket 资源，比如移除缓存中的会话信息。

3. 会话管理器 WsSessionManager

WsSessionManager 用于管理 WebSocket 会话。它使用了 ConcurrentHashMap 实现线程安全的会话存储。这个管理器可以存储、查找、删除 WebSocket 连接，方便服务端主动向客户端推送消息。

package cn.coder4j.study.example.websocket.config;

import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.WebSocketSession;

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

/**
 * WebSocket 会话管理器，用于管理和操作所有的 WebSocket 会话。
 */
@Slf4j
public class WsSessionManager {

    // 使用 ConcurrentHashMap 存储 WebSocket 会话，确保线程安全
    private static final ConcurrentHashMap<String, WebSocketSession> SESSION_POOL = new ConcurrentHashMap<>();

    /**
     * 添加 WebSocket 会话到会话池。
     *
     * @param key     用户标识（通常为 token）
     * @param session WebSocket 会话对象
     */
    public static void add(String key, WebSocketSession session) {
        SESSION_POOL.put(key, session);
    }

    /**
     * 从会话池中移除并返回会话。
     *
     * @param key 用户标识
     * @return 被移除的 WebSocket 会话
     */
    public static WebSocketSession remove(String key) {
        return SESSION_POOL.remove(key);
    }

    /**
     * 移除会话并关闭连接。
     *
     * @param key 用户标识
     */
    public static void removeAndClose(String key) {
        WebSocketSession session = remove(key);
        if (session != null) {
            try {
                session.close(); // 关闭连接
            } catch (IOException e) {
                log.error("关闭 WebSocket 连接时出现异常", e);
            }
        }
    }

    /**
     * 获取会话对象。
     *
     * @param key 用户标识
     * @return 对应的 WebSocket 会话
     */
    public static WebSocketSession get(String key) {
        return SESSION_POOL.get(key);
    }
}

功能：

• 添加会话: 通过 add 方法存储 WebSocket 会话信息。
• 移除和关闭会话: 通过 removeAndClose 方法删除会话并关闭 WebSocket 连接。
• 获取会话: 通过 get 方法获取特定用户的 WebSocket 会话。

4. 握手拦截器 MyInterceptor

MyInterceptor 实现了 HandshakeInterceptor 接口，用于在 WebSocket 握手前后执行特定逻辑，比如身份验证、参数检查等操作。

package cn.coder4j.study.example.websocket.interceptor;

import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpUtil;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.HashMap;
import java.util.Map;

/**
 * WebSocket 握手拦截器，在握手前后进行预处理操作。
 */
@Component
public class MyInterceptor implements HandshakeInterceptor {

    /**
     * 在握手前执行的逻辑，通常用于身份验证。
     *
     * @param request    HTTP 请求
     * @param response   HTTP 响应
     * @param wsHandler  WebSocket 处理器
     * @param attributes 存储握手后的属性
     * @return 是否允许握手成功
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        System.out.println("握手开始");
        // 获取请求参数
        HashMap<String, String> paramMap = HttpUtil.decodeParamMap(request.getURI().getQuery(), "

utf-8");
        String token = paramMap.get("token");
        if (StrUtil.isNotBlank(token)) {
            // 将 token 放入属性中，供后续使用
            attributes.put("token", token);
            System.out.println("用户 token " + token + " 握手成功！");
            return true;
        }
        System.out.println("用户登录已失效，握手失败！");
        return false;
    }

    /**
     * 在握手完成后执行的逻辑。
     *
     * @param request   HTTP 请求
     * @param response  HTTP 响应
     * @param wsHandler WebSocket 处理器
     * @param exception 异常信息
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        System.out.println("握手完成");
    }
}

功能：

• beforeHandshake: 在握手前进行验证，确保客户端提供的 token 有效。如果 token 无效，则拒绝握手。
• afterHandshake: 在握手成功后执行的逻辑，比如记录日志或进行一些后续处理。

5. WebSocket 配置类 WebSocketConfig

通过实现 WebSocketConfigurer 接口，全局注册 WebSocket 处理器和拦截器，配置 WebSocket 的端点路径和跨域设置。

package cn.coder4j.study.example.websocket.config;

import cn.coder4j.study.example.websocket.handler.HttpAuthHandler;
import cn.coder4

j.study.example.websocket.interceptor.MyInterceptor;
import org.springframework.beans.factory.annotation.Autowired;
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;

/**
 * WebSocket 配置类，注册 WebSocket 处理器和拦截器。
 */
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Autowired
    private HttpAuthHandler httpAuthHandler;
    @Autowired
    private MyInterceptor myInterceptor;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // 注册 WebSocket 处理器，并添加握手拦截器
        registry
                .addHandler(httpAuthHandler, "/myWS") // 设置 WebSocket 端点路径
                .addInterceptors(myInterceptor) // 添加握手拦截器
                .setAllowedOrigins("*"); // 允许跨域，适用于本地调试
    }
}

功能：

• addHandler: 注册 WebSocket 处理器，并指定端点路径 /myWS。
• addInterceptors: 添加握手拦截器，用于在握手阶段执行身份验证。
• setAllowedOrigins("*"): 设置跨域策略，允许所有域访问该 WebSocket 端点。在生产环境中，通常应限制跨域策略以确保安全性。

---

总结

通过上述实现方式，Spring Boot 提供了强大的 WebSocket 支持。在本实现中：

• HttpAuthHandler 用于处理 WebSocket 连接建立、消息接收和连接关闭事件。
• WsSessionManager 管理所有的 WebSocket 会话，便于服务端主动推送消息。
• MyInterceptor 在 WebSocket 握手阶段进行身份验证，确保连接的合法性。
• WebSocketConfig 实现了 WebSocket 全局配置，注册处理器和拦截器，并允许跨域请求。

该实现具有较好的扩展性，适合中小型应用。根据需求，还可以进一步引入 Redis 来实现分布式 WebSocket 管理和消息推送等高级功能。