From 1e4312fcb1a7ff969f38af302857ecb0572041ad Mon Sep 17 00:00:00 2001
From: aoshiguchen <1052045476@qq.com>
Date: Fri, 7 Apr 2023 10:48:35 +0800
Subject: [PATCH] =?UTF-8?q?=E5=9F=9F=E5=90=8D=E6=98=A0=E5=B0=84=E9=80=BB?=
 =?UTF-8?q?=E8=BE=91=E4=BC=98=E5=8C=96=EF=BC=8C=E6=94=AF=E6=8C=81websocket?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../server/proxy/core/HttpProxy.java          | 54 ++++++++++---------
 1 file changed, 29 insertions(+), 25 deletions(-)

diff --git a/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/core/HttpProxy.java b/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/core/HttpProxy.java
index 2f36f732..ba6d4a9c 100644
--- a/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/core/HttpProxy.java
+++ b/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/core/HttpProxy.java
@@ -72,7 +72,32 @@ public class HttpProxy implements EventListener<AppLoadEndEvent> {
                 ctx.channel().close();
                 return;
             }
-            String host = getHost(byteBuf);
+
+            byte[] bytes = new byte[byteBuf.readableBytes()];
+            byteBuf.readBytes(bytes);
+            byteBuf.resetReaderIndex();
+            ProxyAttachment proxyAttachment = new ProxyAttachment(ctx.channel(), bytes, (channel, buf) -> {
+                Channel proxyChannel = channel.attr(Constants.NEXT_CHANNEL).get();
+                if (null == proxyChannel) {
+                    // 该端口还没有代理客户端
+                    ctx.channel().close();
+                    return;
+                }
+
+                proxyChannel.writeAndFlush(ProxyMessage.buildTransferMessage(ProxyUtil.getVisitorIdByChannel(channel), bytes));
+
+                // 增加流量计数
+                VisitorChannelAttachInfo visitorChannelAttachInfo = ProxyUtil.getAttachInfo(channel);
+                Solon.context().getBean(FlowReportService.class).addWriteByte(visitorChannelAttachInfo.getLicenseId(), bytes.length);
+            });
+
+            String visitorId = ProxyUtil.getVisitorIdByChannel(ctx.channel());
+            if (StringUtils.isNotBlank(visitorId)) {
+                proxyAttachment.execute();
+                return;
+            }
+
+            String host = getHost(bytes);
             if (StringUtils.isBlank(host)) {
                 ctx.channel().close();
                 return;
@@ -102,26 +127,8 @@ public class HttpProxy implements EventListener<AppLoadEndEvent> {
                 return;
             }
 
-            String visitorId = ProxyUtil.newVisitorId();
-            Channel visitorChannel = ctx.channel();
-
-            byte[] bytes = new byte[byteBuf.readableBytes()];
-            byteBuf.readBytes(bytes);
-            ProxyAttachment proxyAttachment = new ProxyAttachment(ctx.channel(), bytes, (channel, buf) -> {
-                Channel proxyChannel = channel.attr(Constants.NEXT_CHANNEL).get();
-                if (null == proxyChannel) {
-                    // 该端口还没有代理客户端
-                    ctx.channel().close();
-                    return;
-                }
-                proxyChannel.writeAndFlush(ProxyMessage.buildTransferMessage(visitorId, bytes));
-
-                // 增加流量计数
-                VisitorChannelAttachInfo visitorChannelAttachInfo = ProxyUtil.getAttachInfo(visitorChannel);
-                Solon.context().getBean(FlowReportService.class).addWriteByte(visitorChannelAttachInfo.getLicenseId(), bytes.length);
-            });
-
-            ProxyUtil.addVisitorChannelToCmdChannel(cmdChannel, visitorId, visitorChannel, serverPort);
+            visitorId = ProxyUtil.newVisitorId();
+            ProxyUtil.addVisitorChannelToCmdChannel(cmdChannel, visitorId, ctx.channel(), serverPort);
             ProxyUtil.addProxyConnectAttachment(visitorId, proxyAttachment);
             cmdChannel.writeAndFlush(ProxyMessage.buildConnectMessage(visitorId).setData(lanInfo.getBytes()));
         }
@@ -167,10 +174,7 @@ public class HttpProxy implements EventListener<AppLoadEndEvent> {
             ctx.close();
         }
 
-        private String getHost(ByteBuf byteBuf) {
-            byte[] buf = new byte[byteBuf.readableBytes()];
-            byteBuf.readBytes(buf);
-            byteBuf.resetReaderIndex();
+        private String getHost(byte[] buf) {
             String req = new String(buf);
             String[] lines = req.split("\r\n");
             String firstLine = lines[0];