mirror of
https://github.com/fatedier/frp.git
synced 2026-07-24 10:59:18 +08:00
fix(server): prevent control replacement lifecycle leaks (#5424)
This commit is contained in:
@@ -28,6 +28,7 @@ type ClientInfo struct {
|
||||
User string
|
||||
RawClientID string
|
||||
RunID string
|
||||
ControlID uint64
|
||||
Hostname string
|
||||
IP string
|
||||
Version string
|
||||
@@ -64,6 +65,16 @@ func newClientRegistryWithClock(clk clock.PassiveClock) *ClientRegistry {
|
||||
|
||||
// Register stores/updates metadata for a client and returns the registry key plus whether it conflicts with an online client.
|
||||
func (cr *ClientRegistry) Register(user, rawClientID, runID, hostname, version, remoteAddr, wireProtocol string) (key string, conflict bool) {
|
||||
return cr.RegisterWithControlID(user, rawClientID, runID, hostname, version, remoteAddr, wireProtocol, 0)
|
||||
}
|
||||
|
||||
// RegisterWithControlID is the generation-aware form used by ControlManager.
|
||||
// A control ID is process-local and prevents an older control generation from
|
||||
// changing the registry entry now owned by a newer generation with the same run ID.
|
||||
func (cr *ClientRegistry) RegisterWithControlID(
|
||||
user, rawClientID, runID, hostname, version, remoteAddr, wireProtocol string,
|
||||
controlID uint64,
|
||||
) (key string, conflict bool) {
|
||||
if runID == "" {
|
||||
return "", false
|
||||
}
|
||||
@@ -83,6 +94,16 @@ func (cr *ClientRegistry) Register(user, rawClientID, runID, hostname, version,
|
||||
if enforceUnique && exists && info.Online && info.RunID != "" && info.RunID != runID {
|
||||
return key, true
|
||||
}
|
||||
if previousKey, ok := cr.runIndex[runID]; ok && previousKey != key {
|
||||
if previous, ok := cr.clients[previousKey]; ok && previous.RunID == runID {
|
||||
if previous.RawClientID == "" {
|
||||
delete(cr.clients, previousKey)
|
||||
} else {
|
||||
setClientOffline(previous, now)
|
||||
}
|
||||
}
|
||||
delete(cr.runIndex, runID)
|
||||
}
|
||||
|
||||
if !exists {
|
||||
info = &ClientInfo{
|
||||
@@ -97,6 +118,7 @@ func (cr *ClientRegistry) Register(user, rawClientID, runID, hostname, version,
|
||||
|
||||
info.RawClientID = rawClientID
|
||||
info.RunID = runID
|
||||
info.ControlID = controlID
|
||||
info.Hostname = hostname
|
||||
info.IP = remoteAddr
|
||||
info.Version = version
|
||||
@@ -114,6 +136,16 @@ func (cr *ClientRegistry) Register(user, rawClientID, runID, hostname, version,
|
||||
|
||||
// MarkOfflineByRunID marks the client as offline when the corresponding control disconnects.
|
||||
func (cr *ClientRegistry) MarkOfflineByRunID(runID string) {
|
||||
cr.markOfflineByRunID(runID, 0, false)
|
||||
}
|
||||
|
||||
// MarkOfflineByRunIDAndControlID marks a client offline only when the registry
|
||||
// entry still belongs to the supplied control generation.
|
||||
func (cr *ClientRegistry) MarkOfflineByRunIDAndControlID(runID string, controlID uint64) {
|
||||
cr.markOfflineByRunID(runID, controlID, true)
|
||||
}
|
||||
|
||||
func (cr *ClientRegistry) markOfflineByRunID(runID string, controlID uint64, matchControlID bool) {
|
||||
cr.mu.Lock()
|
||||
defer cr.mu.Unlock()
|
||||
|
||||
@@ -121,17 +153,23 @@ func (cr *ClientRegistry) MarkOfflineByRunID(runID string) {
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if info, ok := cr.clients[key]; ok && info.RunID == runID {
|
||||
if info, ok := cr.clients[key]; ok && info.RunID == runID && (!matchControlID || info.ControlID == controlID) {
|
||||
if info.RawClientID == "" {
|
||||
delete(cr.clients, key)
|
||||
} else {
|
||||
info.RunID = ""
|
||||
info.Online = false
|
||||
now := cr.clock.Now()
|
||||
info.DisconnectedAt = now
|
||||
setClientOffline(info, cr.clock.Now())
|
||||
}
|
||||
}
|
||||
delete(cr.runIndex, runID)
|
||||
if info, ok := cr.clients[key]; !ok || info.RunID != runID {
|
||||
delete(cr.runIndex, runID)
|
||||
}
|
||||
}
|
||||
|
||||
func setClientOffline(info *ClientInfo, now time.Time) {
|
||||
info.RunID = ""
|
||||
info.ControlID = 0
|
||||
info.Online = false
|
||||
info.DisconnectedAt = now
|
||||
}
|
||||
|
||||
// List returns a snapshot of all known clients.
|
||||
|
||||
Reference in New Issue
Block a user