websocket

WebSocket 安全与推送通知实现

本文介绍了一个基于 Spring WebSocket 的推送通知系统，包含了身份认证和会话管理的实现，适用于需要进行消息推送的 Web 应用。

项目架构

├── TokenHandshakeInterceptor.java # WebSocket握手拦截器（身份认证）
├── WebSocketConfig.java # WebSocket全局配置
└── WebSocketNotificationHandler.java # WebSocket核心处理器

1. TokenHandshakeInterceptor.java（握手拦截器）

TokenHandshakeInterceptor 是 WebSocket 的握手拦截器，负责在 WebSocket 连接建立之前进行身份认证，确保用户的 Token 合法。

@Slf4j
public class TokenHandshakeInterceptor implements HandshakeInterceptor {
    
    /**
     * 在握手之前进行身份验证（核心安全关卡）
     * @return true=验证通过允许连接，false=拒绝连接
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, 
                                  WebSocketHandler wsHandler, Map<String, Object> attributes) {
        // 1. 从URL查询参数中获取token（如：ws://xxx?token=xxxx）
        String token = getQueryParam(request, "token");  

        // 2. 处理Bearer Token格式（去掉前缀）
        if (token != null && token.startsWith("Bearer ")) {
            token = token.substring(7);  
        }

        try {
            // 3. 解析Token获取用户ID（核心鉴权逻辑）
            Integer userId = JwtUtils.getUserIdByToken(token);
            if (userId != null) {
                // 4. 将用户ID存入WebSocket会话属性（后续处理可用）
                attributes.put("userId", userId.toString());
                log.info("Token验证成功，用户ID: {}", userId); 
                return true;  // 放行连接
            } else {
                log.warn("Token验证失败: 用户ID为空"); 
                return false; // 拦截非法连接
            }
        } catch (Exception e) {
            log.error("Token验证失败: {}", e.getMessage(), e); 
            return false;  // 拦截异常连接
        }
    }

    // 握手完成后处理（此处主要用于异常记录）
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                              WebSocketHandler wsHandler, Exception exception) {
        if (exception != null) {
            log.error("WebSocket握手异常: {}", exception.getMessage(), exception); 
        }
    }

    /**
     * 辅助方法：解析URL查询参数
     * @param paramName 要获取的参数名（比如token）
     * @return 参数值或null
     */
    private String getQueryParam(ServerHttpRequest request, String paramName) {
        String query = request.getURI().getQuery();
        if (query == null) return null;

        // 拆分键值对（示例：name=john&age=20）
        for (String pair : query.split("&")) {
            String[] keyValue = pair.split("=");
            if (keyValue.length == 2 && keyValue[0].equals(paramName)) {
                return keyValue[1]; // 返回解码后的值（注意：此处需处理URL编码）
            }
        }
        return null;
    }
}

2. WebSocketConfig.java（配置中心）

WebSocketConfig 负责注册 WebSocket 端点，并配置 WebSocket 拦截器。通过 TokenHandshakeInterceptor 对 WebSocket 握手进行身份验证。

@Configuration
@EnableWebSocket // 启用WebSocket功能
public class WebSocketConfig implements WebSocketConfigurer {

    // 依赖注入自定义的WebSocket处理器
    private final WebSocketNotificationHandler webSocketNotificationHandler;

    public WebSocketConfig(WebSocketNotificationHandler webSocketNotificationHandler) {
        this.webSocketNotificationHandler = webSocketNotificationHandler;
    }

    /**
     * 注册WebSocket端点与配置
     */
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry
            // 指定处理路径（客户端连接地址）
            .addHandler(webSocketNotificationHandler, "/ws/notifications")
            // 允许所有来源（生产环境应指定具体域名）
            .setAllowedOrigins("*")
            // 添加拦截器（顺序重要！）
            .addInterceptors(
                new TokenHandshakeInterceptor(), // 1.先进行Token验证
                new HttpSessionHandshakeInterceptor() // 2.处理HTTP会话（可选）
            );
    }
}

3. WebSocketNotificationHandler.java（消息处理器）

WebSocketNotificationHandler 负责处理客户端的消息，管理 WebSocket 会话，支持任务通知推送。

@Component
public class WebSocketNotificationHandler extends TextWebSocketHandler {

    // 用户会话存储（线程安全Map）
    private static final Map<String, WebSocketSession> userSessions = new ConcurrentHashMap<>();

    /**
     * 当WebSocket连接成功建立时触发
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 从会话属性中获取用户ID（由拦截器存入）
        String userId = (String) session.getAttributes().get("userId");
        if (userId != null) {
            // 将用户与会话绑定（支持后续定向推送）
            userSessions.put(userId, session);
            System.out.println("用户 " + userId + " 连接已建立"); // 建议改用log
        }
    }

    /**
     * 处理客户端发送的消息
     */
    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String userId = (String) session.getAttributes().get("userId");
        System.out.println("收到用户 " + userId + " 的消息：" + message.getPayload());
        // 此处可扩展消息处理逻辑（如消息转发、业务处理等）
    }

    /**
     * 当连接关闭时清理会话
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String userId = (String) session.getAttributes().get("userId");
        if (userId != null) {
            userSessions.remove(userId);
            System.out.println("用户 " + userId + " 已断开连接");
        }
    }

    /**
     * 向指定用户推送任务通知（核心推送方法）
     * @param userId 目标用户ID
     * @param taskData 要发送的任务数据（JSON字符串）
     */
    public void sendTaskNotificationToFrontend(String userId, String taskData) {
        WebSocketSession session = userSessions.get(userId);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(taskData));
                System.out.println("任务通知已发送至用户 " + userId);
            } catch (Exception e) {
                System.err.println("发送失败：" + e.getMessage()); // 需添加重试机制
            }
        }
    }

    // 获取当前所有活跃会话（调试用）
    public Map<String, WebSocketSession> getSessions() {
        return userSessions;
    }
}

总结

TokenHandshakeInterceptor：拦截 WebSocket 握手过程，验证用户身份。

WebSocketConfig：配置 WebSocket 服务端点及拦截器。

WebSocketNotificationHandler：处理 WebSocket 消息、管理会话，支持消息推送。

该架构可以保证 WebSocket 通信的安全性，并且实现了用户消息的精准推送，适合在需要实时通知的场景中使用。