From 3cf64596c659a1faf13171906a283dcad43e4a1b Mon Sep 17 00:00:00 2001
From: aoshiguchen <1052045476@qq.com>
Date: Sat, 23 Sep 2023 10:56:45 +0800
Subject: [PATCH] =?UTF-8?q?udp=E4=BB=A3=E7=90=86=E4=BC=98=E5=8C=96.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../java/org/dromara/neutrinoproxy/core/Constants.java |  3 +++
 .../server/proxy/core/UdpVisitorChannelHandler.java    |  2 ++
 .../proxy/handler/UdpProxyMessageTransferHandler.java  | 10 +++++++---
 3 files changed, 12 insertions(+), 3 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 59c5175e..472f8ce9 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
@@ -25,6 +25,8 @@ package org.dromara.neutrinoproxy.core;
 import io.netty.channel.Channel;
 import io.netty.util.AttributeKey;
 
+import java.net.InetSocketAddress;
+
 /**
  *
  * @author: aoshiguchen
@@ -45,6 +47,7 @@ public interface Constants {
 
     // 临时解决 udp channel被close的问题
     AttributeKey<Boolean> IS_UDP_KEY =  AttributeKey.newInstance("isUdp");
+    AttributeKey<InetSocketAddress> SENDER = AttributeKey.newInstance("sender");
 
 
     int HEADER_SIZE = 4;
diff --git a/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/core/UdpVisitorChannelHandler.java b/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/core/UdpVisitorChannelHandler.java
index e3f07111..07abc014 100644
--- a/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/core/UdpVisitorChannelHandler.java
+++ b/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/core/UdpVisitorChannelHandler.java
@@ -41,6 +41,8 @@ public class UdpVisitorChannelHandler extends SimpleChannelInboundHandler<Datagr
 //                ctx.channel().close();
                 return;
             }
+
+            proxyChannel.attr(Constants.SENDER).set(datagramPacket.sender());
             String targetIp = proxyChannel.attr(Constants.TARGET_IP).get();
             int targetPort = proxyChannel.attr(Constants.TARGET_PORT).get();
             Integer proxyResponses = proxyChannel.attr(Constants.PROXY_RESPONSES).get();
diff --git a/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/handler/UdpProxyMessageTransferHandler.java b/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/handler/UdpProxyMessageTransferHandler.java
index 7052b5d9..978bae9d 100644
--- a/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/handler/UdpProxyMessageTransferHandler.java
+++ b/neutrino-proxy-server/src/main/java/org/dromara/neutrinoproxy/server/proxy/handler/UdpProxyMessageTransferHandler.java
@@ -35,9 +35,13 @@ public class UdpProxyMessageTransferHandler implements ProxyMessageHandler {
 
         Channel visitorChannel = ctx.channel().attr(Constants.NEXT_CHANNEL).get();
         if (null != visitorChannel) {
-            InetSocketAddress address = new InetSocketAddress(udpBaseInfo.getVisitorIp(), udpBaseInfo.getVisitorPort());
-            ByteBuf byteBuf = Unpooled.copiedBuffer(proxyMessage.getData());
-            visitorChannel.writeAndFlush(new DatagramPacket(byteBuf, address));
+//            InetSocketAddress address = new InetSocketAddress(udpBaseInfo.getVisitorIp(), udpBaseInfo.getVisitorPort());
+//            ByteBuf byteBuf = Unpooled.copiedBuffer(proxyMessage.getData());
+//            visitorChannel.writeAndFlush(new DatagramPacket(byteBuf, address));
+            InetSocketAddress address = ctx.channel().attr(Constants.SENDER).get();
+            if (null != address) {
+                visitorChannel.writeAndFlush(new DatagramPacket(Unpooled.copiedBuffer(proxyMessage.getData()), address));
+            }
 
             // 增加流量计数(TODO 如果UDP映射服务端端口修改，这个似乎不准)
             Integer licenseId = visitorChannel.attr(Constants.LICENSE_ID).get();