From f929513310c5dbbcef8f95a1daf20d8ff3c1f0c1 Mon Sep 17 00:00:00 2001
From: 疯狂的狮子Li <15040126243@163.com>
Date: 星期二, 02 七月 2024 12:00:08 +0800
Subject: [PATCH] update 优化 webscoket 配置与异常拦截

---
 ruoyi-common/ruoyi-common-web/src/main/java/org/dromara/common/web/config/UndertowConfig.java                            |    6 +-----
 ruoyi-common/ruoyi-common-websocket/src/main/java/org/dromara/common/websocket/interceptor/PlusWebSocketInterceptor.java |   35 ++++++++++++++++++++---------------
 2 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/ruoyi-common/ruoyi-common-web/src/main/java/org/dromara/common/web/config/UndertowConfig.java b/ruoyi-common/ruoyi-common-web/src/main/java/org/dromara/common/web/config/UndertowConfig.java
index 0f78928..ac50dd2 100644
--- a/ruoyi-common/ruoyi-common-web/src/main/java/org/dromara/common/web/config/UndertowConfig.java
+++ b/ruoyi-common/ruoyi-common-web/src/main/java/org/dromara/common/web/config/UndertowConfig.java
@@ -16,15 +16,11 @@
 @AutoConfiguration
 public class UndertowConfig implements WebServerFactoryCustomizer<UndertowServletWebServerFactory> {
 
-    /**
-     * 璁剧疆 Undertow 鐨� websocket 缂撳啿姹�
-     */
     @Override
     public void customize(UndertowServletWebServerFactory factory) {
-        // 榛樿涓嶇洿鎺ュ垎閰嶅唴瀛� 濡傛灉椤圭洰涓娇鐢ㄤ簡 websocket 寤鸿鐩存帴鍒嗛厤
         factory.addDeploymentInfoCustomizers(deploymentInfo -> {
             WebSocketDeploymentInfo webSocketDeploymentInfo = new WebSocketDeploymentInfo();
-            webSocketDeploymentInfo.setBuffers(new DefaultByteBufferPool(false, 512));
+            webSocketDeploymentInfo.setBuffers(new DefaultByteBufferPool(true, 1024));
             deploymentInfo.addServletContextAttribute("io.undertow.websockets.jsr.WebSocketDeploymentInfo", webSocketDeploymentInfo);
             // 浣跨敤铏氭嫙绾跨▼
             if (SpringUtils.isVirtual()) {
diff --git a/ruoyi-common/ruoyi-common-websocket/src/main/java/org/dromara/common/websocket/interceptor/PlusWebSocketInterceptor.java b/ruoyi-common/ruoyi-common-websocket/src/main/java/org/dromara/common/websocket/interceptor/PlusWebSocketInterceptor.java
index 1ad8ad6..c70e377 100644
--- a/ruoyi-common/ruoyi-common-websocket/src/main/java/org/dromara/common/websocket/interceptor/PlusWebSocketInterceptor.java
+++ b/ruoyi-common/ruoyi-common-websocket/src/main/java/org/dromara/common/websocket/interceptor/PlusWebSocketInterceptor.java
@@ -35,23 +35,28 @@
      */
     @Override
     public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) {
-        // 妫�鏌ユ槸鍚︾櫥褰� 鏄惁鏈塼oken
-        LoginUser loginUser = LoginHelper.getLoginUser();
+        try {
+            // 妫�鏌ユ槸鍚︾櫥褰� 鏄惁鏈塼oken
+            LoginUser loginUser = LoginHelper.getLoginUser();
 
-        // 瑙e喅 ws 涓嶈蛋 mvc 鎷︽埅鍣ㄩ棶棰�(cloud 鐗堟湰涓嶅彈褰卞搷)
-        // 妫�鏌� header 涓� param 閲岀殑 clientid 涓� token 閲岀殑鏄惁涓�鑷�
-        String headerCid = ServletUtils.getRequest().getHeader(LoginHelper.CLIENT_KEY);
-        String paramCid = ServletUtils.getParameter(LoginHelper.CLIENT_KEY);
-        String clientId = StpUtil.getExtra(LoginHelper.CLIENT_KEY).toString();
-        if (!StringUtils.equalsAny(clientId, headerCid, paramCid)) {
-            // token 鏃犳晥
-            throw NotLoginException.newInstance(StpUtil.getLoginType(),
-                "-100", "瀹㈡埛绔疘D涓嶵oken涓嶅尮閰�",
-                StpUtil.getTokenValue());
+            // 瑙e喅 ws 涓嶈蛋 mvc 鎷︽埅鍣ㄩ棶棰�(cloud 鐗堟湰涓嶅彈褰卞搷)
+            // 妫�鏌� header 涓� param 閲岀殑 clientid 涓� token 閲岀殑鏄惁涓�鑷�
+            String headerCid = ServletUtils.getRequest().getHeader(LoginHelper.CLIENT_KEY);
+            String paramCid = ServletUtils.getParameter(LoginHelper.CLIENT_KEY);
+            String clientId = StpUtil.getExtra(LoginHelper.CLIENT_KEY).toString();
+            if (!StringUtils.equalsAny(clientId, headerCid, paramCid)) {
+                // token 鏃犳晥
+                throw NotLoginException.newInstance(StpUtil.getLoginType(),
+                    "-100", "瀹㈡埛绔疘D涓嶵oken涓嶅尮閰�",
+                    StpUtil.getTokenValue());
+            }
+
+            attributes.put(LOGIN_USER_KEY, loginUser);
+            return true;
+        } catch (NotLoginException e) {
+            log.error("WebSocket 璁よ瘉澶辫触'{}',鏃犳硶璁块棶绯荤粺璧勬簮", e.getMessage());
+            return false;
         }
-
-        attributes.put(LOGIN_USER_KEY, loginUser);
-        return true;
     }
 
     /**

--
Gitblit v1.9.3