fix: track UDP connections in connMap for force-close support
CI / Test (push) Has been cancelled
CI / Build (amd64, linux) (push) Has been cancelled
CI / Build (arm64, linux) (push) Has been cancelled
CI / Build & Push Docker Image (push) Has been cancelled
CI / Release (push) Has been cancelled

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.
This commit is contained in:
will
2026-08-14 11:08:19 +09:00
parent 4bab53f20f
commit 8722fd2658
+13 -7
View File
@@ -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)