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
|
usersMu sync.RWMutex
|
||||||
users map[int]*userStats // userID → stats
|
users map[int]*userStats // userID → stats
|
||||||
uuidMap map[string]int // UUID → userID (for lookup in RoutedConnection)
|
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)
|
uuidConns map[string]map[string]struct{} // UUID → set of connIDs (reverse index for CloseByUUID)
|
||||||
|
|
||||||
idCounter atomic.Int64
|
idCounter atomic.Int64
|
||||||
@@ -147,7 +147,7 @@ func NewConnTracker(_ int) *ConnTracker {
|
|||||||
return &ConnTracker{
|
return &ConnTracker{
|
||||||
users: make(map[int]*userStats),
|
users: make(map[int]*userStats),
|
||||||
uuidMap: make(map[string]int),
|
uuidMap: make(map[string]int),
|
||||||
connMap: make(map[string]net.Conn),
|
connMap: make(map[string]io.Closer),
|
||||||
uuidConns: make(map[string]map[string]struct{}),
|
uuidConns: make(map[string]map[string]struct{}),
|
||||||
globalDevices: make(map[int]map[string]bool),
|
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).
|
// RoutedPacketConnection wraps UDP with per-user counting and force-close support.
|
||||||
// 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().
|
|
||||||
func (t *ConnTracker) RoutedPacketConnection(
|
func (t *ConnTracker) RoutedPacketConnection(
|
||||||
ctx context.Context, conn N.PacketConn,
|
ctx context.Context, conn N.PacketConn,
|
||||||
metadata adapter.InboundContext,
|
metadata adapter.InboundContext,
|
||||||
@@ -300,6 +297,15 @@ func (t *ConnTracker) RoutedPacketConnection(
|
|||||||
|
|
||||||
connID := t.nextID()
|
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
|
var lim *rate.Limiter
|
||||||
if slf := t.speedLimitFunc.Load(); slf != nil {
|
if slf := t.speedLimitFunc.Load(); slf != nil {
|
||||||
lim = (*slf)(uuid)
|
lim = (*slf)(uuid)
|
||||||
@@ -470,7 +476,7 @@ func (t *ConnTracker) CloseByUUID(uuid string) int {
|
|||||||
t.usersMu.RUnlock()
|
t.usersMu.RUnlock()
|
||||||
return 0
|
return 0
|
||||||
}
|
}
|
||||||
conns := make([]net.Conn, 0, len(ids))
|
conns := make([]io.Closer, 0, len(ids))
|
||||||
for id := range ids {
|
for id := range ids {
|
||||||
if c, ok := t.connMap[id]; ok {
|
if c, ok := t.connMap[id]; ok {
|
||||||
conns = append(conns, c)
|
conns = append(conns, c)
|
||||||
|
|||||||
Reference in New Issue
Block a user