From 8722fd26586033042512467c6b3fd0f120535a47 Mon Sep 17 00:00:00 2001
From: will <will@fp-will.xyz>
Date: Fri, 14 Aug 2026 11:08:19 +0900
Subject: [PATCH] fix: track UDP connections in connMap for force-close support

connMap type changed from map[string]net.Conn to map[string]io.Closer.
RoutedPacketConnection now registers in connMap + uuidConns, so
CloseByUUID and CloseByID tear down both TCP and UDP.
---
 internal/kernel/singbox/conntracker.go | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/internal/kernel/singbox/conntracker.go b/internal/kernel/singbox/conntracker.go
index 76fd79b..a1f6d28 100644
--- a/internal/kernel/singbox/conntracker.go
+++ b/internal/kernel/singbox/conntracker.go
@@ -125,7 +125,7 @@ type ConnTracker struct {
 	usersMu sync.RWMutex
 	users   map[int]*userStats  // userID → stats
 	uuidMap map[string]int      // UUID → userID (for lookup in RoutedConnection)
-	connMap map[string]net.Conn // connID → conn (only for force-close support)
+	connMap map[string]io.Closer // connID → conn (TCP + UDP, for force-close)
 	uuidConns map[string]map[string]struct{} // UUID → set of connIDs (reverse index for CloseByUUID)
 
 	idCounter atomic.Int64
@@ -147,7 +147,7 @@ func NewConnTracker(_ int) *ConnTracker {
 	return &ConnTracker{
 		users:         make(map[int]*userStats),
 		uuidMap:       make(map[string]int),
-		connMap:       make(map[string]net.Conn),
+		connMap:       make(map[string]io.Closer),
 		uuidConns:     make(map[string]map[string]struct{}),
 		globalDevices: make(map[int]map[string]bool),
 	}
@@ -265,10 +265,7 @@ func (t *ConnTracker) RoutedConnection(
 	}
 }
 
-// RoutedPacketConnection wraps UDP with per-user counting (UDP not in connMap).
-// Note: UDP connections are NOT stored in connMap because connMap is typed as
-// map[string]net.Conn, but PacketConn is a different interface. Force-close
-// for UDP connections is handled directly via trackedPacketConn.Close().
+// RoutedPacketConnection wraps UDP with per-user counting and force-close support.
 func (t *ConnTracker) RoutedPacketConnection(
 	ctx context.Context, conn N.PacketConn,
 	metadata adapter.InboundContext,
@@ -300,6 +297,15 @@ func (t *ConnTracker) RoutedPacketConnection(
 
 	connID := t.nextID()
 
+	// Store for force-close support (same as TCP)
+	t.usersMu.Lock()
+	t.connMap[connID] = conn
+	if _, ok := t.uuidConns[uuid]; !ok {
+		t.uuidConns[uuid] = make(map[string]struct{})
+	}
+	t.uuidConns[uuid][connID] = struct{}{}
+	t.usersMu.Unlock()
+
 	var lim *rate.Limiter
 	if slf := t.speedLimitFunc.Load(); slf != nil {
 		lim = (*slf)(uuid)
@@ -470,7 +476,7 @@ func (t *ConnTracker) CloseByUUID(uuid string) int {
 		t.usersMu.RUnlock()
 		return 0
 	}
-	conns := make([]net.Conn, 0, len(ids))
+	conns := make([]io.Closer, 0, len(ids))
 	for id := range ids {
 		if c, ok := t.connMap[id]; ok {
 			conns = append(conns, c)