From c10e03727ccced1adc4f5526751745e7d08a297c Mon Sep 17 00:00:00 2001
From: zhongzb <972627721@qq.com>
Date: Sun, 2 Jul 2023 23:04:10 +0800
Subject: [PATCH] =?UTF-8?q?fix:=E6=8F=A1=E6=89=8B=E8=AE=A4=E8=AF=81?=
 =?UTF-8?q?=E4=BC=98=E5=8C=96?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../user/websocket/HttpHeadersHandler.java    |  1 +
 .../NettyWebSocketServerHandler.java          |  6 +++++-
 .../websocket/WebSocketHandshakeHandler.java  | 19 ++++++++++++++++---
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/HttpHeadersHandler.java b/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/HttpHeadersHandler.java
index 6b55b91..30f6cdd 100644
--- a/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/HttpHeadersHandler.java
+++ b/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/HttpHeadersHandler.java
@@ -20,6 +20,7 @@ public class HttpHeadersHandler extends ChannelInboundHandlerAdapter {
                 ip = address.getAddress().getHostAddress();
             }
             NettyUtil.setAttr(ctx.channel(), NettyUtil.IP, ip);
+            ctx.pipeline().remove(this);
         }
         ctx.fireChannelRead(msg);
     }
diff --git a/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/NettyWebSocketServerHandler.java b/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/NettyWebSocketServerHandler.java
index e89030e..3d75069 100644
--- a/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/NettyWebSocketServerHandler.java
+++ b/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/NettyWebSocketServerHandler.java
@@ -1,5 +1,6 @@
 package com.abin.mallchat.custom.user.websocket;
 
+import cn.hutool.core.util.StrUtil;
 import cn.hutool.extra.spring.SpringUtil;
 import cn.hutool.json.JSONUtil;
 import com.abin.mallchat.custom.user.domain.enums.WSReqTypeEnum;
@@ -66,7 +67,10 @@ public class NettyWebSocketServerHandler extends SimpleChannelInboundHandler<Tex
             }
         } else if (evt == WebSocketServerProtocolHandler.ServerHandshakeStateEvent.HANDSHAKE_COMPLETE) {
             getService().connect(ctx.channel());
-            getService().authorize(ctx.channel(), new WSAuthorize(NettyUtil.getAttr(ctx.channel(), NettyUtil.TOKEN)));
+            String token = NettyUtil.getAttr(ctx.channel(), NettyUtil.TOKEN);
+            if (StrUtil.isNotBlank(token)) {
+                getService().authorize(ctx.channel(), new WSAuthorize(token));
+            }
         }
         super.userEventTriggered(ctx, evt);
     }
diff --git a/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/WebSocketHandshakeHandler.java b/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/WebSocketHandshakeHandler.java
index bc1119f..f501602 100644
--- a/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/WebSocketHandshakeHandler.java
+++ b/mallchat-custom-server/src/main/java/com/abin/mallchat/custom/user/websocket/WebSocketHandshakeHandler.java
@@ -1,6 +1,8 @@
 package com.abin.mallchat.custom.user.websocket;
 
 
+import io.netty.channel.ChannelFuture;
+import io.netty.channel.ChannelFutureListener;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.handler.codec.http.FullHttpRequest;
@@ -19,9 +21,20 @@ public class WebSocketHandshakeHandler extends ChannelInboundHandlerAdapter {
             WebSocketServerHandshakerFactory handshakeFactory = new WebSocketServerHandshakerFactory(
                     request.uri(), token, false);
             WebSocketServerHandshaker handshake = handshakeFactory.newHandshaker(request);
-            handshake.handshake(ctx.channel(), request);
-            // 手动触发WebSocket握手状态事件
-            ctx.pipeline().fireUserEventTriggered(WebSocketServerProtocolHandler.ServerHandshakeStateEvent.HANDSHAKE_COMPLETE);
+            final ChannelFuture handshakeFuture = handshake.handshake(ctx.channel(), request);
+            ctx.pipeline().remove(this);
+            handshakeFuture.addListener(new ChannelFutureListener() {
+                @Override
+                public void operationComplete(ChannelFuture future) {
+                    if (!future.isSuccess()) {
+                        ctx.fireExceptionCaught(future.cause());
+                    } else {
+                        // 手动触发WebSocket握手状态事件
+                        ctx.fireUserEventTriggered(
+                                WebSocketServerProtocolHandler.ServerHandshakeStateEvent.HANDSHAKE_COMPLETE);
+                    }
+                }
+            });
         } else {
             super.channelRead(ctx, msg);
         }