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

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

View File

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

View File

@@ -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)
}

543
server/control_test.go Normal file
View File

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

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.

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)
}
}

View File

@@ -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("")
}

View File

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

View File

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

View File

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

View File

@@ -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)
}
}
}