From 9797cb31c06459cbf87771e138be7569aea367a9 Mon Sep 17 00:00:00 2001
From: xjasonlyu <xjasonlyu@gmail.com>
Date: Sun, 3 Apr 2022 23:16:36 +0800
Subject: [PATCH] Refactor: return metadata.Addr only

---
 metadata/metadata.go | 3 ---
 proxy/direct.go      | 4 ++--
 tunnel/udp.go        | 8 +++++++-
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/metadata/metadata.go b/metadata/metadata.go
index a70d73c..63755db 100644
--- a/metadata/metadata.go
+++ b/metadata/metadata.go
@@ -25,9 +25,6 @@ func (m *Metadata) SourceAddress() string {
 }
 
 func (m *Metadata) Addr() net.Addr {
-	if udpAddr := m.UDPAddr(); udpAddr != nil {
-		return udpAddr
-	}
 	return &Addr{metadata: m}
 }
 
diff --git a/proxy/direct.go b/proxy/direct.go
index a24d16d..40ef9b3 100644
--- a/proxy/direct.go
+++ b/proxy/direct.go
@@ -45,8 +45,8 @@ type directPacketConn struct {
 }
 
 func (pc *directPacketConn) WriteTo(b []byte, addr net.Addr) (int, error) {
-	if ma, ok := addr.(*M.Addr); ok && ma.Metadata().DstIP != nil {
-		return pc.PacketConn.WriteTo(b, ma.Metadata().UDPAddr())
+	if udpAddr, ok := addr.(*net.UDPAddr); ok {
+		return pc.PacketConn.WriteTo(b, udpAddr)
 	}
 
 	udpAddr, err := net.ResolveUDPAddr("udp", addr.String())
diff --git a/tunnel/udp.go b/tunnel/udp.go
index 0da6b7a..7ea28cd 100644
--- a/tunnel/udp.go
+++ b/tunnel/udp.go
@@ -48,7 +48,13 @@ func handleUDPConn(uc adapter.UDPConn) {
 	pc = newUDPTracker(pc, metadata)
 	defer pc.Close()
 
-	remote := metadata.Addr()
+	var remote net.Addr
+	if udpAddr := metadata.UDPAddr(); udpAddr != nil {
+		remote = udpAddr
+	} else {
+		remote = metadata.Addr()
+	}
+
 	go handleUDPToRemote(uc, pc, remote)
 	handleUDPToLocal(uc, pc, remote)
 }