From f32adb3bbb72809ba43d6f4c9d85ee1ba87803c0 Mon Sep 17 00:00:00 2001 From: fatedier Date: Tue, 3 Nov 2026 23:31:44 +0800 Subject: [PATCH] fix(server): close the mux session when a control is replaced or exits (#5596) When a client logs in again with the same run ID, frps replaces the old control and closes its control stream, but the yamux session that carried it stays open until keepalive reclaims it. The replaced control's work connections and streams share that session, so they are left behind for up to tcpMuxKeepaliveInterval plus the ping budget. A mux session carries exactly one control connection, so the control now records its session and closes it from the single teardown point shared by replacement, fatal exit and shutdown. QUIC and non-multiplexed transports keep the previous path. Closing the session ends every stream that still belongs to the replaced control. Those proxies are already being torn down, so the cost is that in-flight connections of the old control are interrupted immediately instead of at the keepalive deadline. A replacement control is not affected, and the ownership rules for a still-online control do not change. Tests: replacement and fatal exit close the session while the replacement control keeps working, and concurrent logins with the same run ID leave exactly one live session. --- server/control.go | 8 ++ server/control_session_test.go | 186 +++++++++++++++++++++++++++++++++ server/service.go | 12 ++- server/service_test.go | 14 +-- 4 files changed, 208 insertions(+), 12 deletions(-) create mode 100644 server/control_session_test.go diff --git a/server/control.go b/server/control.go index 8009d8ac..35ded591 100644 --- a/server/control.go +++ b/server/control.go @@ -17,6 +17,7 @@ package server import ( "context" "fmt" + "io" "math" "net" "runtime/debug" @@ -366,6 +367,9 @@ type SessionContext struct { EncryptionKey []byte // control connection Conn *msg.Conn + // A mux session carries exactly one control. Close it when the control is + // replaced or exits fatally to release all remaining streams and work connections. + Session io.Closer // login message LoginMsg *msg.Login // server configuration @@ -575,6 +579,10 @@ func (ctl *Control) markReplaced() <-chan struct{} { func (ctl *Control) interruptReadAndClose() error { ctl.interruptOnce.Do(func() { _ = ctl.sessionCtx.Conn.SetReadDeadline(time.Now()) + if ctl.sessionCtx.Session != nil { + ctl.interruptErr = ctl.sessionCtx.Session.Close() + return + } ctl.interruptErr = ctl.sessionCtx.Conn.Close() }) return ctl.interruptErr diff --git a/server/control_session_test.go b/server/control_session_test.go new file mode 100644 index 00000000..4ab76fd1 --- /dev/null +++ b/server/control_session_test.go @@ -0,0 +1,186 @@ +// Copyright 2026 fatedier, fatedier@gmail.com +// +// 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" + "net" + "testing" + "time" + + fmux "github.com/fatedier/yamux" + "github.com/stretchr/testify/require" + + "github.com/fatedier/frp/pkg/msg" + "github.com/fatedier/frp/pkg/proto/wire" + "github.com/fatedier/frp/pkg/util/util" +) + +func TestControlReplacementClosesYamuxSession(t *testing.T) { + svr := newControlTestService(t) + _, clientA, serverA, _ := newYamuxTestControl(t, svr) + _, _, _, connB := newYamuxTestControl(t, svr) + + waitForSignal(t, serverA.CloseChan(), "replaced control's server yamux session to close") + waitForSignal(t, clientA.CloseChan(), "replaced control's client yamux session to close") + _, err := clientA.OpenStream() + require.ErrorIs(t, err, fmux.ErrSessionShutdown) + + require.NoError(t, connB.SetDeadline(time.Now().Add(3*time.Second))) + require.NoError(t, connB.WriteMsg(&msg.Ping{})) + response, err := connB.ReadMsg() + require.NoError(t, err, "replacement control must still receive a Pong") + require.Equal(t, &msg.Pong{}, response) +} + +func TestControlCloseClosesYamuxSession(t *testing.T) { + ctl, client, server, _ := newYamuxTestControl(t, newControlTestService(t)) + + require.NoError(t, ctl.Close()) + waitForSignal(t, server.CloseChan(), "closed control's server yamux session to close") + waitForSignal(t, client.CloseChan(), "closed control's client yamux session to close") + _, err := client.OpenStream() + require.ErrorIs(t, err, fmux.ErrSessionShutdown) +} + +func TestControlConcurrentReplacementClosesLosingSessions(t *testing.T) { + svr := newControlTestService(t) + ctlA, clientA, serverA, connA := newYamuxTestControl(t, svr) + type result struct { + ctl *Control + client, server *fmux.Session + conn *msg.Conn + err error + } + start := make(chan struct{}) + results := make(chan result, 2) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + for range 2 { + client, server, clientConn, serverConn := newYamuxTestSession(t) + go func() { + r := result{client: client, server: server, conn: clientConn} + select { + case <-start: + r.ctl, r.err = svr.RegisterControl(serverConn, &msg.Login{ + RunID: "shared-run", + ClientID: "client", + Timestamp: 1, + PrivilegeKey: util.GetAuthKey("", 1), + }, false, wire.ProtocolV1, "", server) + if r.err == nil { + r.err = svr.completeControlLogin(r.ctl, func() error { return nil }) + } + if r.err != nil && r.ctl != nil { + svr.ctlManager.Remove(r.ctl) + _ = r.ctl.Close() + } + case <-ctx.Done(): + r.err = ctx.Err() + } + results <- r + }() + } + close(start) + controls := [3]result{{ctl: ctlA, client: clientA, server: serverA, conn: connA}} + for i := range 2 { + r := waitForResult(t, results, "concurrent control login") + if r.ctl != nil { + t.Cleanup(func() { + _ = r.ctl.Close() + waitForControlDone(t, r.ctl) + }) + } + if r.err != nil { + require.ErrorIs(t, r.err, errControlReplaced) + } + controls[i+1] = r + } + + current, ok := svr.ctlManager.GetByID("shared-run") + require.True(t, ok, "a running current control must remain") + require.NotSame(t, ctlA, current) + svr.ctlManager.mu.RLock() + currentCount := len(svr.ctlManager.ctlsByRunID) + svr.ctlManager.mu.RUnlock() + require.Equal(t, 1, currentCount) + winners := 0 + for _, r := range controls { + if r.ctl == current { + winners++ + require.NoError(t, r.err) + require.False(t, r.server.IsClosed(), "winner's server session must stay open") + require.False(t, r.client.IsClosed(), "winner's client session must stay open") + require.NoError(t, r.conn.SetDeadline(time.Now().Add(3*time.Second))) + require.NoError(t, r.conn.WriteMsg(&msg.Ping{})) + response, err := r.conn.ReadMsg() + require.NoError(t, err, "winner must still receive a Pong") + require.Equal(t, &msg.Pong{}, response) + continue + } + waitForSignal(t, r.server.CloseChan(), "losing control's server yamux session to close") + waitForSignal(t, r.client.CloseChan(), "losing control's client yamux session to close") + _, err := r.client.OpenStream() + require.ErrorIs(t, err, fmux.ErrSessionShutdown) + } + require.Equal(t, 1, winners) +} + +func newYamuxTestSession(t *testing.T) (*fmux.Session, *fmux.Session, *msg.Conn, *msg.Conn) { + t.Helper() + listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1)}) + require.NoError(t, err) + t.Cleanup(func() { _ = listener.Close() }) + require.NoError(t, listener.SetDeadline(time.Now().Add(3*time.Second))) + clientTCP, err := net.DialTimeout("tcp", listener.Addr().String(), 3*time.Second) + require.NoError(t, err) + t.Cleanup(func() { _ = clientTCP.Close() }) + serverTCP, err := listener.Accept() + require.NoError(t, err) + t.Cleanup(func() { _ = serverTCP.Close() }) + + client, err := fmux.Client(clientTCP, fmux.DefaultConfig()) + require.NoError(t, err) + t.Cleanup(func() { _ = client.Close() }) + server, err := fmux.Server(serverTCP, fmux.DefaultConfig()) + require.NoError(t, err) + t.Cleanup(func() { _ = server.Close() }) + clientStream, err := client.OpenStream() + require.NoError(t, err) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + serverStream, err := server.AcceptStreamWithContext(ctx) + require.NoError(t, err) + return client, server, msg.NewConn(clientStream, msg.NewV1ReadWriter(clientStream)), + msg.NewConn(serverStream, msg.NewV1ReadWriter(serverStream)) +} + +func newYamuxTestControl(t *testing.T, svr *Service) (*Control, *fmux.Session, *fmux.Session, *msg.Conn) { + t.Helper() + client, server, clientConn, serverConn := newYamuxTestSession(t) + ctl, err := svr.RegisterControl(serverConn, &msg.Login{ + RunID: "shared-run", + ClientID: "client", + Timestamp: 1, + PrivilegeKey: util.GetAuthKey("", 1), + }, false, wire.ProtocolV1, "", server) + require.NoError(t, err) + t.Cleanup(func() { + _ = ctl.Close() + waitForControlDone(t, ctl) + }) + require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil })) + return ctl, client, server, clientConn +} diff --git a/server/service.go b/server/service.go index 3bc20340..203bef14 100644 --- a/server/service.go +++ b/server/service.go @@ -521,7 +521,7 @@ func (svr *Service) closeResources() { svr.ctlManager.Close() } -func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, internal bool) { +func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, internal bool, session io.Closer) { xl := xlog.FromContextSafe(ctx) acceptedConn, err := svr.acceptConnection(ctx, conn) @@ -552,7 +552,7 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, interna } } if err == nil { - ctl, err = svr.RegisterControl(controlConn, m, internal, acceptedConn.wireProtocol, acceptedConn.udpPacketCodec) + ctl, err = svr.RegisterControl(controlConn, m, internal, acceptedConn.wireProtocol, acceptedConn.udpPacketCodec, session) } } @@ -813,10 +813,10 @@ func (svr *Service) HandleListener(l net.Listener, internal bool) { session.Close() return } - go svr.handleConnection(ctx, stream, internal) + go svr.handleConnection(ctx, stream, internal, session) } } else { - svr.handleConnection(ctx, frpConn, internal) + svr.handleConnection(ctx, frpConn, internal, nil) } }(ctx, c) } @@ -839,7 +839,7 @@ func (svr *Service) HandleQUICListener(l *quic.Listener) { _ = frpConn.CloseWithError(0, "") return } - go svr.handleConnection(ctx, netpkg.QuicStreamToNetConn(stream, frpConn), false) + go svr.handleConnection(ctx, netpkg.QuicStreamToNetConn(stream, frpConn), false, nil) } }(context.Background(), c) } @@ -851,6 +851,7 @@ func (svr *Service) RegisterControl( internal bool, wireProtocol string, udpPacketCodec string, + session io.Closer, ) (*Control, error) { switch wireProtocol { case wire.ProtocolV1: @@ -900,6 +901,7 @@ func (svr *Service) RegisterControl( AuthVerifier: authVerifier, EncryptionKey: svr.auth.EncryptionKey(), Conn: ctlConn, + Session: session, LoginMsg: loginMsg, ServerCfg: svr.cfg, WireProtocol: wireProtocol, diff --git a/server/service_test.go b/server/service_test.go index d4fcbe77..8bc055e2 100644 --- a/server/service_test.go +++ b/server/service_test.go @@ -576,7 +576,7 @@ func TestServiceWorkConnRoutingClientHelloPolicy(t *testing.T) { ClientSpec: msg.ClientSpec{ AlwaysAuthPass: true, }, - }, true, wire.ProtocolV2, tc.controlUDPPacketCodec) + }, true, wire.ProtocolV2, tc.controlUDPPacketCodec, nil) require.NoError(t, err) require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil })) waitForSignal(t, controlConn.readStarted, "control reader to start") @@ -634,7 +634,7 @@ func TestServiceRegisterControlRejectsInvalidCodecSelection(t *testing.T) { svr := newControlTestService(t) conn := newDeadlineReadConn() msgConn := msg.NewConn(conn, msg.NewV1ReadWriter(conn)) - ctl, err := svr.RegisterControl(msgConn, &msg.Login{}, true, tc.wireProtocol, tc.udpPacketCodec) + ctl, err := svr.RegisterControl(msgConn, &msg.Login{}, true, tc.wireProtocol, tc.udpPacketCodec, nil) require.Nil(t, ctl) require.ErrorContains(t, err, tc.errorSubstring) }) @@ -650,7 +650,7 @@ func TestServiceRegisterControlRejectsInvalidRunID(t *testing.T) { svr := newControlTestService(t) conn := newDeadlineReadConn() msgConn := msg.NewConn(conn, msg.NewV1ReadWriter(conn)) - ctl, err := svr.RegisterControl(msgConn, &msg.Login{RunID: runID}, true, wire.ProtocolV1, "") + ctl, err := svr.RegisterControl(msgConn, &msg.Login{RunID: runID}, true, wire.ProtocolV1, "", nil) require.Nil(t, ctl) require.ErrorContains(t, err, "invalid run id") }) @@ -684,7 +684,7 @@ func TestServiceRegisterControlPoolCountBoundaries(t *testing.T) { Timestamp: timestamp, PrivilegeKey: util.GetAuthKey("", timestamp), PoolCount: tc.poolCount, - }, false, wire.ProtocolV1, "") + }, false, wire.ProtocolV1, "", nil) if tc.wantErr { require.Nil(t, ctl) require.ErrorContains(t, err, "unexpected error when creating new controller") @@ -758,7 +758,7 @@ func TestServiceVisitorRoutingExcludesPendingUser(t *testing.T) { ClientSpec: msg.ClientSpec{ AlwaysAuthPass: true, }, - }, true, wire.ProtocolV1, "") + }, true, wire.ProtocolV1, "", nil) require.NoError(t, err) timestamp := time.Now().Unix() @@ -802,7 +802,7 @@ func TestServiceVisitorRoutingCarriesControlPacketCodec(t *testing.T) { ClientSpec: msg.ClientSpec{ AlwaysAuthPass: true, }, - }, true, wire.ProtocolV2, wire.UDPPacketCodecBinary) + }, true, wire.ProtocolV2, wire.UDPPacketCodecBinary, nil) require.NoError(t, err) timestamp := time.Now().Unix() @@ -891,7 +891,7 @@ func registerLifecycleTestControl(svr *Service) (*Control, *deadlineReadConn, er ClientSpec: msg.ClientSpec{ AlwaysAuthPass: true, }, - }, true, wire.ProtocolV1, "") + }, true, wire.ProtocolV1, "", nil) return ctl, conn, err }