feat: introduce a new robust approach to handle tproxy udp. (#389)

component/nat/table.go

@@ -1,6 +1,7 @@
 package nat
 
 import (
+	"net"
 	"sync"
 
 	C "github.com/Dreamacro/clash/constant"
@@ -10,16 +11,24 @@ type Table struct {
 	mapping sync.Map
 }
 
-func (t *Table) Set(key string, pc C.PacketConn) {
-	t.mapping.Store(key, pc)
+type Entry struct {
+	PacketConn      C.PacketConn
+	LocalUDPConnMap sync.Map
+}
+
+func (t *Table) Set(key string, e C.PacketConn) {
+	t.mapping.Store(key, &Entry{
+		PacketConn:      e,
+		LocalUDPConnMap: sync.Map{},
+	})
 }
 
 func (t *Table) Get(key string) C.PacketConn {
-	item, exist := t.mapping.Load(key)
+	entry, exist := t.getEntry(key)
 	if !exist {
 		return nil
 	}
-	return item.(C.PacketConn)
+	return entry.PacketConn
 }
 
 func (t *Table) GetOrCreateLock(key string) (*sync.Cond, bool) {
@@ -31,6 +40,62 @@ func (t *Table) Delete(key string) {
 	t.mapping.Delete(key)
 }
 
+func (t *Table) GetLocalConn(lAddr, rAddr string) *net.UDPConn {
+	entry, exist := t.getEntry(lAddr)
+	if !exist {
+		return nil
+	}
+	item, exist := entry.LocalUDPConnMap.Load(rAddr)
+	if !exist {
+		return nil
+	}
+	return item.(*net.UDPConn)
+}
+
+func (t *Table) AddLocalConn(lAddr, rAddr string, conn *net.UDPConn) bool {
+	entry, exist := t.getEntry(lAddr)
+	if !exist {
+		return false
+	}
+	entry.LocalUDPConnMap.Store(rAddr, conn)
+	return true
+}
+
+func (t *Table) RangeLocalConn(lAddr string, f func(key, value any) bool) {
+	entry, exist := t.getEntry(lAddr)
+	if !exist {
+		return
+	}
+	entry.LocalUDPConnMap.Range(f)
+}
+
+func (t *Table) GetOrCreateLockForLocalConn(lAddr, key string) (*sync.Cond, bool) {
+	entry, loaded := t.getEntry(lAddr)
+	if !loaded {
+		return nil, false
+	}
+	item, loaded := entry.LocalUDPConnMap.LoadOrStore(key, sync.NewCond(&sync.Mutex{}))
+	return item.(*sync.Cond), loaded
+}
+
+func (t *Table) DeleteLocalConnMap(lAddr, key string) {
+	entry, loaded := t.getEntry(lAddr)
+	if !loaded {
+		return
+	}
+	entry.LocalUDPConnMap.Delete(key)
+}
+
+func (t *Table) getEntry(key string) (*Entry, bool) {
+	item, ok := t.mapping.Load(key)
+	// This should not happen usually since this function called after PacketConn created
+	if !ok {
+		return nil, false
+	}
+	entry, ok := item.(*Entry)
+	return entry, ok
+}
+
 // New return *Cache
 func New() *Table {
 	return &Table{}