diff --git a/client/proxy/proxy_wrapper.go b/client/proxy/proxy_wrapper.go index 37a71277..30ea5e33 100644 --- a/client/proxy/proxy_wrapper.go +++ b/client/proxy/proxy_wrapper.go @@ -267,8 +267,9 @@ func (pw *Wrapper) InWorkConn(workConn net.Conn, m *msg.StartWorkConn) { xl := pw.xl pw.mu.RLock() pxy := pw.pxy + running := pw.Phase == ProxyPhaseRunning pw.mu.RUnlock() - if pxy != nil && pw.Phase == ProxyPhaseRunning { + if pxy != nil && running { xl.Debugf("start a new work connection, localAddr: %s remoteAddr: %s", workConn.LocalAddr().String(), workConn.RemoteAddr().String()) go pxy.InWorkConn(workConn, m) } else { diff --git a/client/proxy/proxy_wrapper_test.go b/client/proxy/proxy_wrapper_test.go new file mode 100644 index 00000000..639f3cf1 --- /dev/null +++ b/client/proxy/proxy_wrapper_test.go @@ -0,0 +1,76 @@ +//go:build !frps + +package proxy + +import ( + "io" + "net" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/fatedier/frp/pkg/msg" + "github.com/fatedier/frp/pkg/util/xlog" +) + +type closingWorkProxy struct { + Proxy + accepted atomic.Int64 +} + +func (p *closingWorkProxy) InWorkConn(conn net.Conn, _ *msg.StartWorkConn) { + p.accepted.Add(1) + _ = conn.Close() +} + +func assertWrapperWorkConnClosed(t *testing.T, pw *Wrapper) { + t.Helper() + conn, peer := net.Pipe() + defer conn.Close() + defer peer.Close() + require.NoError(t, peer.SetReadDeadline(time.Now().Add(time.Second))) + pw.InWorkConn(conn, &msg.StartWorkConn{}) + _, err := peer.Read(make([]byte, 1)) + require.ErrorIs(t, err, io.EOF) +} + +func TestWrapperInWorkConnPhaseSnapshot(t *testing.T) { + pxy := &closingWorkProxy{} + pw := &Wrapper{pxy: pxy, xl: xlog.New()} + for _, phase := range []string{ProxyPhaseNew, ProxyPhaseWaitStart, ProxyPhaseStartErr, ProxyPhaseCheckFailed, ProxyPhaseClosed} { + pw.Phase = phase + assertWrapperWorkConnClosed(t, pw) + } + require.Zero(t, pxy.accepted.Load(), "inactive proxy received a work connection") + pw.Phase = ProxyPhaseRunning + pw.pxy = nil + assertWrapperWorkConnClosed(t, pw) + pw.pxy = pxy + assertWrapperWorkConnClosed(t, pw) + require.Equal(t, int64(1), pxy.accepted.Load()) + + start, done := make(chan struct{}), make(chan struct{}) + go func() { + <-start + for range 20000 { + pw.mu.Lock() + pw.Phase = ProxyPhaseRunning + pw.mu.Unlock() + pw.mu.Lock() + pw.Phase = ProxyPhaseClosed + pw.mu.Unlock() + } + close(done) + }() + close(start) + for range 10000 { + assertWrapperWorkConnClosed(t, pw) + } + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("phase writer did not finish") + } +} diff --git a/client/service.go b/client/service.go index c18e94dd..055de40d 100644 --- a/client/service.go +++ b/client/service.go @@ -155,6 +155,9 @@ type Service struct { ctx context.Context // call cancel to stop service cancel context.CancelCauseFunc + // Protect cancel publication and shutdown before the first Run. + lifecycleMu sync.Mutex + closed bool connectorCreator func(context.Context, *v1.ClientCommonConfig) Connector handleWorkConnCb func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool @@ -225,8 +228,20 @@ func NewService(options ServiceOptions) (*Service, error) { func (svr *Service) Run(ctx context.Context) error { ctx, cancel := context.WithCancelCause(ctx) + svr.lifecycleMu.Lock() + if svr.closed { + svr.lifecycleMu.Unlock() + cancel(nil) + return ctx.Err() + } svr.ctx = xlog.NewContext(ctx, xlog.FromContextSafe(ctx)) svr.cancel = cancel + svr.lifecycleMu.Unlock() + defer cancel(nil) + defer svr.stop() + if err := ctx.Err(); err != nil { + return err + } // set custom DNSServer if svr.common.DNSServer != "" { @@ -237,7 +252,6 @@ func (svr *Service) Run(ctx context.Context) error { vnetController := svr.vnetController if err := svr.vnetController.Init(); err != nil { log.Errorf("init virtual network controller error: %v", err) - svr.stop() return err } go func() { @@ -263,14 +277,12 @@ func (svr *Service) Run(ctx context.Context) error { if svr.ctl == nil { cancelCause := cancelErr{} _ = errors.As(context.Cause(svr.ctx), &cancelCause) - svr.stop() return fmt.Errorf("login to the server failed: %v. With loginFailExit enabled, no additional retries will be attempted", cancelCause.Err) } go svr.keepControllerWorking() <-svr.ctx.Done() - svr.stop() return nil } @@ -428,7 +440,20 @@ func (svr *Service) Close() { func (svr *Service) GracefulClose(d time.Duration) { svr.gracefulShutdownDuration.Store(int64(d)) - svr.cancel(nil) + svr.lifecycleMu.Lock() + if svr.closed { + svr.lifecycleMu.Unlock() + return + } + svr.closed = true + cancel := svr.cancel + svr.lifecycleMu.Unlock() + if cancel != nil { + // Once Run has claimed startup, it also owns cleanup. + cancel(nil) + } else { + svr.stop() + } } func (svr *Service) stop() { diff --git a/client/service_lifecycle_test.go b/client/service_lifecycle_test.go new file mode 100644 index 00000000..13969506 --- /dev/null +++ b/client/service_lifecycle_test.go @@ -0,0 +1,186 @@ +package client + +import ( + "context" + "io" + "net" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/fatedier/frp/pkg/config/source" + v1 "github.com/fatedier/frp/pkg/config/v1" + "github.com/fatedier/frp/pkg/config/v1/validation" + "github.com/fatedier/frp/pkg/msg" +) + +func newLifecycleTestService(t *testing.T, creator func(context.Context, *v1.ClientCommonConfig) Connector) (*Service, string) { + t.Helper() + svr, err := NewService(ServiceOptions{ + Common: &v1.ClientCommonConfig{ + LoginFailExit: new(false), + WebServer: v1.WebServerConfig{Addr: "127.0.0.1", Port: getFreeTCPPort(t)}, + }, + ConfigSourceAggregator: source.NewAggregator(source.NewConfigSource()), + ConnectorCreator: creator, + }) + require.NoError(t, err) + _, err = validation.NewConfigValidator(nil).ValidateClientCommonConfig(svr.common) + require.NoError(t, err) + address := svr.webServer.Address() + assertLifecycleAdminPort(t, address, false) + return svr, address +} + +func assertLifecycleAdminPort(t *testing.T, address string, released bool) { + t.Helper() + ln, err := net.Listen("tcp", address) + if ln != nil { + _ = ln.Close() + } + if released { + require.NoError(t, err, "admin listener was not released") + } else { + require.Error(t, err, "admin listener was not bound") + } +} + +func waitLifecycleSignal(t *testing.T, done <-chan struct{}) { + t.Helper() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("service lifecycle operation did not finish") + } +} + +func TestServiceRunAfterEarlyShutdown(t *testing.T) { + for _, mode := range []string{"close before Run", "parent already cancelled"} { + t.Run(mode, func(t *testing.T) { + var attempts atomic.Int64 + svr, address := newLifecycleTestService(t, func(context.Context, *v1.ClientCommonConfig) Connector { + attempts.Add(1) + return &failingConnector{err: net.ErrClosed} + }) + // Detect startup's global DNS side effect without making DNS requests. + resolver := net.DefaultResolver + t.Cleanup(func() { net.DefaultResolver = resolver }) + svr.common.DNSServer = "127.0.0.1:53" + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if mode == "close before Run" { + svr.Close() + assertLifecycleAdminPort(t, address, true) + } else { + cancel() + } + done := make(chan struct{}) + var runErr error + go func() { + runErr = svr.Run(ctx) + close(done) + }() + waitLifecycleSignal(t, done) + require.ErrorIs(t, runErr, context.Canceled) + require.Zero(t, attempts.Load(), "closed service attempted login") + require.Same(t, resolver, net.DefaultResolver, "closed service changed DNS") + assertLifecycleAdminPort(t, address, true) + }) + } +} + +func TestServiceStartupShutdownLeavesCleanupToRun(t *testing.T) { + for _, mode := range []string{"Close", "parent cancellation"} { + t.Run(mode, func(t *testing.T) { + entered, resume := make(chan struct{}), make(chan struct{}) + svr, address := newLifecycleTestService(t, func(ctx context.Context, _ *v1.ClientCommonConfig) Connector { + close(entered) + <-resume + return &failingConnector{err: ctx.Err()} + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + defer close(resume) + runDone := make(chan struct{}) + go func() { _ = svr.Run(ctx); close(runDone) }() + waitLifecycleSignal(t, entered) + closeDone := make(chan struct{}) + go func() { + if mode == "Close" { + svr.Close() + } else { + cancel() + } + close(closeDone) + }() + waitLifecycleSignal(t, closeDone) + select { + case <-runDone: + t.Fatal("Run exited while startup still owned its resources") + default: + } + assertLifecycleAdminPort(t, address, false) + resume <- struct{}{} + waitLifecycleSignal(t, runDone) + assertLifecycleAdminPort(t, address, true) + require.ErrorIs(t, svr.ctx.Err(), context.Canceled) + }) + } +} + +func TestServiceShutdownAfterLogin(t *testing.T) { + for _, mode := range []string{"Close", "parent cancellation"} { + t.Run(mode, func(t *testing.T) { + clientConn, serverConn := net.Pipe() + defer clientConn.Close() + defer serverConn.Close() + connector := &testConnector{conn: clientConn} + svr, address := newLifecycleTestService(t, func(context.Context, *v1.ClientCommonConfig) Connector { + return connector + }) + defer svr.Close() + peerDone := make(chan error, 1) + go func() { + rw := msg.NewV1ReadWriter(serverConn) + var login msg.Login + if err := rw.ReadMsgInto(&login); err != nil { + peerDone <- err + return + } + if err := rw.WriteMsg(&msg.LoginResp{RunID: "lifecycle-test"}); err != nil { + peerDone <- err + return + } + _, err := io.Copy(io.Discard, serverConn) + peerDone <- err + }() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + runDone := make(chan struct{}) + var runErr error + go func() { runErr = svr.Run(ctx); close(runDone) }() + require.Eventually(t, func() bool { + svr.ctlMu.RLock() + defer svr.ctlMu.RUnlock() + return svr.ctl != nil + }, 3*time.Second, time.Millisecond, "login did not establish control") + if mode == "Close" { + svr.Close() + } else { + cancel() + } + waitLifecycleSignal(t, runDone) + require.NoError(t, runErr) + require.True(t, connector.closed.Load()) + select { + case err := <-peerDone: + require.NoError(t, err) + case <-time.After(3 * time.Second): + t.Fatal("control connection was not released") + } + assertLifecycleAdminPort(t, address, true) + }) + } +} diff --git a/server/service.go b/server/service.go index 6ea24cc4..4556358b 100644 --- a/server/service.go +++ b/server/service.go @@ -25,6 +25,7 @@ import ( "net/http" "os" "strconv" + "sync" "time" "github.com/fatedier/golib/crypto" @@ -120,6 +121,9 @@ type Service struct { // web server for dashboard UI and apis webServer *httppkg.Server + // HTTP vhost resources are started and owned by NewService. + httpVhostServer *http.Server + httpVhostListener net.Listener sshTunnelGateway *ssh.Gateway @@ -133,7 +137,10 @@ type Service struct { // service context ctx context.Context // call cancel to stop service - cancel context.CancelFunc + cancel context.CancelFunc + lifecycleMu sync.Mutex + closed bool + closeOnce sync.Once } func NewService(cfg *v1.ServerConfig) (*Service, error) { @@ -319,6 +326,8 @@ func NewService(cfg *v1.ServerConfig) (*Service, error) { return nil, fmt.Errorf("create vhost http listener error, %v", err) } } + svr.httpVhostServer = server + svr.httpVhostListener = l go func() { _ = server.Serve(l) }() @@ -365,8 +374,22 @@ func NewService(cfg *v1.ServerConfig) (*Service, error) { func (svr *Service) Run(ctx context.Context) { ctx, cancel := context.WithCancel(ctx) + svr.lifecycleMu.Lock() + if svr.closed { + svr.lifecycleMu.Unlock() + cancel() + return + } svr.ctx = ctx svr.cancel = cancel + svr.lifecycleMu.Unlock() + defer svr.Close() + if ctx.Err() != nil { + return + } + // Cancellation must unblock the synchronous main listener's Accept. + stopClose := context.AfterFunc(ctx, func() { _ = svr.Close() }) + defer stopClose() // run dashboard web server. if svr.webServer != nil { @@ -400,13 +423,21 @@ func (svr *Service) Run(ctx context.Context) { svr.HandleListener(svr.listener, false) <-svr.ctx.Done() - // service context may not be canceled by svr.Close(), we should call it here to release resources - if svr.listener != nil { - svr.Close() - } } func (svr *Service) Close() error { + svr.lifecycleMu.Lock() + svr.closed = true + cancel := svr.cancel + svr.lifecycleMu.Unlock() + if cancel != nil { + cancel() + } + svr.closeOnce.Do(svr.closeResources) + return nil +} + +func (svr *Service) closeResources() { if svr.kcpListener != nil { svr.kcpListener.Close() } @@ -428,16 +459,17 @@ func (svr *Service) Close() error { if svr.webServer != nil { svr.webServer.Close() } + if svr.httpVhostServer != nil { + svr.httpVhostServer.Close() + // Close even if Serve has not registered its listener yet. + svr.httpVhostListener.Close() + } if svr.sshTunnelGateway != nil { svr.sshTunnelGateway.Close() } svr.rc.Close() svr.muxer.Close() svr.ctlManager.Close() - if svr.cancel != nil { - svr.cancel() - } - return nil } func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, internal bool) { diff --git a/server/service_lifecycle_test.go b/server/service_lifecycle_test.go new file mode 100644 index 00000000..61481b32 --- /dev/null +++ b/server/service_lifecycle_test.go @@ -0,0 +1,151 @@ +package server + +import ( + "context" + "fmt" + "net" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + + v1 "github.com/fatedier/frp/pkg/config/v1" + "github.com/fatedier/frp/pkg/config/v1/validation" +) + +type lifecycleListener struct { + net.Listener + accepted, closeEntered, closeResume chan struct{} + acceptOnce, closeOnce sync.Once + accepts, closes atomic.Int64 +} + +func (l *lifecycleListener) Accept() (net.Conn, error) { + l.accepts.Add(1) + l.acceptOnce.Do(func() { close(l.accepted) }) + return l.Listener.Accept() +} + +func (l *lifecycleListener) Close() error { + l.closes.Add(1) + err := l.Listener.Close() + if l.closeEntered != nil { + l.closeOnce.Do(func() { close(l.closeEntered) }) + <-l.closeResume + } + return err +} + +func freeLifecycleTCPPort(t *testing.T) int { + t.Helper() + l, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer l.Close() + return l.Addr().(*net.TCPAddr).Port +} + +func newLifecycleTestService(t *testing.T, web bool) (*Service, *lifecycleListener, []string) { + t.Helper() + cfg := &v1.ServerConfig{ + BindAddr: "127.0.0.1", ProxyBindAddr: "127.0.0.1", + } + if web { + cfg.WebServer = v1.WebServerConfig{Addr: "127.0.0.1", Port: freeLifecycleTCPPort(t)} + cfg.VhostHTTPPort = freeLifecycleTCPPort(t) + } + require.NoError(t, cfg.Complete()) + cfg.BindPort = 0 + _, err := validation.NewConfigValidator(nil).ValidateServerConfig(cfg) + require.NoError(t, err) + svr, err := NewService(cfg) + require.NoError(t, err) + l := &lifecycleListener{Listener: svr.listener, accepted: make(chan struct{})} + svr.listener = l + addresses := []string{l.Addr().String()} + if web { + addresses = append(addresses, svr.webServer.Address()) + addresses = append(addresses, net.JoinHostPort(cfg.ProxyBindAddr, fmt.Sprint(cfg.VhostHTTPPort))) + } + assertLifecyclePorts(t, addresses, false) + return svr, l, addresses +} + +func assertLifecyclePorts(t *testing.T, addresses []string, released bool) { + t.Helper() + for _, address := range addresses { + ln, err := net.Listen("tcp", address) + if ln != nil { + _ = ln.Close() + } + if released { + require.NoError(t, err, "listener %s was not released", address) + } else { + require.Error(t, err, "listener %s was not bound", address) + } + } +} + +func TestServiceRunAfterEarlyShutdown(t *testing.T) { + for _, mode := range []string{"close before Run", "close before Run with web", "parent already cancelled"} { + t.Run(mode, func(t *testing.T) { + svr, listener, addresses := newLifecycleTestService(t, mode == "close before Run with web") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + defer svr.Close() + if mode != "parent already cancelled" { + require.NoError(t, svr.Close()) + assertLifecyclePorts(t, addresses, true) + } else { + cancel() + } + done := make(chan struct{}) + go func() { svr.Run(ctx); close(done) }() + waitForSignal(t, done, "Run after early shutdown") + require.Zero(t, listener.accepts.Load(), "closed service started accepting") + require.Equal(t, int64(1), listener.closes.Load()) + assertLifecyclePorts(t, addresses, true) + }) + } +} + +func TestServiceCloseMarksClosedBeforeCleanup(t *testing.T) { + svr, listener, addresses := newLifecycleTestService(t, false) + listener.closeEntered, listener.closeResume = make(chan struct{}), make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(listener.closeResume) }) } + defer release() + closeDone := make(chan struct{}) + go func() { _ = svr.Close(); close(closeDone) }() + waitForSignal(t, listener.closeEntered, "Close cleanup barrier") + runDone := make(chan struct{}) + go func() { svr.Run(context.Background()); close(runDone) }() + waitForSignal(t, runDone, "Run to observe closed before cleanup finishes") + require.Zero(t, listener.accepts.Load(), "Run started during cleanup") + assertLifecyclePorts(t, addresses, false) + release() + waitForSignal(t, closeDone, "Close cleanup to finish") + assertLifecyclePorts(t, addresses, true) +} + +func TestServiceParentCancellationAndCloseShareCleanup(t *testing.T) { + svr, listener, addresses := newLifecycleTestService(t, false) + listener.closeEntered, listener.closeResume = make(chan struct{}), make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(listener.closeResume) }) } + defer release() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + runDone := make(chan struct{}) + go func() { svr.Run(ctx); close(runDone) }() + waitForSignal(t, listener.accepted, "main listener Accept") + cancel() + waitForSignal(t, listener.closeEntered, "parent cancellation to close main listener") + closeDone := make(chan struct{}) + go func() { _ = svr.Close(); close(closeDone) }() + release() + waitForSignal(t, closeDone, "external Close to finish") + waitForSignal(t, runDone, "Run cleanup to finish") + require.Equal(t, int64(1), listener.closes.Load(), "cleanup ran more than once") + assertLifecyclePorts(t, addresses, true) +}