WebSocket 自定義安全校驗(yàn)優(yōu)化實(shí)踐

作者：一安 2025-05-15 07:11:51

在以Spring Boot與Vue搭建的應(yīng)用體系里，WebSocket作為實(shí)現(xiàn)前后端實(shí)時(shí)交互的得力助手，被廣泛運(yùn)用。然而隨著網(wǎng)絡(luò)安全形勢(shì)日益嚴(yán)峻，為WebSocket交互筑牢安全防線變得至關(guān)重要。

前言

實(shí)現(xiàn)

傳統(tǒng)

@Slf4j
@Component
@ServerEndpoint("/websocket/link/{userId}")
public class OldWebSocketService {
    // 用于存儲(chǔ)在線用戶的會(huì)話，使用ConcurrentHashMap確保線程安全
    private static final Map<String, Session> onlineSessions = new ConcurrentHashMap<>();

    @OnOpen
    public void handleOpen(Session session, @PathParam("userId") String userId) {
        onlineSessions.put(userId, session);
        log.info("用戶ID為 {} 的用戶已連接，當(dāng)前在線用戶數(shù): {}", userId, onlineSessions.size());
        broadcastMessage("系統(tǒng)提示：有新用戶加入");
    }

    @OnMessage
    public void handleMessage(String message, Session session, @PathParam("userId") String userId) {
        log.info("服務(wù)端收到用戶ID為 {} 的消息: {}", userId, message);
        JSONObject jsonMessage = JSON.parseObject(message);
        String targetUserId = jsonMessage.getString("to");
        String content = jsonMessage.getString("content");

        Session targetSession = onlineSessions.get(targetUserId);
        if (targetSession != null) {
            JSONObject responseMessage = new JSONObject();
            responseMessage.put("from", userId);
            responseMessage.put("content", content);
            sendMessage(responseMessage.toJSONString(), targetSession);
            log.info("向用戶ID為 {} 發(fā)送消息: {}", targetUserId, responseMessage.toJSONString());
        } else {
            log.info("未能找到用戶ID為 {} 的會(huì)話，消息發(fā)送失敗", targetUserId);
        }
    }

    private void sendMessage(String message, Session targetSession) {
        try {
            log.info("服務(wù)端向客戶端[{}]發(fā)送消息: {}", targetSession.getId(), message);
            targetSession.getBasicRemote().sendText(message);
        } catch (Exception e) {
            log.error("服務(wù)端向客戶端發(fā)送消息失敗", e);
        }
    }

    private void broadcastMessage(String message) {
        for (Session session : onlineSessions.values()) {
            sendMessage(message, session);
        }
    }

    @OnClose
    public void handleClose(Session session, @PathParam("userId") String userId) {
        onlineSessions.remove(userId);
        log.info("用戶ID為 {} 的連接已關(guān)閉，當(dāng)前在線用戶數(shù): {}", userId, onlineSessions.size());
    }
}

打造安全加固的 WebSocket 體系

實(shí)現(xiàn)流程

客戶端發(fā)起連接：客戶端向/secure-websocket地址發(fā)起WebSocket連接請(qǐng)求，在請(qǐng)求頭中攜帶Authorization（即token）和Unique-User-Key（用戶唯一標(biāo)識(shí)）。
攔截器校驗(yàn)：服務(wù)端的SecurityInterceptor攔截請(qǐng)求，獲取并校驗(yàn)token。若token無(wú)效，阻止握手；若有效，則將用戶唯一標(biāo)識(shí)存入attributes。
連接建立：若攔截器允許握手，連接成功建立。EnhancedWebSocketHandler的afterConnectionEstablished方法被調(diào)用，獲取用戶唯一標(biāo)識(shí)并存儲(chǔ)會(huì)話。
消息交互：客戶端和服務(wù)端進(jìn)行消息收發(fā)。EnhancedWebSocketHandler的handleMessage方法處理消息前，先校驗(yàn)消息來(lái)源的用戶唯一標(biāo)識(shí)是否合法。
連接關(guān)閉：連接關(guān)閉時(shí)，EnhancedWebSocketHandler的afterConnectionClosed方法被調(diào)用，移除對(duì)應(yīng)會(huì)話。

