Mxmilu666/frp
1
1
Fork 1
mirror of https://github.com/fatedier/frp.git synced 2026-07-24 07:29:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(server): prevent control replacement lifecycle leaks (#5424)

Browse Source
This commit is contained in:
fatedier
2026-07-21 18:09:09 +08:00
committed by GitHub
parent fe79598ee4
commit d486018885
10 changed files with 2241 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
server/registry/registry.go
View File

@@ -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.

86
server/registry/registry_test.go
View File

@@ -72,3 +72,89 @@ func TestClientRegistryUsesClockForTimestamps(t *testing.T) {
t.Fatalf("disconnected time mismatch, want %s got %s", disconnectedAt, info.DisconnectedAt)
}
}
func TestClientRegistryControlIDPreventsStaleOffline(t *testing.T) {
registry := NewClientRegistry()
key, conflict := registry.RegisterWithControlID(
"user", "client-id", "run-id", "old-host", "1.0.0", "127.0.0.1", wire.ProtocolV1, 1,
)
if conflict {
t.Fatal("unexpected client conflict")
}
_, conflict = registry.RegisterWithControlID(
"user", "client-id", "run-id", "new-host", "1.0.1", "127.0.0.2", wire.ProtocolV2, 2,
)
if conflict {
t.Fatal("same run ID replacement should not conflict")
}
registry.MarkOfflineByRunIDAndControlID("run-id", 1)
info, ok := registry.GetByKey(key)
if !ok {
t.Fatalf("client %q not found", key)
}
if !info.Online || info.ControlID != 2 || info.Hostname != "new-host" {
t.Fatalf("stale offline changed current generation: %+v", info)
}
registry.MarkOfflineByRunIDAndControlID("run-id", 2)
info, ok = registry.GetByKey(key)
if !ok {
t.Fatalf("client %q not found after disconnect", key)
}
if info.Online || info.ControlID != 0 || info.RunID != "" {
t.Fatalf("current generation was not marked offline: %+v", info)
}
}
func TestClientRegistryClientIDConflictSemantics(t *testing.T) {
registry := NewClientRegistry()
_, conflict := registry.RegisterWithControlID(
"user", "client-id", "run-one", "host", "1.0.0", "127.0.0.1", wire.ProtocolV1, 1,
)
if conflict {
t.Fatal("unexpected initial client conflict")
}
_, conflict = registry.RegisterWithControlID(
"user", "client-id", "run-two", "host", "1.0.0", "127.0.0.2", wire.ProtocolV1, 2,
)
if !conflict {
t.Fatal("different online run IDs with the same explicit client ID must conflict")
}
registry.MarkOfflineByRunIDAndControlID("run-one", 1)
_, conflict = registry.RegisterWithControlID(
"user", "client-id", "run-two", "host", "1.0.0", "127.0.0.2", wire.ProtocolV1, 2,
)
if conflict {
t.Fatal("offline explicit client ID should be reusable")
}
}
func TestClientRegistrySameRunIDMovesBetweenClientKeys(t *testing.T) {
registry := NewClientRegistry()
oldKey, conflict := registry.RegisterWithControlID(
"user", "old-client", "run-id", "old-host", "1.0.0", "127.0.0.1", wire.ProtocolV1, 1,
)
if conflict {
t.Fatal("unexpected initial client conflict")
}
newKey, conflict := registry.RegisterWithControlID(
"user", "new-client", "run-id", "new-host", "1.0.1", "127.0.0.2", wire.ProtocolV2, 2,
)
if conflict {
t.Fatal("same run ID moving to a new client key should not conflict")
}
oldInfo, ok := registry.GetByKey(oldKey)
if !ok {
t.Fatalf("old explicit client %q should remain as offline history", oldKey)
}
if oldInfo.Online || oldInfo.RunID != "" || oldInfo.ControlID != 0 {
t.Fatalf("old client key remained online: %+v", oldInfo)
}
newInfo, ok := registry.GetByKey(newKey)
if !ok || !newInfo.Online || newInfo.RunID != "run-id" || newInfo.ControlID != 2 {
t.Fatalf("new client key was not registered: %+v", newInfo)
}
}