在《Websocket在Java中的實踐——最小可行案例》一文中，我們看到如何用最簡單的方式實現Websocket通信。本文中，我們將介紹如何在握手前后進行干涉，以定制一些特殊需求。
在《Websocket在Java中的實踐——最小可行案例》的基礎上，我們希望建立“用戶”的概念，即不同用戶有自己的用戶名。用戶只能收到別人發的消息，而不能收到自己的消息。
這就要求我們服務可以處理ws://localhost:8080/websocket/{uid}這樣的請求。而對于uid不存在或者不合法的場景，就要拒絕連接。

依賴

在pom.xml中新增

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

握手攔截器

在綁定接口時，我們可以通過addInterceptors方法給WebSocketHandlerRegistry指定一個握手攔截器。

package com.nyctlc.withparam.config;import java.util.Map;import org.springframework.context.annotation.Configuration;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.lang.Nullable;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.HandshakeInterceptor;import jakarta.servlet.http.HttpSession;@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {@Overridepublic void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {registry.addHandler(new com.nyctlc.withparam.handler.WebSocketHandler(), "/websocket/{uid}").setAllowedOrigins("*").addInterceptors(handshakeInterceptor());}

這個攔截器需要重載兩個方法：beforeHandshake和afterHandshake。它們分別在握手前后被調用。
我們的需求要求我們在握手之前獲取uid。如果uid不存在或者不合法，就會拒絕連接；如果合法，則將其保存到session的屬性字段中。

private HandshakeInterceptor handshakeInterceptor() {return new HandshakeInterceptor() {@Overridepublic boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {if (request instanceof ServletServerHttpRequest) {String path = request.getURI().getPath();String prefix = "/websocket/";String uid = path.substring(path.indexOf(prefix) + prefix.length());if (uid.isEmpty()) {return false;}if (uid.contains("/")) {return false;}ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;HttpSession session = servletRequest.getServletRequest().getSession();attributes.put("sessionId", session.getId());attributes.put("uid", uid);}return true;}@Overridepublic void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,WebSocketHandler wsHandler, @Nullable Exception exception) {}};}}

消息處理

這次我們在handleTextMessage方法中，判斷接收消息的Session的Attributes中的uid是否和Session集合中的一致。這個uid是在上一個步驟中，通過握手攔截器設置的。
如果一致，說明它們來自同一個用戶，則不將該消息發送回自己。

package com.nyctlc.withparam.handler;import java.io.IOException;
import java.util.HashSet;
import java.util.Set;import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;public class WebSocketHandler extends TextWebSocketHandler {private static final Set<WebSocketSession> sessions = new HashSet<>();@Overridepublic void afterConnectionEstablished(WebSocketSession session) throws Exception {sessions.add(session);}@Overrideprotected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {for (WebSocketSession webSocketSession : sessions) {String uid = session.getAttributes().get("uid").toString();if (webSocketSession.isOpen()) {String uidInSession = webSocketSession.getAttributes().get("uid").toString();if (uid.equals(uidInSession)) {continue;}try {webSocketSession.sendMessage(message);} catch (IOException e) {e.printStackTrace();}}}}
}

測試

參考資料

• https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/web/socket/server/HandshakeInterceptor.html