自定義 WebSocket 處理器

@Slf4j
@Component
public class EnhancedWebSocketHandler implements WebSocketHandler {
    // 存儲(chǔ)用戶標(biāo)識(shí)與會(huì)話的映射關(guān)系，保證線程安全
    private static final Map<String, WebSocketSession> userSessionMap = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String userKey = (String) session.getAttributes().get("uniqueUserKey");
        session.sendMessage(new TextMessage("用戶:"+userKey+" 認(rèn)證成功"));

        log.info("WebSocket連接已建立，用戶唯一標(biāo)識(shí): {}, 會(huì)話ID: {}", userKey, session.getId());
        userSessionMap.put(userKey, session);
        log.info("新用戶連接，當(dāng)前在線用戶數(shù): {}", userSessionMap.size());
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        JSONObject json = JSONObject.parseObject((String) message.getPayload());
        if(!userSessionMap.containsKey(json.getString("to"))){
            session.sendMessage(new TextMessage("接收用戶不存在!!!"));
            return;
        }
        String userKey = (String) session.getAttributes().get("uniqueUserKey");
        if (!userSessionMap.containsKey(userKey)) {
            session.sendMessage(new TextMessage("發(fā)送用戶不存在!!!"));
            return;
        }
        session.sendMessage(new TextMessage("收到 over"));
        log.info("消息接收成功，內(nèi)容: {}", message.getPayload());
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        String userKey = (String) session.getAttributes().get("uniqueUserKey");
        if (userSessionMap.containsKey(userKey)) {
            log.error("WebSocket傳輸出現(xiàn)錯(cuò)誤，用戶標(biāo)識(shí): {}, 錯(cuò)誤信息: {}", userKey, exception.getMessage());
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        String userKey = (String) session.getAttributes().get("uniqueUserKey");
        log.info("WebSocket連接已關(guān)閉，會(huì)話ID: {}, 關(guān)閉狀態(tài): {}", session.getId(), closeStatus);
        userSessionMap.remove(userKey);
    }

    @Override
    public boolean supportsPartialMessages() {
        returntrue;
    }

    public void sendMessage(String message, WebSocketSession targetSession) {
        try {
            log.info("服務(wù)端向客戶端[{}]發(fā)送消息: {}", targetSession.getId(), message);
            targetSession.sendMessage(new TextMessage(message));
        } catch (Exception e) {
            log.error("服務(wù)端向客戶端發(fā)送消息失敗", e);
        }
    }

    public void broadcastMessage(String message) {
        for (WebSocketSession session : userSessionMap.values()) {
            sendMessage(message, session);
        }
    }
}

自定義 WebSocket 攔截器

@Slf4j
@Component
public class SecurityInterceptor implements HandshakeInterceptor {
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler webSocketHandler, Map<String, Object> attributes) throws Exception {
        // 獲取 HttpServletRequest 對(duì)象
        HttpServletRequest rs=((ServletServerHttpRequest) request).getServletRequest();
        String token = rs.getParameter("Authorization");
        log.info("攔截器獲取到的令牌: {}", token);
        if (token == null ||!isValidToken(token)) {
            log.warn("無(wú)效的令牌，拒絕WebSocket連接");
            returnfalse;
        }
        String userKey = rs.getParameter("UniqueUserKey");
        attributes.put("uniqueUserKey", userKey);
        returntrue;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler webSocketHandler, Exception exception) {
        // 可在此處添加握手成功后的處理邏輯
    }

    private boolean isValidToken(String token) {
        // 實(shí)際應(yīng)用中應(yīng)包含復(fù)雜的令牌驗(yàn)證邏輯，如JWT驗(yàn)證
        // 此處僅為示例，簡(jiǎn)單判斷令牌是否為"validToken"
        return"1234".equals(token);
    }
}

WebSocket 配置類

@Configuration
@EnableWebSocket
public class WebSocketSecurityConfig implements WebSocketConfigurer {
    private final EnhancedWebSocketHandler enhancedWebSocketHandler;
    private final SecurityInterceptor securityInterceptor;

