From f929513310c5dbbcef8f95a1daf20d8ff3c1f0c1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=96=AF=E7=8B=82=E7=9A=84=E7=8B=AE=E5=AD=90Li?=
 <15040126243@163.com>
Date: Tue, 2 Jul 2024 12:00:08 +0800
Subject: [PATCH] =?UTF-8?q?update=20=E4=BC=98=E5=8C=96=20webscoket=20?=
 =?UTF-8?q?=E9=85=8D=E7=BD=AE=E4=B8=8E=E5=BC=82=E5=B8=B8=E6=8B=A6=E6=88=AA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../common/web/config/UndertowConfig.java     |  6 +---
 .../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 0f7892861..ac50dd29a 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 @@ import org.springframework.core.task.VirtualThreadTaskExecutor;
 @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 1ad8ad665..c70e37729 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 @@ public class PlusWebSocketInterceptor implements HandshakeInterceptor {
      */
     @Override
     public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) {
-        // 检查是否登录 是否有token
-        LoginUser loginUser = LoginHelper.getLoginUser();
+        try {
+            // 检查是否登录 是否有token
+            LoginUser loginUser = LoginHelper.getLoginUser();
 
-        // 解决 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", "客户端ID与Token不匹配",
-                StpUtil.getTokenValue());
+            // 解决 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", "客户端ID与Token不匹配",
+                    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;
     }
 
     /**