From fc0830acfed009c4ca09db874d2c764345b9946e Mon Sep 17 00:00:00 2001
From: aoshiguchen <1052045476@qq.com>
Date: Mon, 18 Dec 2023 18:39:23 +0800
Subject: [PATCH] =?UTF-8?q?=E8=A7=A3=E5=86=B3=E4=BD=BF=E7=94=A8=E6=98=A0?=
 =?UTF-8?q?=E5=B0=84=E5=9F=9F=E5=90=8D=E4=B8=8A=E4=BC=A0=E6=96=87=E4=BB=B6?=
 =?UTF-8?q?=E8=BF=9E=E6=8E=A5=E6=96=AD=E5=BC=80=E7=9A=84=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../dromara/neutrinoproxy/core/Constants.java |  2 +
 .../HttpVisitorSecurityChannelHandler.java    | 75 ++++++++++---------
 2 files changed, 42 insertions(+), 35 deletions(-)

diff --git a/neutrino-proxy-core/src/main/java/org/dromara/neutrinoproxy/core/Constants.java b/neutrino-proxy-core/src/main/java/org/dromara/neutrinoproxy/core/Constants.java
index 1f129bf7..a88c3e10 100644
--- a/neutrino-proxy-core/src/main/java/org/dromara/neutrinoproxy/core/Constants.java
+++ b/neutrino-proxy-core/src/main/java/org/dromara/neutrinoproxy/core/Constants.java
@@ -50,6 +50,8 @@ public interface Constants {
     AttributeKey<InetSocketAddress> SENDER = AttributeKey.newInstance("sender");
 
     AttributeKey<Integer> SERVER_PORT = AttributeKey.newInstance("serverPort");
+    AttributeKey<String> REAL_REMOTE_IP = AttributeKey.newInstance("realRemoteIp");
+
 
     AttributeKey<Boolean> FLOW_LIMITER_FLAG = AttributeKey.newInstance("flowLimiterFlag");
 
diff --git a/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/security/HttpVisitorSecurityChannelHandler.java b/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/security/HttpVisitorSecurityChannelHandler.java
index a8128b90..eccd2914 100644
--- a/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/security/HttpVisitorSecurityChannelHandler.java
+++ b/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/security/HttpVisitorSecurityChannelHandler.java
@@ -39,47 +39,52 @@ public class HttpVisitorSecurityChannelHandler extends ChannelInboundHandlerAdap
         }
 
         ByteBuf buf = (ByteBuf) msg;
-        byte[] bytes = new byte[buf.readableBytes()];
-        buf.readBytes(bytes);
 
-        // 获取Host请求头
-        String httpContent = new String(bytes);
-        String host = HttpUtil.getHostIgnorePort(httpContent);
-
-        log.debug("HttpProxy host: {}", host);
-        if (StringUtils.isBlank(host)) {
-            ctx.channel().close();
-            return;
-        }
-
-        // 根据Host匹配端口映射
-        if (!host.endsWith(domainName)) {
-            ctx.channel().close();
-            return;
-        }
-        int index = host.lastIndexOf("." + domainName);
-        String subdomain = host.substring(0, index);
-
-        // 根据域名拿到绑定的映射对应的cmdChannel
-        Integer serverPort = ProxyUtil.getServerPortBySubdomain(subdomain);
+        Integer serverPort = ctx.channel().attr(Constants.SERVER_PORT).get();
         if (null == serverPort) {
-            ctx.channel().close();
-            return;
-        }
+            // 获取Host请求头
+            byte[] bytes = new byte[buf.readableBytes()];
+            buf.readBytes(bytes);
+            String httpContent = new String(bytes);
+            String host = HttpUtil.getHostIgnorePort(httpContent);
 
-        // 判断IP是否在该端口绑定的安全组允许的规则内
-        String ip = IpUtil.getRealRemoteIp(httpContent);
-        if (ip == null) {
-            ip = IpUtil.getRemoteIp(ctx);
-        }
-        if (!securityGroupService.judgeAllow(ip, portMappingService.getSecurityGroupIdByMappingPort(serverPort))) {
-            // 不在安全组规则放行范围内
-            ctx.channel().close();
-            return;
+            log.debug("HttpProxy host: {}", host);
+            if (StringUtils.isBlank(host)) {
+                ctx.channel().close();
+                return;
+            }
+
+            // 根据Host匹配端口映射
+            if (!host.endsWith(domainName)) {
+                ctx.channel().close();
+                return;
+            }
+            int index = host.lastIndexOf("." + domainName);
+            String subdomain = host.substring(0, index);
+
+            // 根据域名拿到绑定的映射对应的cmdChannel
+            serverPort = ProxyUtil.getServerPortBySubdomain(subdomain);
+            if (null == serverPort) {
+                ctx.channel().close();
+                return;
+            }
+
+            // 判断IP是否在该端口绑定的安全组允许的规则内
+            String ip = IpUtil.getRealRemoteIp(httpContent);
+            if (ip == null) {
+                ip = IpUtil.getRemoteIp(ctx);
+            }
+            if (!securityGroupService.judgeAllow(ip, portMappingService.getSecurityGroupIdByMappingPort(serverPort))) {
+                // 不在安全组规则放行范围内
+                ctx.channel().close();
+                return;
+            }
+
+            ctx.channel().attr(Constants.REAL_REMOTE_IP).set(ip);
+            ctx.channel().attr(Constants.SERVER_PORT).set(serverPort);
         }
 
         // 继续传播
-        ctx.channel().attr(Constants.SERVER_PORT).set(serverPort);
         buf.resetReaderIndex();
         ctx.fireChannelRead(buf);
     }