    public WebSocketSecurityConfig(EnhancedWebSocketHandler enhancedWebSocketHandler, SecurityInterceptor securityInterceptor) {
        this.enhancedWebSocketHandler = enhancedWebSocketHandler;
        this.securityInterceptor = securityInterceptor;
    }

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(enhancedWebSocketHandler, "/secure-websocket")
               .setAllowedOrigins("*")
               .addInterceptors(securityInterceptor);
    }

    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

示例頁(yè)面

<!DOCTYPE html>
<html lang="zh-CN">

<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>WebSocket 認(rèn)證交互頁(yè)面</title>
    <link  rel="stylesheet" >
    <style>
        body {
            font-family: Arial, sans-serif;
        }

        #authSection {
            margin-bottom: 10px;
        }

        #tokenInput,
        #userKeyInput {
            width: 200px;
            padding: 8px;
            margin-right: 10px;
        }

        #authButton {
            padding: 8px 16px;
        }

        #messageInput {
            width: 300px;
            padding: 8px;
            margin-right: 10px;
        }

        #targetUserInput {
            width: 200px;
            padding: 8px;
            margin-right: 10px;
        }

        #sendButton {
            padding: 8px 16px;
        }

        #messageList {
            list-style-type: none;
            padding: 0;
        }

        #messageList li {
            margin: 8px 0;
            border: 1px solid #ccc;
            padding: 8px;
            border-radius: 4px;
        }
    </style>
</head>

<body>
<h2>WebSocket 認(rèn)證交互頁(yè)面</h2>
<div id="authSection">
    <label for="tokenInput">輸入認(rèn)證 Token:</label>
    <input type="text" id="tokenInput" placeholder="請(qǐng)輸入認(rèn)證 Token">
    <label for="userKeyInput">輸入用戶唯一標(biāo)識(shí):</label>
    <input type="text" id="userKeyInput" placeholder="請(qǐng)輸入用戶唯一標(biāo)識(shí)">
    <button id="authButton">認(rèn)證并連接</button>
</div>
<input type="text" id="messageInput" placeholder="請(qǐng)輸入要發(fā)送的消息">
<input type="text" id="targetUserInput" placeholder="請(qǐng)輸入接收消息的用戶標(biāo)識(shí)">
<button id="sendButton">發(fā)送消息</button>
<ul id="messageList"></ul>
<script>
    let socket;

    document.getElementById('authButton').addEventListener('click', function () {
        const token = document.getElementById('tokenInput').value;
        const userKey = document.getElementById('userKeyInput').value;
        if (token.trim() === '' || userKey.trim() === '') {
            console.error('Token 或用戶唯一標(biāo)識(shí)不能為空');
            return;
        }

        const socketUrl = 'ws://localhost:8080/secure-websocket?Authorization='+token+'&UniqueUserKey='+userKey ;
        socket = new WebSocket(socketUrl);

        socket.onopen = function () {
             console.log('WebSocket 連接已打開(kāi)');
        };

        socket.onmessage = function (event) {
            const messageItem = document.createElement('li');
            messageItem.textContent = event.data;
            document.getElementById('messageList').appendChild(messageItem);
        };

        socket.onclose = function () {
            console.log('WebSocket 連接已關(guān)閉');
        };

        socket.onerror = function (error) {
            console.error('WebSocket 發(fā)生錯(cuò)誤:', error);
        };
    });

    document.getElementById('sendButton').addEventListener('click', function () {
        if (!socket || socket.readyState!== WebSocket.OPEN) {
            console.error('WebSocket 連接未建立或已關(guān)閉');
            return;
        }
        const message = document.getElementById('messageInput').value;
        const targetUser = document.getElementById('targetUserInput').value;
        if (message.trim() === '' || targetUser.trim() === '') {
            return;
        }
        const messageObj = {
            to: targetUser,
            content: message
        };
        socket.send(JSON.stringify(messageObj));
        document.getElementById('messageInput').value = '';
        document.getElementById('targetUserInput').value = '';
    });
</script>
</body>

</html>