From 4d8e7e322b1b77828e53303df0fe91e5c4dc45d2 Mon Sep 17 00:00:00 2001
From: tsosunchia <59512455+tsosunchia@users.noreply.github.com>
Date: Mon, 14 Apr 2025 22:04:22 +0800
Subject: [PATCH] =?UTF-8?q?=E5=9B=BA=E5=AE=9AUDP=E6=A8=A1=E5=BC=8F?=
 =?UTF-8?q?=E7=9A=84=E6=BA=90=E7=AB=AF=E5=8F=A3=EF=BC=8C=E6=94=B9=E5=96=84?=
 =?UTF-8?q?=E5=9C=A8=E4=B8=80=E4=BA=9B=E8=B4=9F=E8=BD=BD=E5=9D=87=E8=A1=A1?=
 =?UTF-8?q?=E7=BD=91=E7=BB=9C=E7=8E=AF=E5=A2=83=E4=B8=8B=E7=9A=84=E6=95=88?=
 =?UTF-8?q?=E6=9E=9C=EF=BC=8C=E5=9C=A8=E4=BD=BF=E7=94=A8ENV=20"NEXTTRACE?=
 =?UTF-8?q?=5FRANDOMPORT"=20=E6=94=AF=E6=8C=81=E7=BB=A7=E7=BB=AD=E9=9A=8F?=
 =?UTF-8?q?=E6=9C=BA=E5=88=86=E9=85=8D=E6=BA=90=E7=AB=AF=E5=8F=A3=E3=80=82?=
 =?UTF-8?q?=20https://github.com/nxtrace/NTrace-core/issues/296=20https://?=
 =?UTF-8?q?github.com/nxtrace/NTrace-core/pull/297?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Yunlq <vinculo025@gmail.com>
Co-authored-by: tsosunchia <59512455+tsosunchia@users.noreply.github.com>
---
 trace/udp.go | 47 ++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 38 insertions(+), 9 deletions(-)

diff --git a/trace/udp.go b/trace/udp.go
index c9bc423..21edcb8 100644
--- a/trace/udp.go
+++ b/trace/udp.go
@@ -4,6 +4,7 @@ import (
 	"log"
 	"math/rand"
 	"net"
+	"strconv"
 	"sync"
 	"time"
 
@@ -135,9 +136,11 @@ func (t *UDPTracer) handleICMPMessage(msg ReceivedMessage, data []byte) {
 	}
 }
 
-func (t *UDPTracer) getUDPConn(try int) (net.IP, int, net.PacketConn) {
-	srcIP, _ := util.LocalIPPort(t.DestIP)
+var cachedLocalPort int
+var localPortOnce sync.Once
 
+func (t *UDPTracer) getUDPConn(try int) (net.IP, int, net.PacketConn, error) {
+	srcIP, _ := util.LocalIPPort(t.DestIP)
 	var ipString string
 	if srcIP == nil {
 		ipString = ""
@@ -145,14 +148,37 @@ func (t *UDPTracer) getUDPConn(try int) (net.IP, int, net.PacketConn) {
 		ipString = srcIP.String()
 	}
 
-	udpConn, err := net.ListenPacket("udp", ipString+":0")
-	if err != nil {
-		if try > 3 {
-			log.Fatal(err)
+	// Check environment variable to decide caching behavior
+	if util.GetenvDefault("NEXTTRACE_RANDOMPORT", "") == "" {
+		// Use cached random port logic
+		if cachedLocalPort == 0 {
+			// First time: listen on a random port
+			udpConn, err := net.ListenPacket("udp", ipString+":0")
+			if err != nil {
+				if try > 3 {
+					log.Fatal(err)
+				}
+				return srcIP, 0, nil, err
+			}
+			cachedLocalPort = udpConn.LocalAddr().(*net.UDPAddr).Port
+			// Close the initial connection after obtaining the port
+			udpConn.Close()
 		}
-		return t.getUDPConn(try + 1)
+		// Use the cached local port to establish a new connection
+		udpConn, err := net.ListenPacket("udp", ipString+":"+strconv.Itoa(cachedLocalPort))
+		if err != nil {
+			return srcIP, cachedLocalPort, nil, err
+		}
+		return srcIP, cachedLocalPort, udpConn, nil
+	} else {
+		// Without caching: create a new connection each time using a new random port
+		udpConn, err := net.ListenPacket("udp", ipString+":0")
+		if err != nil {
+			return srcIP, 0, nil, err
+		}
+		localPort := udpConn.LocalAddr().(*net.UDPAddr).Port
+		return srcIP, localPort, udpConn, nil
 	}
-	return srcIP, udpConn.LocalAddr().(*net.UDPAddr).Port, udpConn
 }
 
 func (t *UDPTracer) send(ttl int) error {
@@ -167,7 +193,10 @@ func (t *UDPTracer) send(ttl int) error {
 		return nil
 	}
 
-	srcIP, srcPort, udpConn := t.getUDPConn(0)
+	srcIP, srcPort, udpConn, err := t.getUDPConn(0)
+	if err != nil {
+		return err
+	}
 	defer udpConn.Close()
 
 	//var payload []byte