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.
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user