WebSocket 自定義安全校驗優化實踐
作者：一安
在以Spring Boot與Vue搭建的應用體系里，WebSocket作為實現前后端實時交互的得力助手，被廣泛運用。然而隨著網絡安全形勢日益嚴峻，為WebSocket交互筑牢安全防線變得至關重要。
前言
在以Spring Boot與Vue搭建的應用體系里，WebSocket作為實現前后端實時交互的得力助手，被廣泛運用。然而隨著網絡安全形勢日益嚴峻，為WebSocket交互筑牢安全防線變得至關重要。
實現
傳統
@Slf4j
@Component
@ServerEndpoint("/websocket/link/{userId}")
public class OldWebSocketService {
// 用于存儲在線用戶的會話，使用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為 {} 的用戶已連接，當前在線用戶數: {}", userId, onlineSessions.size());
broadcastMessage("系統提示：有新用戶加入");
}
@OnMessage
public void handleMessage(String message, Session session, @PathParam("userId") String userId) {
log.info("服務端收到用戶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為 {} 發送消息: {}", targetUserId, responseMessage.toJSONString());
} else {
log.info("未能找到用戶ID為 {} 的會話，消息發送失敗", targetUserId);
}
}
private void sendMessage(String message, Session targetSession) {
try {
log.info("服務端向客戶端[{}]發送消息: {}", targetSession.getId(), message);
targetSession.getBasicRemote().sendText(message);
} catch (Exception e) {
log.error("服務端向客戶端發送消息失敗", 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為 {} 的連接已關閉，當前在線用戶數: {}", userId, onlineSessions.size());
}
}打造安全加固的 WebSocket 體系
實現流程
- 客戶端發起連接：客戶端向/secure-websocket地址發起WebSocket連接請求，在請求頭中攜帶Authorization（即token）和Unique-User-Key（用戶唯一標識）。
- 攔截器校驗：服務端的SecurityInterceptor攔截請求，獲取并校驗token。若token無效，阻止握手；若有效，則將用戶唯一標識存入attributes。
- 連接建立：若攔截器允許握手，連接成功建立。EnhancedWebSocketHandler的afterConnectionEstablished方法被調用，獲取用戶唯一標識并存儲會話。
- 消息交互：客戶端和服務端進行消息收發。EnhancedWebSocketHandler的handleMessage方法處理消息前，先校驗消息來源的用戶唯一標識是否合法。
- 連接關閉：連接關閉時，EnhancedWebSocketHandler的afterConnectionClosed方法被調用，移除對應會話。
自定義 WebSocket 處理器
@Slf4j
@Component
public class EnhancedWebSocketHandler implements WebSocketHandler {
// 存儲用戶標識與會話的映射關系，保證線程安全
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+" 認證成功"));
log.info("WebSocket連接已建立，用戶唯一標識: {}, 會話ID: {}", userKey, session.getId());
userSessionMap.put(userKey, session);
log.info("新用戶連接，當前在線用戶數: {}", 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("發送用戶不存在!!!"));
return;
}
session.sendMessage(new TextMessage("收到 over"));
log.info("消息接收成功，內容: {}", 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傳輸出現錯誤，用戶標識: {}, 錯誤信息: {}", userKey, exception.getMessage());
}
}
@Override
public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
String userKey = (String) session.getAttributes().get("uniqueUserKey");
log.info("WebSocket連接已關閉，會話ID: {}, 關閉狀態: {}", session.getId(), closeStatus);
userSessionMap.remove(userKey);
}
@Override
public boolean supportsPartialMessages() {
returntrue;
}
public void sendMessage(String message, WebSocketSession targetSession) {
try {
log.info("服務端向客戶端[{}]發送消息: {}", targetSession.getId(), message);
targetSession.sendMessage(new TextMessage(message));
} catch (Exception e) {
log.error("服務端向客戶端發送消息失敗", 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 對象
HttpServletRequest rs=((ServletServerHttpRequest) request).getServletRequest();
String token = rs.getParameter("Authorization");
log.info("攔截器獲取到的令牌: {}", token);
if (token == null ||!isValidToken(token)) {
log.warn("無效的令牌，拒絕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) {
// 實際應用中應包含復雜的令牌驗證邏輯，如JWT驗證
// 此處僅為示例，簡單判斷令牌是否為"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();
}
}示例頁面
<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>WebSocket 認證交互頁面</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 認證交互頁面</h2>
<div id="authSection">
<label for="tokenInput">輸入認證 Token:</label>
<input type="text" id="tokenInput" placeholder="請輸入認證 Token">
<label for="userKeyInput">輸入用戶唯一標識:</label>
<input type="text" id="userKeyInput" placeholder="請輸入用戶唯一標識">
<button id="authButton">認證并連接</button>
</div>
<input type="text" id="messageInput" placeholder="請輸入要發送的消息">
<input type="text" id="targetUserInput" placeholder="請輸入接收消息的用戶標識">
<button id="sendButton">發送消息</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 或用戶唯一標識不能為空');
return;
}
const socketUrl = 'ws://localhost:8080/secure-websocket?Authorization='+token+'&UniqueUserKey='+userKey ;
socket = new WebSocket(socketUrl);
socket.onopen = function () {
console.log('WebSocket 連接已打開');
};
socket.onmessage = function (event) {
const messageItem = document.createElement('li');
messageItem.textContent = event.data;
document.getElementById('messageList').appendChild(messageItem);
};
socket.onclose = function () {
console.log('WebSocket 連接已關閉');
};
socket.onerror = function (error) {
console.error('WebSocket 發生錯誤:', error);
};
});
document.getElementById('sendButton').addEventListener('click', function () {
if (!socket || socket.readyState!== WebSocket.OPEN) {
console.error('WebSocket 連接未建立或已關閉');
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>測試
圖片
責任編輯：武曉燕
來源：
一安未來
