diff --git a/Release.md b/Release.md index 2858be11..d4ad3208 100644 --- a/Release.md +++ b/Release.md @@ -1,9 +1,3 @@ -## Features - -* Expanded the frps dashboard API v2 migration across Clients, Proxies, Server Overview, Client Detail, and Proxy Detail, covering paginated users/clients/proxies, detail data, proxy traffic history, server system info, offline proxy statistics pruning, server-side pagination, search, and proxy type filtering. - ## Fixes -* WebSocket and WSS tunnel payloads are now sent as binary frames, avoiding disconnects through RFC-compliant intermediaries that validate text frames as UTF-8. -* The `tls2raw` client plugin now writes the proxy protocol header to the local raw connection when proxy protocol is enabled. -* frpc now rejects duplicate proxy and visitor names in config files instead of silently overwriting earlier entries. +* Fixed control-session replacement leaks when frpc reconnects through a half-open TCP multiplexed connection. diff --git a/server/control.go b/server/control.go index 084b4eea..6a5d3ce1 100644 --- a/server/control.go +++ b/server/control.go @@ -17,6 +17,7 @@ package server import ( "context" "fmt" + "net" "runtime/debug" "sync" "sync/atomic" @@ -40,55 +41,311 @@ import ( "github.com/fatedier/frp/server/registry" ) +type ControlID uint64 + +var nextControlID atomic.Uint64 + +type controlEntry struct { + ctl *Control + id ControlID + // runMu serializes lifecycle and routing decisions for one run ID. + // Replacements inherit it; removing the entry releases the manager's reference. + runMu *sync.Mutex + + registryOnline bool + registryControlID ControlID +} + type ControlManager struct { // controls indexed by run id - ctlsByRunID map[string]*Control + ctlsByRunID map[string]*controlEntry + registry *registry.ClientRegistry + closed bool mu sync.RWMutex } -func NewControlManager() *ControlManager { +func NewControlManager(clientRegistry *registry.ClientRegistry) *ControlManager { return &ControlManager{ - ctlsByRunID: make(map[string]*Control), + ctlsByRunID: make(map[string]*controlEntry), + registry: clientRegistry, } } -func (cm *ControlManager) Add(runID string, ctl *Control) (old *Control) { - cm.mu.Lock() - defer cm.mu.Unlock() - - var ok bool - old, ok = cm.ctlsByRunID[runID] - if ok { - old.Replaced(ctl) +// lockCurrentRun returns the current entry with its run gate held. It never +// waits for the gate while holding cm.mu and revalidates the gate after waiting. +// The global order is runMu, cm.mu, ctl.lifecycleMu, then registry locks. +func (cm *ControlManager) lockCurrentRun(runID string, allowClosed bool) (*controlEntry, bool) { + cm.mu.RLock() + entry, ok := cm.ctlsByRunID[runID] + if cm.closed && !allowClosed { + ok = false } - cm.ctlsByRunID[runID] = ctl - return + cm.mu.RUnlock() + if !ok { + return nil, false + } + + runMu := entry.runMu + runMu.Lock() + cm.mu.RLock() + entry, ok = cm.ctlsByRunID[runID] + if (cm.closed && !allowClosed) || !ok || entry.runMu != runMu { + ok = false + } + cm.mu.RUnlock() + if !ok { + runMu.Unlock() + return nil, false + } + return entry, true } -// we should make sure if it's the same control to prevent delete a new one -func (cm *ControlManager) Del(runID string, ctl *Control) { +// Add makes ctl the pending current generation and records the predecessor +// finalization barrier it must wait for before activation. +func (cm *ControlManager) Add(ctl *Control) error { + for { + // Never wait for a run gate while holding cm.mu. + cm.mu.RLock() + old := cm.ctlsByRunID[ctl.runID] + cm.mu.RUnlock() + if old != nil { + old.runMu.Lock() + } + + cm.mu.Lock() + if cm.closed { + cm.mu.Unlock() + if old != nil { + old.runMu.Unlock() + } + return fmt.Errorf("control manager is closed") + } + if cm.ctlsByRunID[ctl.runID] != old { + cm.mu.Unlock() + if old != nil { + old.runMu.Unlock() + } + continue + } + + id := ControlID(nextControlID.Add(1)) + if err := ctl.admit(cm, id); err != nil { + cm.mu.Unlock() + if old != nil { + old.runMu.Unlock() + } + return err + } + + runMu := &sync.Mutex{} + if old != nil { + runMu = old.runMu + } + entry := &controlEntry{ctl: ctl, id: id, runMu: runMu} + var ( + oldCtl *Control + barrier <-chan struct{} + ) + if old != nil { + oldCtl = old.ctl + barrier = oldCtl.markReplaced() + ctl.setHandoffBarrier(barrier) + entry.registryOnline = old.registryOnline + entry.registryControlID = old.registryControlID + } + cm.ctlsByRunID[ctl.runID] = entry + cm.mu.Unlock() + if old != nil { + old.runMu.Unlock() + } + + if oldCtl != nil { + oldCtl.Replaced(ctl) + } + return nil + } +} + +// Activate registers ctl as online only if it is still the pending current +// generation. +func (cm *ControlManager) Activate(ctl *Control) (bool, error) { + entry, ok := cm.lockCurrentRun(ctl.runID, false) + if !ok { + return false, nil + } + defer entry.runMu.Unlock() cm.mu.Lock() defer cm.mu.Unlock() - if c, ok := cm.ctlsByRunID[runID]; ok && c == ctl { - delete(cm.ctlsByRunID, runID) + + if cm.closed || cm.ctlsByRunID[ctl.runID] != entry || entry.ctl != ctl || entry.id != ctl.controlID { + return false, nil } + + ctl.lifecycleMu.Lock() + defer ctl.lifecycleMu.Unlock() + if ctl.state != controlStatePending { + return false, nil + } + if ctl.activated { + return true, nil + } + + loginMsg := ctl.sessionCtx.LoginMsg + remoteAddr := ctl.sessionCtx.Conn.RemoteAddr().String() + if host, _, err := net.SplitHostPort(remoteAddr); err == nil { + remoteAddr = host + } + _, conflict := cm.registry.RegisterWithControlID( + loginMsg.User, + loginMsg.ClientID, + ctl.runID, + loginMsg.Hostname, + loginMsg.Version, + remoteAddr, + ctl.sessionCtx.WireProtocol, + uint64(entry.id), + ) + if conflict { + return true, fmt.Errorf("client_id [%s] for user [%s] is already online", loginMsg.ClientID, loginMsg.User) + } + + entry.registryOnline = true + entry.registryControlID = entry.id + ctl.activated = true + return true, nil +} + +// completeLogin reserves ctl's current ownership with its run gate while the +// bounded successful LoginResp write runs, then transitions it to running. +// The callback must only perform that bounded write; it must not call back into +// the control manager or the same control lifecycle. +func (cm *ControlManager) completeLogin(ctl *Control, writeSuccess func() error) (bool, error) { + entry, ok := cm.lockCurrentRun(ctl.runID, false) + if !ok { + return false, nil + } + defer entry.runMu.Unlock() + if entry.ctl != ctl || entry.id != ctl.controlID { + return false, nil + } + + ctl.lifecycleMu.Lock() + defer ctl.lifecycleMu.Unlock() + if ctl.state != controlStatePending || !ctl.activated { + return false, nil + } + if err := writeSuccess(); err != nil { + return false, err + } + if !ctl.startLocked() { + return false, nil + } + return true, nil +} + +// Remove deletes and offlines ctl only if it is still the current generation. +func (cm *ControlManager) Remove(ctl *Control) bool { + entry, ok := cm.lockCurrentRun(ctl.runID, true) + if !ok { + return false + } + defer entry.runMu.Unlock() + cm.mu.Lock() + defer cm.mu.Unlock() + + if cm.ctlsByRunID[ctl.runID] != entry || entry.ctl != ctl || entry.id != ctl.controlID { + return false + } + delete(cm.ctlsByRunID, ctl.runID) + if entry.registryOnline { + cm.registry.MarkOfflineByRunIDAndControlID(ctl.runID, uint64(entry.registryControlID)) + } + return true } func (cm *ControlManager) GetByID(runID string) (ctl *Control, ok bool) { - cm.mu.RLock() - defer cm.mu.RUnlock() - ctl, ok = cm.ctlsByRunID[runID] - return + entry, ok := cm.lockCurrentRun(runID, false) + if !ok { + return nil, false + } + defer entry.runMu.Unlock() + ctl = entry.ctl + + ctl.lifecycleMu.Lock() + defer ctl.lifecycleMu.Unlock() + if ctl.state != controlStateRunning { + return nil, false + } + return ctl, true +} + +// admitVisitorByRunID commits a visitor admission against the current running +// control while its run and lifecycle ownership are held. The callback must +// only perform the in-memory, buffered visitor admission. +func (cm *ControlManager) admitVisitorByRunID(runID string, admit func(user string) error) (bool, error) { + entry, ok := cm.lockCurrentRun(runID, false) + if !ok { + return false, nil + } + defer entry.runMu.Unlock() + ctl := entry.ctl + + ctl.lifecycleMu.Lock() + defer ctl.lifecycleMu.Unlock() + if ctl.state != controlStateRunning { + return false, nil + } + return true, admit(ctl.sessionCtx.LoginMsg.User) +} + +// RegisterWorkConn transfers conn to ctl only if ctl is still the current +// running generation. On error, ownership remains with the caller. +func (cm *ControlManager) RegisterWorkConn(ctl *Control, conn *proxy.WorkConn) error { + entry, ok := cm.lockCurrentRun(ctl.runID, false) + if !ok { + cm.mu.RLock() + closed := cm.closed + cm.mu.RUnlock() + if closed { + return fmt.Errorf("control manager is closed") + } + return fmt.Errorf("client control for run id [%s] is no longer current", ctl.runID) + } + defer entry.runMu.Unlock() + if entry.ctl != ctl || entry.id != ctl.controlID { + return fmt.Errorf("client control for run id [%s] is no longer current", ctl.runID) + } + + ctl.lifecycleMu.Lock() + defer ctl.lifecycleMu.Unlock() + if ctl.state != controlStateRunning { + return fmt.Errorf("client control for run id [%s] is not running", ctl.runID) + } + + select { + case ctl.workConnCh <- conn: + ctl.xl.Debugf("new work connection registered") + return nil + default: + ctl.xl.Debugf("work connection pool is full, discarding") + return fmt.Errorf("work connection pool is full, discarding") + } } func (cm *ControlManager) Close() error { cm.mu.Lock() - defer cm.mu.Unlock() - for _, ctl := range cm.ctlsByRunID { - ctl.Close() + cm.closed = true + ctls := make([]*Control, 0, len(cm.ctlsByRunID)) + for _, entry := range cm.ctlsByRunID { + ctls = append(ctls, entry.ctl) + } + cm.mu.Unlock() + + for _, ctl := range ctls { + cm.Remove(ctl) + _ = ctl.Close() } - cm.ctlsByRunID = make(map[string]*Control) return nil } @@ -110,12 +367,20 @@ type SessionContext struct { LoginMsg *msg.Login // server configuration ServerCfg *v1.ServerConfig - // client registry - ClientRegistry *registry.ClientRegistry // negotiated wire protocol for this client session WireProtocol string } +type controlState uint8 + +const ( + controlStateCreated controlState = iota + controlStatePending + controlStateRunning + controlStateClosing + controlStateClosed +) + type Control struct { // session context sessionCtx *SessionContext @@ -142,30 +407,42 @@ type Control struct { // last time got the Ping message lastPing atomic.Value - // A new run id will be generated when a new client login. - // If run id got from login message has same run id, it means it's the same client, so we can - // replace old controller instantly. - runID string + // runID never changes during the lifetime of a control. controlID is assigned + // once by ControlManager and distinguishes same-runID generations. + runID string + controlID ControlID + manager *ControlManager + + lifecycleMu sync.Mutex + state controlState + activated bool + handoffBarrier <-chan struct{} + + interruptOnce sync.Once + interruptErr error mu sync.RWMutex - xl *xlog.Logger - ctx context.Context - doneCh chan struct{} + xl *xlog.Logger + ctx context.Context + doneCh chan struct{} + serverMetrics metrics.ServerMetrics } func NewControl(ctx context.Context, sessionCtx *SessionContext) (*Control, error) { poolCount := min(sessionCtx.LoginMsg.PoolCount, int(sessionCtx.ServerCfg.Transport.MaxPoolCount)) ctl := &Control{ - sessionCtx: sessionCtx, - workConnCh: make(chan *proxy.WorkConn, poolCount+10), - proxies: make(map[string]proxy.Proxy), - poolCount: poolCount, - portsUsedNum: 0, - runID: sessionCtx.LoginMsg.RunID, - xl: xlog.FromContextSafe(ctx), - ctx: ctx, - doneCh: make(chan struct{}), + sessionCtx: sessionCtx, + workConnCh: make(chan *proxy.WorkConn, poolCount+10), + proxies: make(map[string]proxy.Proxy), + poolCount: poolCount, + portsUsedNum: 0, + runID: sessionCtx.LoginMsg.RunID, + state: controlStateCreated, + xl: xlog.FromContextSafe(ctx), + ctx: ctx, + doneCh: make(chan struct{}), + serverMetrics: metrics.Server, } ctl.lastPing.Store(time.Now()) @@ -175,48 +452,121 @@ func NewControl(ctx context.Context, sessionCtx *SessionContext) (*Control, erro return ctl, nil } -// Start starts the control session workers after login succeeds. -func (ctl *Control) Start() { - go func() { - for i := 0; i < ctl.poolCount; i++ { - // ignore error here, that means that this control is closed - _ = ctl.msgDispatcher.Send(&msg.ReqWorkConn{}) - } - }() - go ctl.worker() +func (ctl *Control) RunID() string { + return ctl.runID } -func (ctl *Control) Close() error { - ctl.sessionCtx.Conn.Close() +func (ctl *Control) ID() ControlID { + ctl.lifecycleMu.Lock() + defer ctl.lifecycleMu.Unlock() + return ctl.controlID +} + +func (ctl *Control) admit(manager *ControlManager, id ControlID) error { + ctl.lifecycleMu.Lock() + defer ctl.lifecycleMu.Unlock() + if ctl.state != controlStateCreated { + return fmt.Errorf("control [%s] is not in created state", ctl.runID) + } + ctl.manager = manager + ctl.controlID = id + ctl.state = controlStatePending return nil } -func (ctl *Control) Replaced(newCtl *Control) { - xl := ctl.xl - xl.Infof("replaced by client [%s]", newCtl.runID) - ctl.runID = "" - ctl.sessionCtx.Conn.Close() +func (ctl *Control) setHandoffBarrier(barrier <-chan struct{}) { + ctl.lifecycleMu.Lock() + ctl.handoffBarrier = barrier + ctl.lifecycleMu.Unlock() } -func (ctl *Control) RegisterWorkConn(conn *proxy.WorkConn) error { - xl := ctl.xl - defer func() { - if err := recover(); err != nil { - xl.Errorf("panic error: %v", err) - xl.Errorf(string(debug.Stack())) - } - }() - - select { - case ctl.workConnCh <- conn: - xl.Debugf("new work connection registered") - return nil - default: - xl.Debugf("work connection pool is full, discarding") - return fmt.Errorf("work connection pool is full, discarding") +func (ctl *Control) WaitForHandoff() { + ctl.lifecycleMu.Lock() + barrier := ctl.handoffBarrier + ctl.lifecycleMu.Unlock() + if barrier != nil { + <-barrier } } +// Start starts the control session workers after login succeeds. +func (ctl *Control) Start() bool { + ctl.lifecycleMu.Lock() + defer ctl.lifecycleMu.Unlock() + return ctl.startLocked() +} + +func (ctl *Control) startLocked() bool { + if ctl.state != controlStatePending || !ctl.activated { + return false + } + ctl.state = controlStateRunning + go ctl.worker() + return true +} + +func (ctl *Control) Close() error { + ctl.lifecycleMu.Lock() + switch ctl.state { + case controlStateCreated, controlStatePending: + ctl.state = controlStateClosing + ctl.finishLocked() + case controlStateRunning: + ctl.state = controlStateClosing + } + ctl.lifecycleMu.Unlock() + return ctl.interruptReadAndClose() +} + +func (ctl *Control) Replaced(newCtl *Control) { + ctl.markReplaced() + ctl.xl.Infof("replaced by client [%s] (control ID %d)", newCtl.runID, newCtl.ID()) + _ = ctl.interruptReadAndClose() +} + +// markReplaced returns the transitive predecessor barrier. A pending control +// has no worker, so it finishes immediately and passes its inherited barrier +// to the replacement. A running control is finished only by its worker. +func (ctl *Control) markReplaced() <-chan struct{} { + ctl.lifecycleMu.Lock() + defer ctl.lifecycleMu.Unlock() + + switch ctl.state { + case controlStateCreated: + ctl.state = controlStateClosing + ctl.finishLocked() + return nil + case controlStatePending: + barrier := ctl.handoffBarrier + ctl.state = controlStateClosing + ctl.finishLocked() + return barrier + case controlStateRunning: + ctl.state = controlStateClosing + return ctl.doneCh + case controlStateClosing, controlStateClosed: + return ctl.doneCh + default: + return ctl.doneCh + } +} + +func (ctl *Control) interruptReadAndClose() error { + ctl.interruptOnce.Do(func() { + _ = ctl.sessionCtx.Conn.SetReadDeadline(time.Now()) + ctl.interruptErr = ctl.sessionCtx.Conn.Close() + }) + return ctl.interruptErr +} + +func (ctl *Control) finishLocked() { + if ctl.state == controlStateClosed { + return + } + ctl.state = controlStateClosed + close(ctl.doneCh) +} + // When frps get one user connection, we get one work connection from the pool and return it. // If no workConn available in the pool, send message to frpc to get one or more // and wait until it is available. @@ -271,10 +621,10 @@ func (ctl *Control) heartbeatWorker() { } xl := ctl.xl - go wait.Until(func() { + wait.Until(func() { if time.Since(ctl.lastPing.Load().(time.Time)) > time.Duration(ctl.sessionCtx.ServerCfg.Transport.HeartbeatTimeout)*time.Second { xl.Warnf("heartbeat timeout") - ctl.sessionCtx.Conn.Close() + _ = ctl.Close() return } }, time.Second, ctl.doneCh) @@ -289,14 +639,14 @@ func (ctl *Control) loginUserInfo() plugin.UserInfo { return plugin.UserInfo{ User: ctl.sessionCtx.LoginMsg.User, Metas: ctl.sessionCtx.LoginMsg.Metas, - RunID: ctl.sessionCtx.LoginMsg.RunID, + RunID: ctl.runID, } } func (ctl *Control) closeProxy(pxy proxy.Proxy) { pxy.Close() ctl.sessionCtx.PxyManager.Del(pxy.GetName()) - metrics.Server.CloseProxy(pxy.GetName(), pxy.GetConfigurer().GetBaseConfig().Type) + ctl.serverMetrics.CloseProxy(pxy.GetName(), pxy.GetConfigurer().GetBaseConfig().Type) notifyContent := &plugin.CloseProxyContent{ User: ctl.loginUserInfo(), @@ -311,12 +661,24 @@ func (ctl *Control) closeProxy(pxy proxy.Proxy) { func (ctl *Control) worker() { xl := ctl.xl + ctl.serverMetrics.NewClient() go ctl.heartbeatWorker() go ctl.msgDispatcher.Run() + go func() { + for i := 0; i < ctl.poolCount; i++ { + // Ignore the error: it means this control is already closing. + _ = ctl.msgDispatcher.Send(&msg.ReqWorkConn{}) + } + }() <-ctl.msgDispatcher.Done() - ctl.sessionCtx.Conn.Close() + ctl.lifecycleMu.Lock() + if ctl.state == controlStateRunning { + ctl.state = controlStateClosing + } + ctl.lifecycleMu.Unlock() + _ = ctl.interruptReadAndClose() ctl.mu.Lock() close(ctl.workConnCh) @@ -331,10 +693,14 @@ func (ctl *Control) worker() { ctl.closeProxy(pxy) } - metrics.Server.CloseClient() - ctl.sessionCtx.ClientRegistry.MarkOfflineByRunID(ctl.runID) + ctl.serverMetrics.CloseClient() + if ctl.manager != nil { + ctl.manager.Remove(ctl) + } xl.Infof("client exit success") - close(ctl.doneCh) + ctl.lifecycleMu.Lock() + ctl.finishLocked() + ctl.lifecycleMu.Unlock() } func (ctl *Control) registerMsgHandlers() { @@ -374,9 +740,9 @@ func (ctl *Control) handleNewProxy(m msg.Message) { xl.Infof("new proxy [%s] type [%s] success", inMsg.ProxyName, inMsg.ProxyType) clientID := ctl.sessionCtx.LoginMsg.ClientID if clientID == "" { - clientID = ctl.sessionCtx.LoginMsg.RunID + clientID = ctl.runID } - metrics.Server.NewProxy(inMsg.ProxyName, inMsg.ProxyType, ctl.sessionCtx.LoginMsg.User, clientID) + ctl.serverMetrics.NewProxy(inMsg.ProxyName, inMsg.ProxyType, ctl.sessionCtx.LoginMsg.User, clientID) } _ = ctl.msgDispatcher.Send(resp) } diff --git a/server/control_test.go b/server/control_test.go new file mode 100644 index 00000000..9dce3d24 --- /dev/null +++ b/server/control_test.go @@ -0,0 +1,543 @@ +// Copyright 2026 The frp Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "context" + "errors" + "net" + "os" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/fatedier/frp/pkg/auth" + v1 "github.com/fatedier/frp/pkg/config/v1" + "github.com/fatedier/frp/pkg/msg" + plugin "github.com/fatedier/frp/pkg/plugin/server" + "github.com/fatedier/frp/server/controller" + "github.com/fatedier/frp/server/proxy" + "github.com/fatedier/frp/server/registry" +) + +func TestControlPendingReplacementFinishesWithoutStarting(t *testing.T) { + clientRegistry := registry.NewClientRegistry() + manager := NewControlManager(clientRegistry) + metrics := newCountingServerMetrics() + oldCtl, oldConn := newLifecycleTestControl(t, "same-run", "client", metrics) + newCtl, _ := newLifecycleTestControl(t, "same-run", "client", metrics) + + mustAddAndActivate(t, manager, oldCtl) + + err := manager.Add(newCtl) + require.NoError(t, err) + waitForControlDone(t, oldCtl) + require.False(t, oldCtl.Start()) + require.Equal(t, []string{"deadline", "close"}, oldConn.eventsSnapshot()) + require.Equal(t, int64(0), metrics.newClients()) + require.Equal(t, int64(0), metrics.closedClients()) +} + +func TestControlRunningReplacementFinishesInWorker(t *testing.T) { + clientRegistry := registry.NewClientRegistry() + manager := NewControlManager(clientRegistry) + metrics := newCountingServerMetrics() + oldCtl, oldConn := newLifecycleTestControl(t, "same-run", "client", metrics) + newCtl, _ := newLifecycleTestControl(t, "same-run", "client", metrics) + + mustAddAndActivate(t, manager, oldCtl) + require.True(t, oldCtl.Start()) + waitForSignal(t, oldConn.readStarted, "control reader to start") + + err := manager.Add(newCtl) + require.NoError(t, err) + waitForControlDone(t, oldCtl) + require.Equal(t, []string{"deadline", "close"}, oldConn.eventsSnapshot()) + require.Equal(t, int64(1), metrics.newClients()) + require.Equal(t, int64(1), metrics.closedClients()) + + _, ok := manager.GetByID("same-run") + require.False(t, ok) + require.Same(t, newCtl, currentControlForTest(manager, "same-run")) + info, ok := clientRegistry.GetByKey("client") + require.True(t, ok) + require.True(t, info.Online) + require.Equal(t, uint64(oldCtl.ID()), info.ControlID) + + active, err := manager.Activate(newCtl) + require.NoError(t, err) + require.True(t, active) + _, ok = manager.GetByID("same-run") + require.False(t, ok) + info, ok = clientRegistry.GetByKey("client") + require.True(t, ok) + require.Equal(t, uint64(newCtl.ID()), info.ControlID) +} + +func TestControlClosePendingAndRunning(t *testing.T) { + t.Run("pending", func(t *testing.T) { + manager := NewControlManager(registry.NewClientRegistry()) + metrics := newCountingServerMetrics() + ctl, conn := newLifecycleTestControl(t, "pending", "pending", metrics) + err := manager.Add(ctl) + require.NoError(t, err) + + require.NoError(t, ctl.Close()) + waitForControlDone(t, ctl) + require.Equal(t, []string{"deadline", "close"}, conn.eventsSnapshot()) + require.Equal(t, int64(0), metrics.newClients()) + require.Equal(t, int64(0), metrics.closedClients()) + }) + + t.Run("running", func(t *testing.T) { + manager := NewControlManager(registry.NewClientRegistry()) + metrics := newCountingServerMetrics() + ctl, conn := newLifecycleTestControl(t, "running", "running", metrics) + mustAddAndActivate(t, manager, ctl) + require.True(t, ctl.Start()) + waitForSignal(t, conn.readStarted, "control reader to start") + + require.NoError(t, ctl.Close()) + waitForControlDone(t, ctl) + require.Equal(t, []string{"deadline", "close"}, conn.eventsSnapshot()) + require.Equal(t, int64(1), metrics.newClients()) + require.Equal(t, int64(1), metrics.closedClients()) + }) +} + +func TestControlCloseAndReplacedAreIdempotent(t *testing.T) { + manager := NewControlManager(registry.NewClientRegistry()) + metrics := newCountingServerMetrics() + ctl, conn := newLifecycleTestControl(t, "same-run", "client", metrics) + replacement, _ := newLifecycleTestControl(t, "same-run", "client", metrics) + + err := manager.Add(ctl) + require.NoError(t, err) + err = manager.Add(replacement) + require.NoError(t, err) + require.NoError(t, ctl.Close()) + ctl.Replaced(replacement) + require.NoError(t, ctl.Close()) + waitForControlDone(t, ctl) + + require.Equal(t, []string{"deadline", "close"}, conn.eventsSnapshot()) + require.Equal(t, int64(0), metrics.newClients()) + require.Equal(t, int64(0), metrics.closedClients()) +} + +func TestControlHeartbeatTimeoutInterruptsRead(t *testing.T) { + manager := NewControlManager(registry.NewClientRegistry()) + metrics := newCountingServerMetrics() + ctl, conn := newLifecycleTestControl(t, "heartbeat", "heartbeat", metrics) + ctl.sessionCtx.ServerCfg.Transport.HeartbeatTimeout = 1 + ctl.lastPing.Store(time.Now().Add(-2 * time.Second)) + + mustAddAndActivate(t, manager, ctl) + require.True(t, ctl.Start()) + waitForSignal(t, conn.readStarted, "control reader to start") + waitForControlDone(t, ctl) + + require.Equal(t, []string{"deadline", "close"}, conn.eventsSnapshot()) + require.Equal(t, int64(1), metrics.newClients()) + require.Equal(t, int64(1), metrics.closedClients()) +} + +func TestControlStartReplacementRacePairsMetrics(t *testing.T) { + for range 100 { + clientRegistry := registry.NewClientRegistry() + manager := NewControlManager(clientRegistry) + metrics := newCountingServerMetrics() + ctl, _ := newLifecycleTestControl(t, "same-run", "client", metrics) + replacement, _ := newLifecycleTestControl(t, "same-run", "client", metrics) + + mustAddAndActivate(t, manager, ctl) + + startGate := make(chan struct{}) + startedCh := make(chan bool, 1) + addErrCh := make(chan error, 1) + go func() { + <-startGate + startedCh <- ctl.Start() + }() + go func() { + <-startGate + addErr := manager.Add(replacement) + addErrCh <- addErr + }() + close(startGate) + + started := <-startedCh + require.NoError(t, <-addErrCh) + waitForControlDone(t, ctl) + if started { + require.Equal(t, int64(1), metrics.newClients()) + require.Equal(t, int64(1), metrics.closedClients()) + } else { + require.Equal(t, int64(0), metrics.newClients()) + require.Equal(t, int64(0), metrics.closedClients()) + } + } +} + +func TestControlManagerRejectsStaleActivateAndRemove(t *testing.T) { + clientRegistry := registry.NewClientRegistry() + manager := NewControlManager(clientRegistry) + metrics := newCountingServerMetrics() + oldCtl, _ := newLifecycleTestControl(t, "same-run", "client", metrics) + newCtl, _ := newLifecycleTestControl(t, "same-run", "client", metrics) + + mustAddAndActivate(t, manager, oldCtl) + err := manager.Add(newCtl) + require.NoError(t, err) + require.Greater(t, uint64(newCtl.ID()), uint64(oldCtl.ID())) + + active, err := manager.Activate(oldCtl) + require.NoError(t, err) + require.False(t, active) + require.False(t, manager.Remove(oldCtl)) + + _, ok := manager.GetByID("same-run") + require.False(t, ok) + require.Same(t, newCtl, currentControlForTest(manager, "same-run")) + info, ok := clientRegistry.GetByKey("client") + require.True(t, ok) + require.True(t, info.Online) + require.Equal(t, uint64(oldCtl.ID()), info.ControlID) + + active, err = manager.Activate(newCtl) + require.NoError(t, err) + require.True(t, active) + info, ok = clientRegistry.GetByKey("client") + require.True(t, ok) + require.True(t, info.Online) + require.Equal(t, uint64(newCtl.ID()), info.ControlID) +} + +func TestControlManagerPreservesClientIDConflict(t *testing.T) { + clientRegistry := registry.NewClientRegistry() + manager := NewControlManager(clientRegistry) + metrics := newCountingServerMetrics() + first, _ := newLifecycleTestControl(t, "run-one", "shared-client", metrics) + conflicting, _ := newLifecycleTestControl(t, "run-two", "shared-client", metrics) + + mustAddAndActivate(t, manager, first) + err := manager.Add(conflicting) + require.NoError(t, err) + active, err := manager.Activate(conflicting) + require.True(t, active) + require.ErrorContains(t, err, "already online") + + require.True(t, manager.Remove(conflicting)) + info, ok := clientRegistry.GetByKey("shared-client") + require.True(t, ok) + require.True(t, info.Online) + require.Equal(t, "run-one", info.RunID) +} + +func TestControlManagerFailedLoginWriteReleasesRunWithoutStarting(t *testing.T) { + clientRegistry := registry.NewClientRegistry() + manager := NewControlManager(clientRegistry) + metrics := newCountingServerMetrics() + ctl, _ := newLifecycleTestControl(t, "same-run", "client", metrics) + replacement, _ := newLifecycleTestControl(t, "same-run", "client", metrics) + + mustAddAndActivate(t, manager, ctl) + + writeErr := errors.New("write failed") + committed, err := manager.completeLogin(ctl, func() error { return writeErr }) + require.ErrorIs(t, err, writeErr) + require.False(t, committed) + + err = manager.Add(replacement) + require.NoError(t, err) + waitForControlDone(t, ctl) + require.Same(t, replacement, currentControlForTest(manager, "same-run")) + require.Equal(t, int64(0), metrics.newClients()) + require.Equal(t, int64(0), metrics.closedClients()) + require.True(t, manager.Remove(replacement)) + info, ok := clientRegistry.GetByKey("client") + require.True(t, ok) + require.False(t, info.Online) + require.Empty(t, info.RunID) + require.Zero(t, info.ControlID) + require.False(t, info.DisconnectedAt.IsZero()) + require.NoError(t, replacement.Close()) +} + +func TestControlManagerCloseWaitsForInFlightLoginRun(t *testing.T) { + clientRegistry := registry.NewClientRegistry() + manager := NewControlManager(clientRegistry) + metrics := newCountingServerMetrics() + ctl, _ := newLifecycleTestControl(t, "same-run", "client", metrics) + + mustAddAndActivate(t, manager, ctl) + + writeEntered := make(chan struct{}) + resumeWrite := make(chan struct{}) + loginDone := make(chan struct { + committed bool + err error + }, 1) + go func() { + committed, loginErr := manager.completeLogin(ctl, func() error { + close(writeEntered) + <-resumeWrite + return nil + }) + loginDone <- struct { + committed bool + err error + }{committed: committed, err: loginErr} + }() + waitForSignal(t, writeEntered, "LoginResp write") + + closeDone := make(chan error, 1) + go func() { closeDone <- manager.Close() }() + waitForManagerClosed(t, manager) + select { + case err := <-closeDone: + t.Fatalf("manager close completed during LoginResp write: %v", err) + default: + } + + close(resumeWrite) + result := <-loginDone + require.NoError(t, result.err) + require.True(t, result.committed) + require.NoError(t, <-closeDone) + waitForControlDone(t, ctl) + require.Nil(t, currentControlForTest(manager, "same-run")) + require.Equal(t, int64(1), metrics.newClients()) + require.Equal(t, int64(1), metrics.closedClients()) + info, ok := clientRegistry.GetByKey("client") + require.True(t, ok) + require.False(t, info.Online) +} + +func newLifecycleTestControl( + t *testing.T, + runID string, + clientID string, + serverMetrics *countingServerMetrics, +) (*Control, *deadlineReadConn) { + t.Helper() + conn := newDeadlineReadConn() + msgConn := msg.NewConn(conn, msg.NewV1ReadWriter(conn)) + ctl, err := NewControl(context.Background(), &SessionContext{ + RC: &controller.ResourceController{}, + PxyManager: proxy.NewManager(), + PluginManager: plugin.NewManager(), + AuthVerifier: auth.AlwaysPassVerifier, + Conn: msgConn, + LoginMsg: &msg.Login{ + RunID: runID, + ClientID: clientID, + }, + ServerCfg: &v1.ServerConfig{}, + }) + require.NoError(t, err) + ctl.serverMetrics = serverMetrics + t.Cleanup(func() { _ = ctl.Close() }) + return ctl, conn +} + +func mustAddAndActivate(t *testing.T, manager *ControlManager, ctl *Control) { + t.Helper() + require.NoError(t, manager.Add(ctl)) + active, err := manager.Activate(ctl) + require.NoError(t, err) + require.True(t, active) +} + +func waitForControlDone(t *testing.T, ctl *Control) { + t.Helper() + done := make(chan struct{}) + go func() { + ctl.WaitClosed() + close(done) + }() + waitForSignal(t, done, "control to finish") +} + +func currentControlForTest(manager *ControlManager, runID string) *Control { + manager.mu.RLock() + defer manager.mu.RUnlock() + entry := manager.ctlsByRunID[runID] + if entry == nil { + return nil + } + return entry.ctl +} + +func currentRunGateForTest(manager *ControlManager, runID string) *sync.Mutex { + manager.mu.RLock() + defer manager.mu.RUnlock() + entry := manager.ctlsByRunID[runID] + if entry == nil { + return nil + } + return entry.runMu +} + +func waitForManagerClosed(t *testing.T, manager *ControlManager) { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + manager.mu.RLock() + closed := manager.closed + manager.mu.RUnlock() + if closed { + return + } + } + t.Fatal("timed out waiting for control manager to close") +} + +func waitForSignal(t *testing.T, ch <-chan struct{}, description string) { + t.Helper() + select { + case <-ch: + case <-time.After(3 * time.Second): + t.Fatalf("timed out waiting for %s", description) + } +} + +type deadlineReadConn struct { + readStarted chan struct{} + unblockRead chan struct{} + + readOnce sync.Once + unblockOnce sync.Once + deadlineOnce sync.Once + closeOnce sync.Once + + eventsMu sync.Mutex + events []string +} + +func newDeadlineReadConn() *deadlineReadConn { + return &deadlineReadConn{ + readStarted: make(chan struct{}), + unblockRead: make(chan struct{}), + } +} + +func (c *deadlineReadConn) Read([]byte) (int, error) { + c.readOnce.Do(func() { close(c.readStarted) }) + <-c.unblockRead + return 0, os.ErrDeadlineExceeded +} + +func (*deadlineReadConn) Write(p []byte) (int, error) { return len(p), nil } + +func (c *deadlineReadConn) Close() error { + c.closeOnce.Do(func() { + c.recordEvent("close") + c.unblockOnce.Do(func() { close(c.unblockRead) }) + }) + return nil +} + +func (*deadlineReadConn) LocalAddr() net.Addr { return lifecycleTestAddr("local") } +func (*deadlineReadConn) RemoteAddr() net.Addr { return lifecycleTestAddr("remote") } + +func (c *deadlineReadConn) SetDeadline(deadline time.Time) error { + if err := c.SetReadDeadline(deadline); err != nil { + return err + } + return c.SetWriteDeadline(deadline) +} + +func (c *deadlineReadConn) SetReadDeadline(deadline time.Time) error { + if deadline.IsZero() { + return nil + } + c.deadlineOnce.Do(func() { + c.recordEvent("deadline") + c.unblockOnce.Do(func() { close(c.unblockRead) }) + }) + return nil +} + +func (*deadlineReadConn) SetWriteDeadline(time.Time) error { return nil } + +func (c *deadlineReadConn) recordEvent(event string) { + c.eventsMu.Lock() + c.events = append(c.events, event) + c.eventsMu.Unlock() +} + +func (c *deadlineReadConn) eventsSnapshot() []string { + c.eventsMu.Lock() + defer c.eventsMu.Unlock() + return append([]string(nil), c.events...) +} + +type lifecycleTestAddr string + +func (a lifecycleTestAddr) Network() string { return string(a) } +func (a lifecycleTestAddr) String() string { return string(a) } + +type countingServerMetrics struct { + mu sync.Mutex + newCount int64 + closeCount int64 + closeEnter chan struct{} + closeResume chan struct{} + closeOnce sync.Once +} + +func newCountingServerMetrics() *countingServerMetrics { + return &countingServerMetrics{} +} + +func (m *countingServerMetrics) NewClient() { + m.mu.Lock() + m.newCount++ + m.mu.Unlock() +} + +func (m *countingServerMetrics) CloseClient() { + m.mu.Lock() + m.closeCount++ + closeEnter := m.closeEnter + closeResume := m.closeResume + m.mu.Unlock() + if closeEnter != nil { + m.closeOnce.Do(func() { close(closeEnter) }) + <-closeResume + } +} + +func (*countingServerMetrics) NewProxy(string, string, string, string) {} +func (*countingServerMetrics) CloseProxy(string, string) {} +func (*countingServerMetrics) OpenConnection(string, string) {} +func (*countingServerMetrics) CloseConnection(string, string) {} +func (*countingServerMetrics) AddTrafficIn(string, string, int64) {} +func (*countingServerMetrics) AddTrafficOut(string, string, int64) {} + +func (m *countingServerMetrics) newClients() int64 { + m.mu.Lock() + defer m.mu.Unlock() + return m.newCount +} + +func (m *countingServerMetrics) closedClients() int64 { + m.mu.Lock() + defer m.mu.Unlock() + return m.closeCount +} diff --git a/server/registry/registry.go b/server/registry/registry.go index 21771bce..a3632f5d 100644 --- a/server/registry/registry.go +++ b/server/registry/registry.go @@ -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. diff --git a/server/registry/registry_test.go b/server/registry/registry_test.go index 0ff083b8..bacac949 100644 --- a/server/registry/registry_test.go +++ b/server/registry/registry_test.go @@ -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) + } +} diff --git a/server/service.go b/server/service.go index f3f9e323..39399ff9 100644 --- a/server/service.go +++ b/server/service.go @@ -18,6 +18,7 @@ import ( "bytes" "context" "crypto/tls" + "errors" "fmt" "io" "net" @@ -51,7 +52,6 @@ import ( "github.com/fatedier/frp/pkg/util/xlog" "github.com/fatedier/frp/server/controller" "github.com/fatedier/frp/server/group" - "github.com/fatedier/frp/server/metrics" "github.com/fatedier/frp/server/ports" "github.com/fatedier/frp/server/proxy" "github.com/fatedier/frp/server/registry" @@ -64,6 +64,8 @@ const ( vhostReadWriteTimeout time.Duration = 30 * time.Second ) +var errControlReplaced = errors.New("control was replaced during login") + func init() { crypto.DefaultSalt = "frp" // Disable quic-go's receive buffer warning. @@ -161,9 +163,10 @@ func NewService(cfg *v1.ServerConfig) (*Service, error) { return nil, err } + clientRegistry := registry.NewClientRegistry() svr := &Service{ - ctlManager: NewControlManager(), - clientRegistry: registry.NewClientRegistry(), + ctlManager: NewControlManager(clientRegistry), + clientRegistry: clientRegistry, pxyManager: proxy.NewManager(), pluginManager: plugin.NewManager(), rc: &controller.ResourceController{ @@ -469,6 +472,9 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, interna if err != nil { xl.Warnf("register control error: %v", err) + if ctl != nil { + svr.ctlManager.Remove(ctl) + } if writeErr := writeWithDeadline(conn, connWriteTimeout, func() error { return acceptedConn.conn.WriteMsg(&msg.LoginResp{ Version: version.Full(), @@ -477,29 +483,27 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, interna }); writeErr != nil { xl.Warnf("write login error response error: %v", writeErr) } - conn.Close() + if ctl != nil { + _ = ctl.Close() + } else { + conn.Close() + } return } - if err = writeWithDeadline(conn, connWriteTimeout, func() error { - return acceptedConn.conn.WriteMsg(&msg.LoginResp{ - Version: version.Full(), - RunID: ctl.runID, - Error: "", + if err = svr.completeControlLogin(ctl, func() error { + return writeWithDeadline(conn, connWriteTimeout, func() error { + return acceptedConn.conn.WriteMsg(&msg.LoginResp{ + Version: version.Full(), + RunID: ctl.runID, + Error: "", + }) }) }); err != nil { - xl.Warnf("write login response error: %v", err) - svr.ctlManager.Del(m.RunID, ctl) - svr.clientRegistry.MarkOfflineByRunID(m.RunID) - conn.Close() + xl.Warnf("complete control login error: %v", err) + svr.ctlManager.Remove(ctl) + _ = ctl.Close() return } - ctl.Start() - metrics.Server.NewClient() - go func() { - // block until control closed - ctl.WaitClosed() - svr.ctlManager.Del(m.RunID, ctl) - }() case *msg.NewWorkConn: if err := svr.RegisterWorkConn(acceptedConn.conn, m); err != nil { _ = acceptedConn.conn.WriteMsg(&msg.StartWorkConn{ @@ -527,6 +531,17 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, interna } } +func (svr *Service) completeControlLogin(ctl *Control, writeSuccess func() error) error { + committed, err := svr.ctlManager.completeLogin(ctl, writeSuccess) + if err != nil { + return err + } + if !committed { + return errControlReplaced + } + return nil +} + type acceptedConnection struct { conn *msg.Conn wireProtocol string @@ -768,16 +783,15 @@ func (svr *Service) RegisterControl( } ctl, err := NewControl(ctx, &SessionContext{ - RC: svr.rc, - PxyManager: svr.pxyManager, - PluginManager: svr.pluginManager, - AuthVerifier: authVerifier, - EncryptionKey: svr.auth.EncryptionKey(), - Conn: ctlConn, - LoginMsg: loginMsg, - ServerCfg: svr.cfg, - ClientRegistry: svr.clientRegistry, - WireProtocol: wireProtocol, + RC: svr.rc, + PxyManager: svr.pxyManager, + PluginManager: svr.pluginManager, + AuthVerifier: authVerifier, + EncryptionKey: svr.auth.EncryptionKey(), + Conn: ctlConn, + LoginMsg: loginMsg, + ServerCfg: svr.cfg, + WireProtocol: wireProtocol, }) if err != nil { xl.Warnf("create new controller error: %v", err) @@ -785,18 +799,17 @@ func (svr *Service) RegisterControl( return nil, fmt.Errorf("unexpected error when creating new controller") } - if oldCtl := svr.ctlManager.Add(loginMsg.RunID, ctl); oldCtl != nil { - oldCtl.WaitClosed() + if err := svr.ctlManager.Add(ctl); err != nil { + return ctl, err } + ctl.WaitForHandoff() - remoteAddr := ctlConn.RemoteAddr().String() - if host, _, err := net.SplitHostPort(remoteAddr); err == nil { - remoteAddr = host + active, err := svr.ctlManager.Activate(ctl) + if err != nil { + return ctl, err } - _, conflict := svr.clientRegistry.Register(loginMsg.User, loginMsg.ClientID, loginMsg.RunID, loginMsg.Hostname, loginMsg.Version, remoteAddr, wireProtocol) - if conflict { - svr.ctlManager.Del(loginMsg.RunID, ctl) - return nil, fmt.Errorf("client_id [%s] for user [%s] is already online", loginMsg.ClientID, loginMsg.User) + if !active { + return ctl, errControlReplaced } return ctl, nil @@ -830,20 +843,25 @@ func (svr *Service) RegisterWorkConn(workConn *msg.Conn, newMsg *msg.NewWorkConn xl.Warnf("invalid NewWorkConn with run id [%s]", newMsg.RunID) return err } - return ctl.RegisterWorkConn(proxy.NewWorkConn(workConn)) + return svr.ctlManager.RegisterWorkConn(ctl, proxy.NewWorkConn(workConn)) } func (svr *Service) RegisterVisitorConn(visitorConn net.Conn, newMsg *msg.NewVisitorConn, wireProtocol string) error { - visitorUser := "" + admit := func(visitorUser string) error { + return svr.rc.VisitorManager.NewConn(newMsg.ProxyName, visitorConn, newMsg.Timestamp, newMsg.SignKey, + newMsg.UseEncryption, newMsg.UseCompression, visitorUser, wireProtocol) + } // TODO(deprecation): Compatible with old versions, can be without runID, user is empty. In later versions, it will be mandatory to include runID. // If runID is required, it is not compatible with versions prior to v0.50.0. if newMsg.RunID != "" { - ctl, exist := svr.ctlManager.GetByID(newMsg.RunID) - if !exist { + admitted, err := svr.ctlManager.admitVisitorByRunID(newMsg.RunID, admit) + if err != nil { + return err + } + if !admitted { return fmt.Errorf("no client control found for run id [%s]", newMsg.RunID) } - visitorUser = ctl.sessionCtx.LoginMsg.User + return nil } - return svr.rc.VisitorManager.NewConn(newMsg.ProxyName, visitorConn, newMsg.Timestamp, newMsg.SignKey, - newMsg.UseEncryption, newMsg.UseCompression, visitorUser, wireProtocol) + return admit("") } diff --git a/server/service_test.go b/server/service_test.go index cf5eec5e..10b1a3dd 100644 --- a/server/service_test.go +++ b/server/service_test.go @@ -15,12 +15,27 @@ package server import ( + "context" "errors" "net" + "runtime" + "sync" + "sync/atomic" "testing" "time" "github.com/stretchr/testify/require" + + "github.com/fatedier/frp/pkg/auth" + v1 "github.com/fatedier/frp/pkg/config/v1" + "github.com/fatedier/frp/pkg/msg" + plugin "github.com/fatedier/frp/pkg/plugin/server" + "github.com/fatedier/frp/pkg/proto/wire" + "github.com/fatedier/frp/pkg/util/util" + "github.com/fatedier/frp/server/controller" + "github.com/fatedier/frp/server/proxy" + "github.com/fatedier/frp/server/registry" + "github.com/fatedier/frp/server/visitor" ) func TestWriteWithDeadlineTimesOutAndClearsDeadline(t *testing.T) { @@ -61,3 +76,504 @@ func TestWriteWithDeadlineTimesOutAndClearsDeadline(t *testing.T) { t.Fatal("timed out waiting for write after deadline reset") } } + +func TestServiceControlHandoffSkipsStalePendingGeneration(t *testing.T) { + svr := newControlTestService(t) + metrics := newCountingServerMetrics() + metrics.closeEnter = make(chan struct{}) + metrics.closeResume = make(chan struct{}) + + ctlA, connA, err := registerLifecycleTestControl(svr) + require.NoError(t, err) + ctlA.serverMetrics = metrics + require.NoError(t, svr.completeControlLogin(ctlA, func() error { return nil })) + waitForSignal(t, connA.readStarted, "A reader to start") + + require.NoError(t, ctlA.Close()) + waitForSignal(t, metrics.closeEnter, "A finalization barrier") + + type registerResult struct { + ctl *Control + conn *deadlineReadConn + err error + } + resultB := make(chan registerResult, 1) + go func() { + ctl, conn, registerErr := registerLifecycleTestControl(svr) + resultB <- registerResult{ctl: ctl, conn: conn, err: registerErr} + }() + ctlB := waitForDifferentCurrentControl(t, svr.ctlManager, "shared-run", ctlA) + ctlB.serverMetrics = metrics + + resultC := make(chan registerResult, 1) + go func() { + ctl, conn, registerErr := registerLifecycleTestControl(svr) + resultC <- registerResult{ctl: ctl, conn: conn, err: registerErr} + }() + ctlC := waitForDifferentCurrentControl(t, svr.ctlManager, "shared-run", ctlB) + ctlC.serverMetrics = metrics + waitForControlDone(t, ctlB) + + select { + case result := <-resultB: + t.Fatalf("B returned before A finalized: %v", result.err) + default: + } + select { + case result := <-resultC: + t.Fatalf("C returned before A finalized: %v", result.err) + default: + } + + close(metrics.closeResume) + waitForControlDone(t, ctlA) + + b := <-resultB + require.Same(t, ctlB, b.ctl) + require.ErrorIs(t, b.err, errControlReplaced) + require.False(t, svr.ctlManager.Remove(ctlB)) + require.NoError(t, ctlB.Close()) + + c := <-resultC + require.NoError(t, c.err) + require.Same(t, ctlC, c.ctl) + _, ok := svr.ctlManager.GetByID("shared-run") + require.False(t, ok) + require.Same(t, ctlC, currentControlForTest(svr.ctlManager, "shared-run")) + + info, ok := svr.clientRegistry.GetByKey("client") + require.True(t, ok) + require.True(t, info.Online) + require.Equal(t, uint64(ctlC.ID()), info.ControlID) + + var staleWrites atomic.Int64 + err = svr.completeControlLogin(ctlB, func() error { + staleWrites.Add(1) + return nil + }) + require.ErrorIs(t, err, errControlReplaced) + require.Equal(t, int64(0), staleWrites.Load()) + + require.NoError(t, svr.completeControlLogin(ctlC, func() error { return nil })) + waitForSignal(t, c.conn.readStarted, "C reader to start") + current, ok := svr.ctlManager.GetByID("shared-run") + require.True(t, ok) + require.Same(t, ctlC, current) + require.Equal(t, int64(2), metrics.newClients()) + require.Equal(t, int64(1), metrics.closedClients()) + + require.NoError(t, ctlC.Close()) + waitForControlDone(t, ctlC) + require.Equal(t, int64(2), metrics.newClients()) + require.Equal(t, int64(2), metrics.closedClients()) + _, ok = svr.ctlManager.GetByID("shared-run") + require.False(t, ok) +} + +func TestServiceLoginResponseSynchronizationIsScopedToRun(t *testing.T) { + svr := newControlTestService(t) + metrics := newCountingServerMetrics() + ctlA, connA, err := registerLifecycleTestControl(svr) + require.NoError(t, err) + ctlA.serverMetrics = metrics + + writeEntered := make(chan struct{}) + resumeWrite := make(chan struct{}) + var resumeWriteOnce sync.Once + resume := func() { + resumeWriteOnce.Do(func() { close(resumeWrite) }) + } + t.Cleanup(resume) + writeCount := atomic.Int64{} + loginDone := make(chan error, 1) + go func() { + loginDone <- svr.completeControlLogin(ctlA, func() error { + close(writeEntered) + <-resumeWrite + writeCount.Add(1) + return nil + }) + }() + waitForSignal(t, writeEntered, "A LoginResp write") + + runMu := currentRunGateForTest(svr.ctlManager, "shared-run") + require.NotNil(t, runMu) + if !svr.ctlManager.mu.TryLock() { + t.Fatal("ControlManager mutex was held while LoginResp write was in progress") + } + svr.ctlManager.mu.Unlock() + + ctlB, connB := newLifecycleTestControl(t, "shared-run", "client", metrics) + gateAvailable := make(chan bool) + addDone := make(chan error, 1) + go func() { + if runMu.TryLock() { + runMu.Unlock() + gateAvailable <- true + } else { + gateAvailable <- false + } + addErr := svr.ctlManager.Add(ctlB) + addDone <- addErr + }() + available := waitForResult(t, gateAvailable, "same-run replacement gate probe") + require.False(t, available, "same-run gate was available to replacement during LoginResp write") + select { + case addErr := <-addDone: + t.Fatalf("same-run replacement completed during LoginResp write: %v", addErr) + case <-time.After(20 * time.Millisecond): + } + require.Same(t, ctlA, currentControlForTest(svr.ctlManager, "shared-run")) + + otherMetrics := newCountingServerMetrics() + otherCtl, otherConn := newLifecycleTestControl(t, "other-run", "other-client", otherMetrics) + type unrelatedResult struct { + addErr error + active bool + activateErr error + loginErr error + current *Control + found bool + } + unrelatedDone := make(chan unrelatedResult, 1) + go func() { + result := unrelatedResult{} + result.addErr = svr.ctlManager.Add(otherCtl) + if result.addErr == nil { + result.active, result.activateErr = svr.ctlManager.Activate(otherCtl) + } + if result.activateErr == nil && result.active { + result.loginErr = svr.completeControlLogin(otherCtl, func() error { return nil }) + } + result.current, result.found = svr.ctlManager.GetByID("other-run") + unrelatedDone <- result + }() + result := waitForResult(t, unrelatedDone, "unrelated run lifecycle") + require.NoError(t, result.addErr) + require.NoError(t, result.activateErr) + require.True(t, result.active) + require.NoError(t, result.loginErr) + require.True(t, result.found) + require.Same(t, otherCtl, result.current) + waitForSignal(t, otherConn.readStarted, "unrelated control reader to start") + require.Equal(t, int64(1), otherMetrics.newClients()) + + resume() + require.NoError(t, waitForResult(t, loginDone, "LoginResp completion")) + require.NoError(t, waitForResult(t, addDone, "replacement")) + waitForControlDone(t, ctlA) + require.Same(t, ctlB, currentControlForTest(svr.ctlManager, "shared-run")) + require.Equal(t, int64(1), writeCount.Load()) + require.Equal(t, int64(1), metrics.newClients()) + require.Equal(t, int64(1), metrics.closedClients()) + require.Equal(t, []string{"deadline", "close"}, connA.eventsSnapshot()) + + require.False(t, svr.ctlManager.Remove(ctlA)) + require.NoError(t, ctlA.Close()) + require.True(t, svr.ctlManager.Remove(ctlB)) + require.NoError(t, ctlB.Close()) + require.Equal(t, []string{"deadline", "close"}, connB.eventsSnapshot()) + + require.NoError(t, otherCtl.Close()) + waitForControlDone(t, otherCtl) + require.Equal(t, int64(1), otherMetrics.newClients()) + require.Equal(t, int64(1), otherMetrics.closedClients()) +} + +func TestServiceVisitorAdmissionSerializesReplacement(t *testing.T) { + svr := newControlTestService(t) + ctlA, controlConn, err := registerLifecycleTestControl(svr) + require.NoError(t, err) + ctlA.sessionCtx.LoginMsg.User = "old-user" + require.NoError(t, svr.completeControlLogin(ctlA, func() error { return nil })) + waitForSignal(t, controlConn.readStarted, "A reader to start") + + admissionEntered := make(chan struct{}) + resumeAdmission := make(chan struct{}) + var resumeOnce sync.Once + resume := func() { + resumeOnce.Do(func() { close(resumeAdmission) }) + } + t.Cleanup(resume) + type admissionResult struct { + admitted bool + user string + err error + } + admissionDone := make(chan admissionResult, 1) + go func() { + var admittedUser string + admitted, admitErr := svr.ctlManager.admitVisitorByRunID("shared-run", func(user string) error { + admittedUser = user + close(admissionEntered) + <-resumeAdmission + return nil + }) + admissionDone <- admissionResult{admitted: admitted, user: admittedUser, err: admitErr} + }() + waitForSignal(t, admissionEntered, "visitor admission callback") + runMu := currentRunGateForTest(svr.ctlManager, "shared-run") + require.NotNil(t, runMu) + + type registerResult struct { + ctl *Control + err error + } + gateAvailable := make(chan bool) + replacementDone := make(chan registerResult, 1) + go func() { + if runMu.TryLock() { + runMu.Unlock() + gateAvailable <- true + } else { + gateAvailable <- false + } + ctl, _, registerErr := registerLifecycleTestControl(svr) + replacementDone <- registerResult{ctl: ctl, err: registerErr} + }() + available := waitForResult(t, gateAvailable, "visitor replacement gate probe") + require.False(t, available, "same-run gate was available during visitor admission") + select { + case result := <-replacementDone: + t.Fatalf("replacement completed during visitor admission: %v", result.err) + case <-time.After(20 * time.Millisecond): + } + require.Same(t, ctlA, currentControlForTest(svr.ctlManager, "shared-run")) + + resume() + admission := waitForResult(t, admissionDone, "visitor admission") + require.NoError(t, admission.err) + require.True(t, admission.admitted) + require.Equal(t, "old-user", admission.user) + replacement := waitForResult(t, replacementDone, "replacement") + require.NoError(t, replacement.err) + ctlB := replacement.ctl + require.Same(t, ctlB, currentControlForTest(svr.ctlManager, "shared-run")) + waitForControlDone(t, ctlA) + require.True(t, svr.ctlManager.Remove(ctlB)) + require.NoError(t, ctlB.Close()) +} + +func TestServiceWorkConnRoutingRequiresCurrentRunningControl(t *testing.T) { + svr := newControlTestService(t) + ctl, controlConn, err := registerLifecycleTestControl(svr) + require.NoError(t, err) + + pendingConn := newCountingCloseConn() + pendingMsgConn := msg.NewConn(pendingConn, msg.NewV1ReadWriter(pendingConn)) + err = registerWorkConnAsCaller(svr, pendingMsgConn, &msg.NewWorkConn{RunID: "shared-run"}) + require.Error(t, err) + require.Equal(t, int64(1), pendingConn.closeCount.Load()) + require.Len(t, ctl.workConnCh, 0) + + require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil })) + waitForSignal(t, controlConn.readStarted, "control reader to start") + current, ok := svr.ctlManager.GetByID("shared-run") + require.True(t, ok) + require.Same(t, ctl, current) + require.Len(t, ctl.workConnCh, 0) + + runningConn := newCountingCloseConn() + runningMsgConn := msg.NewConn(runningConn, msg.NewV1ReadWriter(runningConn)) + require.NoError(t, svr.RegisterWorkConn(runningMsgConn, &msg.NewWorkConn{RunID: "shared-run"})) + require.Len(t, ctl.workConnCh, 1) + + require.NoError(t, ctl.Close()) + waitForControlDone(t, ctl) + require.Equal(t, int64(1), runningConn.closeCount.Load()) +} + +func TestServiceWorkConnRoutingRejectsLostGeneration(t *testing.T) { + for _, action := range []string{"replace", "close"} { + t.Run(action, func(t *testing.T) { + svr := newControlTestService(t) + ctl, controlConn, err := registerLifecycleTestControl(svr) + require.NoError(t, err) + require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil })) + waitForSignal(t, controlConn.readStarted, "control reader to start") + + barrier := newWorkConnBarrierPlugin() + svr.pluginManager.Register(barrier) + workConn := newCountingCloseConn() + workMsgConn := msg.NewConn(workConn, msg.NewV1ReadWriter(workConn)) + routeDone := make(chan error, 1) + go func() { + routeDone <- registerWorkConnAsCaller(svr, workMsgConn, &msg.NewWorkConn{RunID: "shared-run"}) + }() + waitForSignal(t, barrier.entered, "work connection plugin barrier") + + var replacement *Control + switch action { + case "replace": + replacement, _, err = registerLifecycleTestControl(svr) + require.NoError(t, err) + case "close": + require.NoError(t, ctl.Close()) + waitForControlDone(t, ctl) + } + + close(barrier.resume) + require.Error(t, waitForResult(t, routeDone, "work connection route to finish")) + require.Equal(t, int64(1), workConn.closeCount.Load()) + require.Len(t, ctl.workConnCh, 0) + + if replacement != nil { + require.Len(t, replacement.workConnCh, 0) + require.True(t, svr.ctlManager.Remove(replacement)) + require.NoError(t, replacement.Close()) + waitForControlDone(t, replacement) + } + }) + } +} + +func TestServiceVisitorRoutingExcludesPendingUser(t *testing.T) { + svr := newControlTestService(t) + listener, err := svr.rc.VisitorManager.Listen("visitor", "secret", []string{"pending-user"}) + require.NoError(t, err) + t.Cleanup(func() { _ = listener.Close() }) + + controlConn := newDeadlineReadConn() + controlMsgConn := msg.NewConn(controlConn, msg.NewV1ReadWriter(controlConn)) + ctl, err := svr.RegisterControl(controlMsgConn, &msg.Login{ + RunID: "visitor-run", + User: "pending-user", + ClientID: "visitor-client", + ClientSpec: msg.ClientSpec{ + AlwaysAuthPass: true, + }, + }, true, wire.ProtocolV1) + require.NoError(t, err) + + timestamp := time.Now().Unix() + visitorMsg := &msg.NewVisitorConn{ + RunID: "visitor-run", + ProxyName: "visitor", + Timestamp: timestamp, + SignKey: util.GetAuthKey("secret", timestamp), + } + pendingConn := newCountingCloseConn() + err = svr.RegisterVisitorConn(pendingConn, visitorMsg, wire.ProtocolV1) + require.ErrorContains(t, err, "no client control found") + require.NoError(t, pendingConn.Close()) + require.Equal(t, int64(1), pendingConn.closeCount.Load()) + + require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil })) + waitForSignal(t, controlConn.readStarted, "control reader to start") + runningConn := newCountingCloseConn() + require.NoError(t, svr.RegisterVisitorConn(runningConn, visitorMsg, wire.ProtocolV1)) + accepted, err := listener.Accept() + require.NoError(t, err) + require.NoError(t, accepted.Close()) + require.Equal(t, int64(1), runningConn.closeCount.Load()) + + require.NoError(t, ctl.Close()) + waitForControlDone(t, ctl) +} + +func newControlTestService(t *testing.T) *Service { + t.Helper() + cfg := &v1.ServerConfig{} + cfg.Auth.Method = v1.AuthMethodToken + authRuntime, err := auth.BuildServerAuth(&cfg.Auth) + require.NoError(t, err) + clientRegistry := registry.NewClientRegistry() + return &Service{ + ctlManager: NewControlManager(clientRegistry), + clientRegistry: clientRegistry, + pxyManager: proxy.NewManager(), + pluginManager: plugin.NewManager(), + rc: &controller.ResourceController{ + VisitorManager: visitor.NewManager(), + }, + auth: authRuntime, + cfg: cfg, + } +} + +func registerLifecycleTestControl(svr *Service) (*Control, *deadlineReadConn, error) { + conn := newDeadlineReadConn() + msgConn := msg.NewConn(conn, msg.NewReadWriter(conn, wire.ProtocolV1)) + ctl, err := svr.RegisterControl(msgConn, &msg.Login{ + RunID: "shared-run", + ClientID: "client", + ClientSpec: msg.ClientSpec{ + AlwaysAuthPass: true, + }, + }, true, wire.ProtocolV1) + return ctl, conn, err +} + +func waitForDifferentCurrentControl(t *testing.T, manager *ControlManager, runID string, old *Control) *Control { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if ctl := currentControlForTest(manager, runID); ctl != nil && ctl != old { + return ctl + } + runtime.Gosched() + } + t.Fatalf("timed out waiting for a new current control after ID %d", old.ID()) + return nil +} + +func registerWorkConnAsCaller(svr *Service, workConn *msg.Conn, newMsg *msg.NewWorkConn) error { + err := svr.RegisterWorkConn(workConn, newMsg) + if err != nil { + _ = workConn.Close() + } + return err +} + +func waitForResult[T any](t *testing.T, ch <-chan T, description string) T { + t.Helper() + select { + case result := <-ch: + return result + case <-time.After(3 * time.Second): + t.Fatalf("timed out waiting for %s", description) + var zero T + return zero + } +} + +type workConnBarrierPlugin struct { + entered chan struct{} + resume chan struct{} +} + +func newWorkConnBarrierPlugin() *workConnBarrierPlugin { + return &workConnBarrierPlugin{ + entered: make(chan struct{}), + resume: make(chan struct{}), + } +} + +func (*workConnBarrierPlugin) Name() string { return "work-conn-barrier" } + +func (*workConnBarrierPlugin) IsSupport(op string) bool { return op == plugin.OpNewWorkConn } + +func (p *workConnBarrierPlugin) Handle( + context.Context, + string, + any, +) (*plugin.Response, any, error) { + close(p.entered) + <-p.resume + return &plugin.Response{Unchange: true}, nil, nil +} + +type countingCloseConn struct { + closeCount atomic.Int64 +} + +func newCountingCloseConn() *countingCloseConn { return &countingCloseConn{} } + +func (*countingCloseConn) Read([]byte) (int, error) { return 0, net.ErrClosed } +func (*countingCloseConn) Write(p []byte) (int, error) { return len(p), nil } +func (c *countingCloseConn) Close() error { c.closeCount.Add(1); return nil } +func (*countingCloseConn) LocalAddr() net.Addr { return lifecycleTestAddr("local") } +func (*countingCloseConn) RemoteAddr() net.Addr { return lifecycleTestAddr("remote") } +func (*countingCloseConn) SetDeadline(time.Time) error { return nil } +func (*countingCloseConn) SetReadDeadline(time.Time) error { return nil } +func (*countingCloseConn) SetWriteDeadline(time.Time) error { return nil } diff --git a/test/e2e/framework/process.go b/test/e2e/framework/process.go index 8d27e887..ffbec33b 100644 --- a/test/e2e/framework/process.go +++ b/test/e2e/framework/process.go @@ -94,11 +94,9 @@ func (f *Framework) RunProcessesWithBinaries( } func (f *Framework) RunFrps(args ...string) (*process.Process, string, error) { - p := process.NewWithEnvs(TestContext.FRPServerPath, args, f.osEnvs) - f.serverProcesses = append(f.serverProcesses, p) - err := p.Start() + p, output, err := f.StartFrps(args...) if err != nil { - return p, p.Output(), err + return p, output, err } select { case <-p.Done(): @@ -107,17 +105,39 @@ func (f *Framework) RunFrps(args ...string) (*process.Process, string, error) { return p, p.Output(), nil } +// StartFrps starts frps without an implicit sleep so tests can wait on an +// explicit readiness event. +func (f *Framework) StartFrps(args ...string) (*process.Process, string, error) { + p := process.NewWithEnvs(TestContext.FRPServerPath, args, f.osEnvs) + f.serverProcesses = append(f.serverProcesses, p) + err := p.Start() + if err != nil { + return p, p.Output(), err + } + return p, p.Output(), nil +} + func (f *Framework) RunFrpc(args ...string) (*process.Process, string, error) { + p, output, err := f.StartFrpc(args...) + if err != nil { + return p, output, err + } + select { + case <-p.Done(): + case <-time.After(1500 * time.Millisecond): + } + return p, p.Output(), nil +} + +// StartFrpc starts frpc without an implicit sleep so tests can wait on an +// explicit login or proxy-readiness event. +func (f *Framework) StartFrpc(args ...string) (*process.Process, string, error) { p := process.NewWithEnvs(TestContext.FRPClientPath, args, f.osEnvs) f.clientProcesses = append(f.clientProcesses, p) err := p.Start() if err != nil { return p, p.Output(), err } - select { - case <-p.Done(): - case <-time.After(1500 * time.Millisecond): - } return p, p.Output(), nil } diff --git a/test/e2e/pkg/relay/halfopen.go b/test/e2e/pkg/relay/halfopen.go new file mode 100644 index 00000000..34389e53 --- /dev/null +++ b/test/e2e/pkg/relay/halfopen.go @@ -0,0 +1,208 @@ +// Copyright 2026 The frp Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package relay + +import ( + "fmt" + "io" + "net" + "strconv" + "sync" + "time" +) + +// HalfOpen forwards TCP connections until Blackhole is called. A blackholed +// pair stops forwarding but deliberately retains the upstream socket so the +// peer sees a real half-open connection until the relay is closed. +type HalfOpen struct { + bindAddr string + bindPort int + upstreamAddr string + + listener net.Listener + done chan struct{} + accepted chan struct{} + + mu sync.Mutex + pairs []*connectionPair + + wg sync.WaitGroup + closeOnce sync.Once +} + +type connectionPair struct { + downstream net.Conn + upstream net.Conn + + mu sync.Mutex + blackholed bool +} + +func New(upstreamAddr string) *HalfOpen { + return &HalfOpen{ + bindAddr: "127.0.0.1", + upstreamAddr: upstreamAddr, + done: make(chan struct{}), + accepted: make(chan struct{}, 1), + } +} + +func (r *HalfOpen) Run() error { + listener, err := net.Listen("tcp", net.JoinHostPort(r.bindAddr, strconv.Itoa(r.bindPort))) + if err != nil { + return err + } + r.listener = listener + r.bindPort = listener.Addr().(*net.TCPAddr).Port + + r.wg.Add(1) + go r.acceptLoop() + return nil +} + +func (r *HalfOpen) acceptLoop() { + defer r.wg.Done() + for { + downstream, err := r.listener.Accept() + if err != nil { + return + } + upstream, err := net.DialTimeout("tcp", r.upstreamAddr, 3*time.Second) + if err != nil { + _ = downstream.Close() + continue + } + + pair := &connectionPair{downstream: downstream, upstream: upstream} + r.mu.Lock() + r.pairs = append(r.pairs, pair) + r.mu.Unlock() + select { + case r.accepted <- struct{}{}: + default: + } + + r.wg.Add(1) + go r.servePair(pair) + } +} + +func (r *HalfOpen) servePair(pair *connectionPair) { + defer r.wg.Done() + copyDone := make(chan struct{}, 2) + go func() { + _, _ = io.Copy(pair.upstream, pair.downstream) + copyDone <- struct{}{} + }() + go func() { + _, _ = io.Copy(pair.downstream, pair.upstream) + copyDone <- struct{}{} + }() + + completed := 0 + select { + case <-copyDone: + completed = 1 + if !pair.isBlackholed() { + _ = pair.downstream.Close() + _ = pair.upstream.Close() + } + case <-r.done: + _ = pair.downstream.Close() + _ = pair.upstream.Close() + } + for completed < 2 { + <-copyDone + completed++ + } + + if pair.isBlackholed() { + <-r.done + _ = pair.downstream.Close() + _ = pair.upstream.Close() + } +} + +func (r *HalfOpen) WaitForConnections(count int, timeout time.Duration) error { + timer := time.NewTimer(timeout) + defer timer.Stop() + for { + r.mu.Lock() + accepted := len(r.pairs) + r.mu.Unlock() + if accepted >= count { + return nil + } + select { + case <-r.accepted: + case <-r.done: + return fmt.Errorf("relay closed after accepting %d of %d connections", accepted, count) + case <-timer.C: + return fmt.Errorf("timed out after accepting %d of %d connections", accepted, count) + } + } +} + +// Blackhole uses a one-based connection index in accept order. +func (r *HalfOpen) Blackhole(index int) error { + r.mu.Lock() + if index <= 0 || index > len(r.pairs) { + accepted := len(r.pairs) + r.mu.Unlock() + return fmt.Errorf("connection %d is unavailable; accepted %d", index, accepted) + } + pair := r.pairs[index-1] + r.mu.Unlock() + + pair.mu.Lock() + if pair.blackholed { + pair.mu.Unlock() + return nil + } + pair.blackholed = true + pair.mu.Unlock() + + now := time.Now() + _ = pair.downstream.SetDeadline(now) + _ = pair.upstream.SetDeadline(now) + return nil +} + +func (r *HalfOpen) Close() error { + r.closeOnce.Do(func() { + close(r.done) + if r.listener != nil { + _ = r.listener.Close() + } + r.mu.Lock() + pairs := append([]*connectionPair(nil), r.pairs...) + r.mu.Unlock() + for _, pair := range pairs { + _ = pair.downstream.Close() + _ = pair.upstream.Close() + } + r.wg.Wait() + }) + return nil +} + +func (r *HalfOpen) BindAddr() string { return r.bindAddr } +func (r *HalfOpen) BindPort() int { return r.bindPort } + +func (p *connectionPair) isBlackholed() bool { + p.mu.Lock() + defer p.mu.Unlock() + return p.blackholed +} diff --git a/test/e2e/v1/features/control_replacement.go b/test/e2e/v1/features/control_replacement.go new file mode 100644 index 00000000..dbb95f41 --- /dev/null +++ b/test/e2e/v1/features/control_replacement.go @@ -0,0 +1,300 @@ +// Copyright 2026 The frp Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package features + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "strconv" + "strings" + "time" + + "github.com/onsi/ginkgo/v2" + + "github.com/fatedier/frp/test/e2e/framework" + "github.com/fatedier/frp/test/e2e/framework/consts" + "github.com/fatedier/frp/test/e2e/pkg/relay" + "github.com/fatedier/frp/test/e2e/pkg/request" +) + +var _ = ginkgo.Describe("[Feature: ControlReplacement]", func() { + f := framework.NewDefaultFramework() + + for _, wireProtocol := range []string{"v1", "v2"} { + for _, tcpMux := range []bool{true, false} { + ginkgo.It(fmt.Sprintf("recovers a %s control through a half-open relay with tcpMux=%t", wireProtocol, tcpMux), func() { + runHalfOpenControlReplacement(f, wireProtocol, tcpMux) + }) + } + } +}) + +func runHalfOpenControlReplacement(f *framework.Framework, wireProtocol string, tcpMux bool) { + serverPort := f.AllocPort() + dashboardPort := f.AllocPort() + remotePort := f.AllocPort() + heartbeatTimeout := int64(-1) + if !tcpMux { + heartbeatTimeout = 3 + } + + serverConfig := fmt.Sprintf(` +bindAddr = "127.0.0.1" +bindPort = %d +log.level = "trace" +transport.tcpMux = %t +transport.tcpMuxKeepaliveInterval = 30 +transport.heartbeatTimeout = %d +webServer.addr = "127.0.0.1" +webServer.port = %d +webServer.pprofEnable = true +enablePrometheus = true +`, serverPort, tcpMux, heartbeatTimeout, dashboardPort) + serverConfigPath := f.WriteTempFile("issue-5391-frps.toml", serverConfig) + serverProcess, _, err := f.StartFrps("-c", serverConfigPath) + framework.ExpectNoError(err) + framework.ExpectNoError(framework.WaitForTCPReady(fmt.Sprintf("127.0.0.1:%d", serverPort), 5*time.Second)) + + halfOpenRelay := relay.New(fmt.Sprintf("127.0.0.1:%d", serverPort)) + f.RunServer("", halfOpenRelay) + + heartbeatInterval := int64(-1) + clientHeartbeatTimeout := int64(-1) + if !tcpMux { + heartbeatInterval = 1 + clientHeartbeatTimeout = 3 + } + clientConfig := fmt.Sprintf(` +serverAddr = "127.0.0.1" +serverPort = %d +clientID = "issue-5391" +loginFailExit = false +log.level = "trace" +transport.wireProtocol = %q +transport.tcpMux = %t +transport.tcpMuxKeepaliveInterval = 1 +transport.heartbeatInterval = %d +transport.heartbeatTimeout = %d +transport.tls.enable = false + +[[proxies]] +name = "issue-5391-tcp" +type = "tcp" +localPort = %d +remotePort = %d +`, halfOpenRelay.BindPort(), wireProtocol, tcpMux, heartbeatInterval, clientHeartbeatTimeout, + f.PortByName(framework.TCPEchoServerPort), remotePort) + clientConfigPath := f.WriteTempFile("issue-5391-frpc.toml", clientConfig) + clientProcess, _, err := f.StartFrpc("-c", clientConfigPath) + framework.ExpectNoError(err) + framework.ExpectNoError(halfOpenRelay.WaitForConnections(1, 5*time.Second)) + framework.ExpectNoError(clientProcess.WaitForOutput("login to server success", 1, 10*time.Second)) + framework.ExpectNoError(clientProcess.WaitForOutput("[issue-5391-tcp] start proxy success", 1, 10*time.Second)) + framework.ExpectNoError(waitForReplacementState(dashboardPort, remotePort, 10*time.Second)) + + replacementCount := 1 + if tcpMux { + replacementCount = 3 + } + for i := 0; i < replacementCount; i++ { + connectionIndex := 1 + if tcpMux { + connectionIndex = i + 1 + } + framework.ExpectNoError(halfOpenRelay.Blackhole(connectionIndex)) + if tcpMux { + framework.ExpectNoError(halfOpenRelay.WaitForConnections(connectionIndex+1, 15*time.Second)) + } + framework.ExpectNoError(clientProcess.WaitForOutput("login to server success", i+2, 15*time.Second)) + framework.ExpectNoError(clientProcess.WaitForOutput("[issue-5391-tcp] start proxy success", i+2, 15*time.Second)) + framework.ExpectNoError(waitForReplacementState(dashboardPort, remotePort, 10*time.Second)) + } + + _ = clientProcess.Stop() + select { + case <-clientProcess.Done(): + case <-time.After(5 * time.Second): + framework.Failf("frpc did not exit") + } + framework.ExpectNoError(waitForReplacementShutdown(dashboardPort, 10*time.Second)) + framework.ExpectNoError(halfOpenRelay.Close()) + framework.ExpectNoError(waitForNoHandoffWaiters(dashboardPort, 5*time.Second)) + + _ = serverProcess.Stop() + select { + case <-serverProcess.Done(): + case <-time.After(5 * time.Second): + framework.Failf("frps did not exit") + } +} + +func waitForReplacementState(dashboardPort, remotePort int, timeout time.Duration) error { + return waitForLifecycleCondition(timeout, func() error { + metricsBody, err := getLifecycleEndpoint(dashboardPort, "/metrics") + if err != nil { + return err + } + if err := expectMetricValue(metricsBody, "frp_server_client_counts", "", 1); err != nil { + return err + } + if err := expectMetricValue(metricsBody, "frp_server_proxy_counts", `type="tcp"`, 1); err != nil { + return err + } + clients, err := getOnlineLifecycleClients(dashboardPort) + if err != nil { + return err + } + if len(clients) != 1 || clients[0].ClientID != "issue-5391" { + return fmt.Errorf("expected one online client, got %+v", clients) + } + profile, err := getLifecycleEndpoint(dashboardPort, "/debug/pprof/goroutine?debug=2") + if err != nil { + return err + } + if err := expectNoHandoffWaiter(profile, "after replacement"); err != nil { + return err + } + resp, err := request.New(). + TCP(). + Port(remotePort). + Timeout(time.Second). + Body([]byte(consts.TestString)). + Do() + if err != nil { + return err + } + if string(resp.Content) != consts.TestString { + return fmt.Errorf("unexpected proxy response %q", resp.Content) + } + return nil + }) +} + +func waitForReplacementShutdown(dashboardPort int, timeout time.Duration) error { + return waitForLifecycleCondition(timeout, func() error { + metricsBody, err := getLifecycleEndpoint(dashboardPort, "/metrics") + if err != nil { + return err + } + if err := expectMetricValue(metricsBody, "frp_server_client_counts", "", 0); err != nil { + return err + } + if err := expectMetricValue(metricsBody, "frp_server_proxy_counts", `type="tcp"`, 0); err != nil { + return err + } + clients, err := getOnlineLifecycleClients(dashboardPort) + if err != nil { + return err + } + if len(clients) != 0 { + return fmt.Errorf("expected no online clients, got %+v", clients) + } + return nil + }) +} + +func waitForNoHandoffWaiters(dashboardPort int, timeout time.Duration) error { + return waitForLifecycleCondition(timeout, func() error { + profile, err := getLifecycleEndpoint(dashboardPort, "/debug/pprof/goroutine?debug=2") + if err != nil { + return err + } + return expectNoHandoffWaiter(profile, "after relay shutdown") + }) +} + +type lifecycleClient struct { + ClientID string `json:"clientID"` +} + +func expectNoHandoffWaiter(profile, phase string) error { + if strings.Contains(profile, "(*Control).WaitForHandoff") { + return fmt.Errorf("control handoff waiter remained %s", phase) + } + return nil +} + +func getOnlineLifecycleClients(dashboardPort int) ([]lifecycleClient, error) { + body, err := getLifecycleEndpoint(dashboardPort, "/api/clients?status=online") + if err != nil { + return nil, err + } + var clients []lifecycleClient + if err := json.Unmarshal([]byte(body), &clients); err != nil { + return nil, err + } + return clients, nil +} + +func getLifecycleEndpoint(port int, path string) (string, error) { + client := &http.Client{Timeout: time.Second} + resp, err := client.Get(fmt.Sprintf("http://127.0.0.1:%d%s", port, path)) + if err != nil { + return "", err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("GET %s returned %s", path, resp.Status) + } + body, err := io.ReadAll(resp.Body) + if err != nil { + return "", err + } + return string(body), nil +} + +func expectMetricValue(body, name, labels string, want float64) error { + prefix := name + if labels != "" { + prefix += "{" + labels + "}" + } + for line := range strings.SplitSeq(body, "\n") { + fields := strings.Fields(line) + if len(fields) != 2 || fields[0] != prefix { + continue + } + got, err := strconv.ParseFloat(fields[1], 64) + if err != nil { + return err + } + if got != want { + return fmt.Errorf("metric %s = %v, want %v", prefix, got, want) + } + return nil + } + return fmt.Errorf("metric %s not found", prefix) +} + +func waitForLifecycleCondition(timeout time.Duration, condition func() error) error { + timer := time.NewTimer(timeout) + defer timer.Stop() + ticker := time.NewTicker(25 * time.Millisecond) + defer ticker.Stop() + var lastErr error + for { + err := condition() + if err == nil { + return nil + } + lastErr = err + select { + case <-ticker.C: + case <-timer.C: + return fmt.Errorf("condition was not met: %w", lastErr) + } + } +}