增加 spring websocket 示例

lab-25/lab-websocket-25-02/src/main/java/cn/iocoder/springboot/lab25/springwebsocket/config/WebSocketConfiguration.java

@@ -1,16 +1,32 @@
 package cn.iocoder.springboot.lab25.springwebsocket.config;
 
+import cn.iocoder.springboot.lab25.springwebsocket.websocket.DemoWebSocketHandler;
+import cn.iocoder.springboot.lab25.springwebsocket.websocket.DemoWebSocketShakeInterceptor;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.web.socket.server.standard.ServerEndpointExporter;
+import org.springframework.web.socket.config.annotation.EnableWebSocket;
+import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
+import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
 
 @Configuration
-// @EnableWebSocket // 无需添加该注解，因为我们并不是使用 Spring WebSocket
-public class WebSocketConfiguration {
+@EnableWebSocket // 开启 Spring WebSocket
+public class WebSocketConfiguration implements WebSocketConfigurer {
+
+    @Override
+    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
+        registry.addHandler(this.webSocketHandler(), "/") // 配置处理器
+                .addInterceptors(new DemoWebSocketShakeInterceptor()) // 配置拦截器
+                .setAllowedOrigins("*"); // 解决跨域问题
+    }
 
     @Bean
-    public ServerEndpointExporter serverEndpointExporter() {
-        return new ServerEndpointExporter();
+    public DemoWebSocketHandler webSocketHandler() {
+        return new DemoWebSocketHandler();
+    }
+
+    @Bean
+    public DemoWebSocketShakeInterceptor webSocketShakeInterceptor() {
+        return new DemoWebSocketShakeInterceptor();
     }
 
 }

lab-25/lab-websocket-25-02/src/main/java/cn/iocoder/springboot/lab25/springwebsocket/websocket/DemoWebSocketHandler.java

@@ -12,40 +12,36 @@ import org.springframework.aop.framework.AopProxyUtils;
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.ApplicationContext;
-import org.springframework.stereotype.Controller;
-import org.springframework.util.CollectionUtils;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.TextMessage;
+import org.springframework.web.socket.WebSocketSession;
+import org.springframework.web.socket.handler.TextWebSocketHandler;
 
-import javax.websocket.*;
-import javax.websocket.server.ServerEndpoint;
 import java.lang.reflect.ParameterizedType;
 import java.lang.reflect.Type;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
-@Controller
-@ServerEndpoint("/")
-public class WebsocketServerEndpoint implements InitializingBean {
+public class DemoWebSocketHandler extends TextWebSocketHandler implements InitializingBean {
 
     private Logger logger = LoggerFactory.getLogger(getClass());
 
     /**
      * 消息类型与 MessageHandler 的映射
      *
-     * 注意，这里设置成静态变量。虽然说 WebsocketServerEndpoint 是单例，但是 Spring Boot 还是会为每个 WebSocket 创建一个 WebsocketServerEndpoint Bean 。
+     * 无需设置成静态变量
      */
-    private static final Map<String, MessageHandler> HANDLERS = new HashMap<>();
+    private final Map<String, MessageHandler> HANDLERS = new HashMap<>();
 
     @Autowired
     private ApplicationContext applicationContext;
 
-    @OnOpen
-    public void onOpen(Session session, EndpointConfig config) {
-        logger.info("[onOpen][session({}) 接入]", session);
+    @Override // 对应 open 事件
+    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
+        logger.info("[afterConnectionEstablished][session({}) 接入]", session);
         // 解析 accessToken
-        List<String> accessTokenValues = session.getRequestParameterMap().get("accessToken");
-        String accessToken = !CollectionUtils.isEmpty(accessTokenValues) ? accessTokenValues.get(0) : null;
+        String accessToken = (String) session.getAttributes().get("accessToken");
         // 创建 AuthRequest 消息类型
         AuthRequest authRequest = new AuthRequest().setAccessToken(accessToken);
         // 获得消息处理器
@@ -57,12 +53,12 @@ public class WebsocketServerEndpoint implements InitializingBean {
         messageHandler.execute(session, authRequest);
     }
 
-    @OnMessage
-    public void onMessage(Session session, String message) {
-        logger.info("[onOpen][session({}) 接收到一条消息({})]", session, message); // 生产环境下，请设置成 debug 级别
+    @Override // 对应 message 事件
+    public void handleTextMessage(WebSocketSession session, TextMessage textMessage) throws Exception {
+        logger.info("[handleMessage][session({}) 接收到一条消息({})]", session, textMessage); // 生产环境下，请设置成 debug 级别
         try {
             // 获得消息类型
-            JSONObject jsonMessage = JSON.parseObject(message);
+            JSONObject jsonMessage = JSON.parseObject(textMessage.getPayload());
             String messageType = jsonMessage.getString("type");
             // 获得消息处理器
             MessageHandler messageHandler = HANDLERS.get(messageType);
@@ -80,15 +76,15 @@ public class WebsocketServerEndpoint implements InitializingBean {
         }
     }
 
-    @OnClose
-    public void onClose(Session session, CloseReason closeReason) {
-        logger.info("[onClose][session({}) 连接关闭。关闭原因是({})}]", session, closeReason);
+    @Override // 对应 close 事件
+    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
+        logger.info("[afterConnectionClosed][session({}) 连接关闭。关闭原因是({})}]", session, status);
         WebSocketUtil.removeSession(session);
     }
 
-    @OnError
-    public void onError(Session session, Throwable throwable) {
-        logger.info("[onClose][session({}) 发生异常]", session, throwable);
+    @Override // 对应 error 事件
+    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
+        logger.info("[handleTransportError][session({}) 发生异常]", session, exception);
     }
 
     @Override

lab-25/lab-websocket-25-02/src/main/java/cn/iocoder/springboot/lab25/springwebsocket/websocket/DemoWebSocketShakeInterceptor.java

@@ -0,0 +1,30 @@
+package cn.iocoder.springboot.lab25.springwebsocket.websocket;
+
+import org.springframework.http.server.ServerHttpRequest;
+import org.springframework.http.server.ServerHttpResponse;
+import org.springframework.http.server.ServletServerHttpRequest;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;
+
+import java.util.Map;
+
+/**
+ * 自定义 HttpSessionHandshakeInterceptor 拦截器
+ *
+ * 因为 WebSocketSession 无法获得 ws 地址上的请求参数，所以只好通过该拦截器，获得 accessToken 请求参数，设置到 attributes 中
+ */
+public class DemoWebSocketShakeInterceptor extends HttpSessionHandshakeInterceptor {
+
+    @Override // 拦截 Handshake 事件
+    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
+                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
+        // 获得 accessToken
+        if (request instanceof ServletServerHttpRequest) {
+            ServletServerHttpRequest serverRequest = (ServletServerHttpRequest) request;
+            attributes.put("accessToken", serverRequest.getServletRequest().getParameter("accessToken"));
+        }
+        // 调用父方法，继续执行逻辑
+        return super.beforeHandshake(request, response, wsHandler, attributes);
+    }
+
+}