Spring Boot WebSocket 两种集成方式详解

Spring容器启动
    └── ServerEndpointExporter.afterPropertiesSet()
            └── 扫描 @ServerEndpoint 类
                    └── 注册到 ServerContainer（Tomcat WebSocket 运行时）

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

@Configuration
public class WebSocketConfig {
    /**
     * 向 Spring 容器注册 ServerEndpointExporter
     * 它会在应用启动后，将所有 @ServerEndpoint 注解的类注册到底层 Servlet 容器
     * 注意：使用外部容器（如独立部署的 Tomcat）时，不需要注册此 Bean，
     *       外部容器会自行完成注册
     */
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

@Component  // ① 必须交给 Spring 容器，才能在内部注入 Service 等 Bean
@ServerEndpoint("/ws/chat/{roomId}")  // ② 定义 WebSocket 连接路径
public class ChatWebSocketServer {
    // ③ 核心踩坑点：@ServerEndpoint 每个连接都会 new 一个新实例
    //    因此不能用普通的 @Autowired 字段注入，必须用 static 字段 + setter 注入
    private static MessageService messageService;
    @Autowired
    public void setMessageService(MessageService messageService) {
        ChatWebSocketServer.messageService = messageService;
    }
    // ④ 线程安全：用 ConcurrentHashMap 管理所有在线 Session
    private static final ConcurrentHashMap<String, Session> SESSION_MAP
            = new ConcurrentHashMap<>();
    private Session session;
    private String userId;
    /**
     * 连接建立成功时触发
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("roomId") String roomId) {
        this.session = session;
        this.userId = session.getId();
        SESSION_MAP.put(userId, session);
        System.out.printf("用户 [%s] 加入房间 [%s]，当前在线人数：%d%n",
                userId, roomId, SESSION_MAP.size());
    }
    /**
     * 收到客户端消息时触发
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        System.out.printf("收到用户 [%s] 的消息：%s%n", userId, message);
        // 调用业务 Service 处理消息（static 注入，可正常使用）
        messageService.saveMessage(userId, message);
        // 广播给所有在线用户
        broadcastMessage(userId + ": " + message);
    }
    /**
     * 连接关闭时触发
     */
    @OnClose
    public void onClose() {
        SESSION_MAP.remove(userId);
        System.out.printf("用户 [%s] 断开连接，当前在线人数：%d%n",
                userId, SESSION_MAP.size());
    }
    /**
     * 发生错误时触发
     */
    @OnError
    public void onError(Session session, Throwable error) {
        System.err.printf("用户 [%s] 发生错误：%s%n", userId, error.getMessage());
        error.printStackTrace();
    }
    /**
     * 广播消息给所有在线用户
     */
    private void broadcastMessage(String message) {
        SESSION_MAP.values().forEach(s -> {
            try {
                if (s.isOpen()) {
                    s.getBasicRemote().sendText(message);
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        });
    }
    /**
     * 向指定用户发送消息（可供外部调用）
     */
    public static void sendMessageToUser(String userId, String message) {
        Session session = SESSION_MAP.get(userId);
        if (session != null && session.isOpen()) {
            try {
                session.getBasicRemote().sendText(message);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }
}

HTTP 请求升级为 WebSocket
    └── DispatcherServlet
            └── WebSocketHandlerMapping（路径路由）
                    └── 你的 WebSocketHandler（处理具体逻辑）

@Configuration
@EnableWebSocket  // ① 开启 Spring WebSocket 支持
public class WebSocketConfig implements WebSocketConfigurer {  // ② 实现此接口
    @Autowired
    private ChatWebSocketHandler chatWebSocketHandler;
    @Autowired
    private WebSocketAuthInterceptor authInterceptor;
    /**
     * ③ 重写此方法，将处理器注册到指定路径
     */
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry
            .addHandler(chatWebSocketHandler, "/ws/chat")  // 注册处理器和路径
            .addInterceptors(authInterceptor)               // 可添加握手拦截器
            .setAllowedOrigins("*");                        // 跨域配置
    }
}

@Component
public class WebSocketAuthInterceptor implements HandshakeInterceptor {
    /**
     * WebSocket 握手前执行：返回 false 则拒绝连接
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request,
                                   ServerHttpResponse response,
                                   WebSocketHandler wsHandler,
                                   Map<String, Object> attributes) throws Exception {
        // 从请求参数或 Header 中获取 Token，验证用户身份
        String token = ((ServletServerHttpRequest) request)
                .getServletRequest().getParameter("token");
        if (token == null || !isValidToken(token)) {
            response.setStatusCode(HttpStatus.UNAUTHORIZED);
            return false;  // 拒绝握手
        }
        // 将用户信息存入 attributes，后续 Handler 中可以取到
        attributes.put("userId", parseUserId(token));
        return true;
    }
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
        // 握手后执行，一般留空
    }
    private boolean isValidToken(String token) {
        // 实际项目中调用 JWT 解析或 Redis 校验
        return token.startsWith("valid_");
    }
    private String parseUserId(String token) {
        return token.replace("valid_", "");
    }
}

@Component  // ① 交给 Spring 容器管理，正常 @Autowired 注入无任何问题
public class ChatWebSocketHandler extends TextWebSocketHandler {  // ② 继承此类处理文本消息
    @Autowired
    private MessageService messageService;  // ③ 单例 Handler，直接 @Autowired 完全没问题
    // 维护在线 Session 的线程安全 Map
    private static final ConcurrentHashMap<String, WebSocketSession> SESSION_MAP
            = new ConcurrentHashMap<>();
    /**
     * 连接建立成功时触发
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String userId = (String) session.getAttributes().get("userId");
        SESSION_MAP.put(userId, session);
        System.out.printf("用户 [%s] 已连接，当前在线：%d%n", userId, SESSION_MAP.size());
    }
    /**
     * 收到文本消息时触发
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message)
            throws Exception {
        String userId = (String) session.getAttributes().get("userId");
        String payload = message.getPayload();
        System.out.printf("收到 [%s] 的消息：%s%n", userId, payload);
        // 调用业务 Service
        messageService.saveMessage(userId, payload);
        // 广播消息
        broadcastMessage(userId + ": " + payload);
    }
    /**
     * 连接关闭时触发
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status)
            throws Exception {
        String userId = (String) session.getAttributes().get("userId");
        SESSION_MAP.remove(userId);
        System.out.printf("用户 [%s] 已断开，当前在线：%d%n", userId, SESSION_MAP.size());
    }
    /**
     * 传输异常时触发
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception)
            throws Exception {
        System.err.println("传输错误：" + exception.getMessage());
        session.close(CloseStatus.SERVER_ERROR);
    }
    private void broadcastMessage(String message) {
        SESSION_MAP.values().forEach(s -> {
            try {
                if (s.isOpen()) {
                    s.sendMessage(new TextMessage(message));
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        });
    }
}

java.lang.ClassCastException: class X cannot be cast to class Y

@ServerEndpoint("/ws/chat")
@Component
public class ChatServer {
    @Autowired
    private UserService userService;  // 运行时是 null！
    @OnMessage
    public void onMessage(String msg) {
        userService.doSomething(msg);  // NullPointerException
    }
}

@ServerEndpoint("/ws/chat")
@Component
public class ChatServer {
    private static UserService userService;
    @Autowired  // Spring 对它管理的那个实例执行此方法，写入 static 字段
    public void setUserService(UserService userService) {
        ChatServer.userService = userService;
    }
}

// 嵌入式容器：需要
// 外部容器：删掉这个 Bean
@Bean
public ServerEndpointExporter serverEndpointExporter() {
    return new ServerEndpointExporter();
}

// 原生方式路径示例
const ws1 = new WebSocket('ws://localhost:8080/ws/chat/room123');
// Spring 整合方式路径示例（附带 token 参数用于握手拦截器验证）
const ws2 = new WebSocket('ws://localhost:8080/ws/chat?token=valid_user001');
ws2.onopen = () => {
    console.log('连接已建立');
    ws2.send(JSON.stringify({ type: 'chat', content: 'Hello!' }));
};
ws2.onmessage = (event) => {
    console.log('收到消息：', event.data);
};
ws2.onclose = (event) => {
    console.log('连接已关闭，code:', event.code);
};
ws2.onerror = (error) => {
    console.error('连接错误：', error);
};

需要权限校验、Session 管理、与 Spring Security 集成？
    └── 选 Spring 整合方式（WebSocketHandler）
快速实现、团队熟悉 Java EE 规范、无复杂 Spring 生态依赖？
    └── 选原生方式（@ServerEndpoint）
需要打 war 包部署到外部 Tomcat？
    └── 两种都行，但原生方式记得去掉 ServerEndpointExporter Bean
追求更强的消息抽象（发布订阅、广播频道）？
    └── 考虑 Spring WebSocket + STOMP 协议（本文未涉及，可作进阶方向）