fix: synchronize service startup and shutdown (#5560)

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 {

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

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() {

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

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
 
@@ -134,6 +138,9 @@ type Service struct {
 	ctx context.Context
 	// call cancel to stop service
 	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) {

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