refactor: clean up code (#5308)

server/registry/registry.go

@@ -18,6 +18,8 @@ import (
 	"fmt"
 	"sync"
 	"time"
+
+	"k8s.io/utils/clock"
 )
 
 // ClientInfo captures metadata about a connected frpc instance.
@@ -42,12 +44,21 @@ type ClientRegistry struct {
 	mu       sync.RWMutex
 	clients  map[string]*ClientInfo
 	runIndex map[string]string
+	clock    clock.PassiveClock
 }
 
 func NewClientRegistry() *ClientRegistry {
+	return newClientRegistryWithClock(clock.RealClock{})
+}
+
+func newClientRegistryWithClock(clk clock.PassiveClock) *ClientRegistry {
+	if clk == nil {
+		clk = clock.RealClock{}
+	}
 	return &ClientRegistry{
 		clients:  make(map[string]*ClientInfo),
 		runIndex: make(map[string]string),
+		clock:    clk,
 	}
 }
 
@@ -64,7 +75,7 @@ func (cr *ClientRegistry) Register(user, rawClientID, runID, hostname, version,
 	key = cr.composeClientKey(user, effectiveID)
 	enforceUnique := rawClientID != ""
 
-	now := time.Now()
+	now := cr.clock.Now()
 	cr.mu.Lock()
 	defer cr.mu.Unlock()
 
@@ -116,7 +127,7 @@ func (cr *ClientRegistry) MarkOfflineByRunID(runID string) {
 		} else {
 			info.RunID = ""
 			info.Online = false
-			now := time.Now()
+			now := cr.clock.Now()
 			info.DisconnectedAt = now
 		}
 	}

server/registry/registry_test.go

@@ -16,6 +16,9 @@ package registry
 
 import (
 	"testing"
+	"time"
+
+	clocktesting "k8s.io/utils/clock/testing"
 
 	"github.com/fatedier/frp/pkg/proto/wire"
 )
@@ -35,3 +38,37 @@ func TestClientRegistryRegisterStoresWireProtocol(t *testing.T) {
 		t.Fatalf("wire protocol mismatch, want %q got %q", wire.ProtocolV2, info.WireProtocol)
 	}
 }
+
+func TestClientRegistryUsesClockForTimestamps(t *testing.T) {
+	start := time.Date(2026, time.May, 8, 12, 30, 0, 0, time.UTC)
+	clk := clocktesting.NewFakeClock(start)
+	registry := newClientRegistryWithClock(clk)
+
+	key, conflict := registry.Register("user", "client-id", "run-id", "host", "1.0.0", "127.0.0.1", wire.ProtocolV2)
+	if conflict {
+		t.Fatal("unexpected client conflict")
+	}
+
+	info, ok := registry.GetByKey(key)
+	if !ok {
+		t.Fatalf("client %q not found", key)
+	}
+	if !info.FirstConnectedAt.Equal(start) {
+		t.Fatalf("first connected time mismatch, want %s got %s", start, info.FirstConnectedAt)
+	}
+	if !info.LastConnectedAt.Equal(start) {
+		t.Fatalf("last connected time mismatch, want %s got %s", start, info.LastConnectedAt)
+	}
+
+	disconnectedAt := start.Add(time.Minute)
+	clk.SetTime(disconnectedAt)
+	registry.MarkOfflineByRunID("run-id")
+
+	info, ok = registry.GetByKey(key)
+	if !ok {
+		t.Fatalf("client %q not found after disconnect", key)
+	}
+	if !info.DisconnectedAt.Equal(disconnectedAt) {
+		t.Fatalf("disconnected time mismatch, want %s got %s", disconnectedAt, info.DisconnectedAt)
+	}
+}