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 }