chore: ensures packets can be sent without blocking the tunnel

component/nat/table.go

@@ -10,47 +10,30 @@ import (
 )
 
 type Table struct {
-	mapping *xsync.MapOf[string, *Entry]
-	lockMap *xsync.MapOf[string, *sync.Cond]
+	mapping *xsync.MapOf[string, *entry]
 }
 
-type Entry struct {
-	PacketConn      C.PacketConn
-	WriteBackProxy  C.WriteBackProxy
+type entry struct {
+	PacketSender    C.PacketSender
 	LocalUDPConnMap *xsync.MapOf[string, *net.UDPConn]
 	LocalLockMap    *xsync.MapOf[string, *sync.Cond]
 }
 
-func (t *Table) Set(key string, e C.PacketConn, w C.WriteBackProxy) {
-	t.mapping.Store(key, &Entry{
-		PacketConn:      e,
-		WriteBackProxy:  w,
-		LocalUDPConnMap: xsync.NewMapOf[string, *net.UDPConn](),
-		LocalLockMap:    xsync.NewMapOf[string, *sync.Cond](),
+func (t *Table) GetOrCreate(key string, maker func() C.PacketSender) (C.PacketSender, bool) {
+	item, loaded := t.mapping.LoadOrCompute(key, func() *entry {
+		return &entry{
+			PacketSender:    maker(),
+			LocalUDPConnMap: xsync.NewMapOf[string, *net.UDPConn](),
+			LocalLockMap:    xsync.NewMapOf[string, *sync.Cond](),
+		}
 	})
-}
-
-func (t *Table) Get(key string) (C.PacketConn, C.WriteBackProxy) {
-	entry, exist := t.getEntry(key)
-	if !exist {
-		return nil, nil
-	}
-	return entry.PacketConn, entry.WriteBackProxy
-}
-
-func (t *Table) GetOrCreateLock(key string) (*sync.Cond, bool) {
-	item, loaded := t.lockMap.LoadOrCompute(key, makeLock)
-	return item, loaded
+	return item.PacketSender, loaded
 }
 
 func (t *Table) Delete(key string) {
 	t.mapping.Delete(key)
 }
 
-func (t *Table) DeleteLock(lockKey string) {
-	t.lockMap.Delete(lockKey)
-}
-
 func (t *Table) GetForLocalConn(lAddr, rAddr string) *net.UDPConn {
 	entry, exist := t.getEntry(lAddr)
 	if !exist {
@@ -105,7 +88,7 @@ func (t *Table) DeleteLockForLocalConn(lAddr, key string) {
 	entry.LocalLockMap.Delete(key)
 }
 
-func (t *Table) getEntry(key string) (*Entry, bool) {
+func (t *Table) getEntry(key string) (*entry, bool) {
 	return t.mapping.Load(key)
 }
 
@@ -116,7 +99,6 @@ func makeLock() *sync.Cond {
 // New return *Cache
 func New() *Table {
 	return &Table{
-		mapping: xsync.NewMapOf[string, *Entry](),
-		lockMap: xsync.NewMapOf[string, *sync.Cond](),
+		mapping: xsync.NewMapOf[string, *entry](),
 	}
 }

constant/adapters.go

@@ -255,23 +255,33 @@ type UDPPacketInAddr interface {
 // PacketAdapter is a UDP Packet adapter for socks/redir/tun
 type PacketAdapter interface {
 	UDPPacket
+	// Metadata returns destination metadata
 	Metadata() *Metadata
+	// Key is a SNAT key
+	Key() string
 }
 
 type packetAdapter struct {
 	UDPPacket
 	metadata *Metadata
+	key      string
 }
 
 // Metadata returns destination metadata
 func (s *packetAdapter) Metadata() *Metadata {
 	return s.metadata
 }
 
+// Key is a SNAT key
+func (s *packetAdapter) Key() string {
+	return s.key
+}
+
 func NewPacketAdapter(packet UDPPacket, metadata *Metadata) PacketAdapter {
 	return &packetAdapter{
 		packet,
 		metadata,
+		packet.LocalAddr().String(),
 	}
 }
 
@@ -284,17 +294,19 @@ type WriteBackProxy interface {
 	UpdateWriteBack(wb WriteBack)
 }
 
-type NatTable interface {
-	Set(key string, e PacketConn, w WriteBackProxy)
-
-	Get(key string) (PacketConn, WriteBackProxy)
+type PacketSender interface {
+	// Send will send PacketAdapter nonblocking
+	// the implement must call UDPPacket.Drop() inside Send
+	Send(PacketAdapter)
+	Process(PacketConn, WriteBackProxy)
+	Close()
+}
 
-	GetOrCreateLock(key string) (*sync.Cond, bool)
+type NatTable interface {
+	GetOrCreate(key string, maker func() PacketSender) (PacketSender, bool)
 
 	Delete(key string)
 
-	DeleteLock(key string)
-
 	GetForLocalConn(lAddr, rAddr string) *net.UDPConn
 
 	AddForLocalConn(lAddr, rAddr string, conn *net.UDPConn) bool

tunnel/connection.go

@@ -1,6 +1,7 @@
 package tunnel
 
 import (
+	"context"
 	"errors"
 	"net"
 	"net/netip"
@@ -11,7 +12,78 @@ import (
 	"github.com/metacubex/mihomo/log"
 )
 
+type packetSender struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+	ch     chan C.PacketAdapter
+}
+
+// newPacketSender return a chan based C.PacketSender
+// It ensures that packets can be sent sequentially and without blocking
+func newPacketSender() C.PacketSender {
+	ctx, cancel := context.WithCancel(context.Background())
+	ch := make(chan C.PacketAdapter, senderCapacity)
+	return &packetSender{
+		ctx:    ctx,
+		cancel: cancel,
+		ch:     ch,
+	}
+}
+
+func (s *packetSender) Process(pc C.PacketConn, proxy C.WriteBackProxy) {
+	for {
+		select {
+		case <-s.ctx.Done():
+			return // sender closed
+		case packet := <-s.ch:
+			if proxy != nil {
+				proxy.UpdateWriteBack(packet)
+			}
+			_ = handleUDPToRemote(packet, pc, packet.Metadata())
+			packet.Drop()
+		}
+	}
+}
+
+func (s *packetSender) dropAll() {
+	for {
+		select {
+		case data := <-s.ch:
+			data.Drop() // drop all data still in chan
+		default:
+			return // no data, exit goroutine
+		}
+	}
+}
+
+func (s *packetSender) Send(packet C.PacketAdapter) {
+	select {
+	case <-s.ctx.Done():
+		packet.Drop() // sender closed before Send()
+		return
+	default:
+	}
+
+	select {
+	case s.ch <- packet:
+		// put ok, so don't drop packet, will process by other side of chan
+	case <-s.ctx.Done():
+		packet.Drop() // sender closed when putting data to chan
+	default:
+		packet.Drop() // chan is full
+	}
+}
+
+func (s *packetSender) Close() {
+	s.cancel()
+	s.dropAll()
+}
+
 func handleUDPToRemote(packet C.UDPPacket, pc C.PacketConn, metadata *C.Metadata) error {
+	if err := resolveUDP(metadata); err != nil {
+		return err
+	}
+
 	addr := metadata.UDPAddr()
 	if addr == nil {
 		return errors.New("udp addr invalid")
@@ -26,8 +98,9 @@ func handleUDPToRemote(packet C.UDPPacket, pc C.PacketConn, metadata *C.Metadata
 	return nil
 }
 
-func handleUDPToLocal(writeBack C.WriteBack, pc N.EnhancePacketConn, key string, oAddrPort netip.AddrPort, fAddr netip.Addr) {
+func handleUDPToLocal(writeBack C.WriteBack, pc N.EnhancePacketConn, sender C.PacketSender, key string, oAddrPort netip.AddrPort, fAddr netip.Addr) {
 	defer func() {
+		sender.Close()
 		_ = pc.Close()
 		closeAllLocalCoon(key)
 		natTable.Delete(key)