From 0b23d6f7d7ae2a01d31e56b2197d0dae8151e411 Mon Sep 17 00:00:00 2001 From: sagit <36596628+Sagit-chu@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:35:30 +0800 Subject: [PATCH] fix(agent): retire replaced tunnel sessions (#541) --- go-gost/x/chain/chain.go | 52 +++++++++++++-- go-gost/x/chain/chain_test.go | 64 +++++++++++++++++++ go-gost/x/chain/transport.go | 33 +++++++++- go-gost/x/chain/transport_test.go | 45 +++++++++++++ go-gost/x/config/parsing/chain/parse.go | 4 +- go-gost/x/dialer/kcp/conn.go | 7 +- go-gost/x/dialer/kcp/dialer.go | 38 +++++++++++ go-gost/x/dialer/kcp/dialer_test.go | 17 +++++ go-gost/x/dialer/mtcp/conn.go | 14 ++++ go-gost/x/dialer/mtcp/dialer.go | 40 ++++++++++++ go-gost/x/dialer/mtcp/dialer_test.go | 30 +++++++++ go-gost/x/dialer/mtls/conn.go | 14 ++++ go-gost/x/dialer/mtls/dialer.go | 40 ++++++++++++ go-gost/x/dialer/mtls/dialer_test.go | 30 +++++++++ go-gost/x/dialer/mws/conn.go | 14 ++++ go-gost/x/dialer/mws/dialer.go | 40 ++++++++++++ go-gost/x/dialer/mws/dialer_test.go | 30 +++++++++ go-gost/x/hop/hop.go | 57 ++++++++++++++--- .../util/sessionretire/sessionretire.go | 58 +++++++++++++++++ .../util/sessionretire/sessionretire_test.go | 59 +++++++++++++++++ go-gost/x/registry/chain.go | 8 ++- go-gost/x/registry/chain_test.go | 26 ++++++++ 22 files changed, 701 insertions(+), 19 deletions(-) create mode 100644 go-gost/x/chain/chain_test.go create mode 100644 go-gost/x/chain/transport_test.go create mode 100644 go-gost/x/dialer/kcp/dialer_test.go create mode 100644 go-gost/x/dialer/mtcp/dialer_test.go create mode 100644 go-gost/x/dialer/mtls/dialer_test.go create mode 100644 go-gost/x/dialer/mws/dialer_test.go create mode 100644 go-gost/x/internal/util/sessionretire/sessionretire.go create mode 100644 go-gost/x/internal/util/sessionretire/sessionretire_test.go diff --git a/go-gost/x/chain/chain.go b/go-gost/x/chain/chain.go index 458045a..b55bc48 100644 --- a/go-gost/x/chain/chain.go +++ b/go-gost/x/chain/chain.go @@ -2,6 +2,8 @@ package chain import ( "context" + "errors" + "io" "github.com/go-gost/core/chain" "github.com/go-gost/core/hop" @@ -38,11 +40,12 @@ type chainNamer interface { } type Chain struct { - name string - hops []hop.Hop - marker selector.Marker - metadata metadata.Metadata - logger logger.Logger + name string + hops []hop.Hop + ownedHops []hop.Hop + marker selector.Marker + metadata metadata.Metadata + logger logger.Logger } func NewChain(name string, opts ...ChainOption) *Chain { @@ -61,8 +64,15 @@ func NewChain(name string, opts ...ChainOption) *Chain { } } -func (c *Chain) AddHop(hop hop.Hop) { +func (c *Chain) AddHop(hop hop.Hop, owned ...bool) { c.hops = append(c.hops, hop) + isOwned := true + if len(owned) > 0 { + isOwned = owned[0] + } + if isOwned { + c.ownedHops = append(c.ownedHops, hop) + } } // Metadata implements metadata.Metadatable interface. @@ -112,6 +122,36 @@ func (c *Chain) Route(ctx context.Context, network, address string, opts ...chai return rt } +// Retire gracefully drains resources owned by a chain that has been replaced. +func (c *Chain) Retire() { + if c == nil { + return + } + for _, h := range c.ownedHops { + if retirer, ok := h.(interface{ Retire() }); ok { + retirer.Retire() + continue + } + if closer, ok := h.(io.Closer); ok { + _ = closer.Close() + } + } +} + +// Close immediately releases all resources owned by the chain. +func (c *Chain) Close() error { + if c == nil { + return nil + } + var errs []error + for _, h := range c.ownedHops { + if closer, ok := h.(io.Closer); ok { + errs = append(errs, closer.Close()) + } + } + return errors.Join(errs...) +} + type chainGroup struct { chains []chain.Chainer selector selector.Selector[chain.Chainer] diff --git a/go-gost/x/chain/chain_test.go b/go-gost/x/chain/chain_test.go new file mode 100644 index 0000000..5d7413f --- /dev/null +++ b/go-gost/x/chain/chain_test.go @@ -0,0 +1,64 @@ +package chain + +import ( + "context" + "testing" + + corechain "github.com/go-gost/core/chain" + corehop "github.com/go-gost/core/hop" +) + +type lifecycleTestHop struct { + selected int + retired int + closed int +} + +func (h *lifecycleTestHop) Select(context.Context, ...corehop.SelectOption) *corechain.Node { + h.selected++ + return nil +} + +func (h *lifecycleTestHop) Retire() { + h.retired++ +} + +func (h *lifecycleTestHop) Close() error { + h.closed++ + return nil +} + +func TestChainRoutesThroughSharedHopWithoutOwningLifecycle(t *testing.T) { + hop := &lifecycleTestHop{} + chain := NewChain("shared-hop") + chain.AddHop(hop, false) + + if route := chain.Route(context.Background(), "tcp", "example.com:443"); route == nil { + t.Fatal("route is nil") + } + if hop.selected != 1 { + t.Fatalf("shared hop selected %d times, want 1", hop.selected) + } + + chain.Retire() + if err := chain.Close(); err != nil { + t.Fatalf("close chain: %v", err) + } + if hop.retired != 0 || hop.closed != 0 { + t.Fatalf("shared hop lifecycle changed: retired=%d closed=%d", hop.retired, hop.closed) + } +} + +func TestChainRetiresAndClosesOwnedHop(t *testing.T) { + hop := &lifecycleTestHop{} + chain := NewChain("owned-hop") + chain.AddHop(hop) + + chain.Retire() + if err := chain.Close(); err != nil { + t.Fatalf("close chain: %v", err) + } + if hop.retired != 1 || hop.closed != 1 { + t.Fatalf("owned hop lifecycle: retired=%d closed=%d, want 1/1", hop.retired, hop.closed) + } +} diff --git a/go-gost/x/chain/transport.go b/go-gost/x/chain/transport.go index ed999f4..1976a3a 100644 --- a/go-gost/x/chain/transport.go +++ b/go-gost/x/chain/transport.go @@ -2,6 +2,8 @@ package chain import ( "context" + "errors" + "io" "net" "github.com/go-gost/core/chain" @@ -102,5 +104,34 @@ func (tr *Transport) Options() *chain.TransportOptions { func (tr *Transport) Copy() chain.Transporter { tr2 := &Transport{} *tr2 = *tr - return tr + return tr2 +} + +// Retire prevents long-lived dialer sessions owned by an obsolete chain from +// accepting new streams while allowing existing streams to drain. +func (tr *Transport) Retire() { + if tr == nil { + return + } + if retirer, ok := tr.dialer.(interface{ Retire() }); ok { + retirer.Retire() + } + if retirer, ok := tr.connector.(interface{ Retire() }); ok { + retirer.Retire() + } +} + +// Close immediately releases transport-owned dialer and connector resources. +func (tr *Transport) Close() error { + if tr == nil { + return nil + } + var errs []error + if closer, ok := tr.dialer.(io.Closer); ok { + errs = append(errs, closer.Close()) + } + if closer, ok := tr.connector.(io.Closer); ok { + errs = append(errs, closer.Close()) + } + return errors.Join(errs...) } diff --git a/go-gost/x/chain/transport_test.go b/go-gost/x/chain/transport_test.go new file mode 100644 index 0000000..c80bc5b --- /dev/null +++ b/go-gost/x/chain/transport_test.go @@ -0,0 +1,45 @@ +package chain + +import ( + "context" + "net" + "testing" + + corechain "github.com/go-gost/core/chain" +) + +type copyTestRoute struct{} + +func (copyTestRoute) Dial(context.Context, string, string, ...corechain.DialOption) (net.Conn, error) { + return nil, nil +} + +func (copyTestRoute) Bind(context.Context, string, string, ...corechain.BindOption) (net.Listener, error) { + return nil, nil +} + +func (copyTestRoute) Nodes() []*corechain.Node { + return nil +} + +func TestTransportCopyReturnsIndependentTransport(t *testing.T) { + originalRoute := copyTestRoute{} + replacementRoute := ©TestRoute{} + original := NewTransport(nil, nil, corechain.RouteTransportOption(originalRoute)) + + copied, ok := original.Copy().(*Transport) + if !ok { + t.Fatalf("copy type = %T, want *Transport", original.Copy()) + } + if copied == original { + t.Fatal("Copy returned the original transport") + } + + copied.Options().Route = replacementRoute + if original.Options().Route != originalRoute { + t.Fatal("mutating copied transport changed original route") + } + if copied.Options().Route != replacementRoute { + t.Fatal("copied transport did not retain its independent route") + } +} diff --git a/go-gost/x/config/parsing/chain/parse.go b/go-gost/x/config/parsing/chain/parse.go index b801188..b4e1e0c 100644 --- a/go-gost/x/config/parsing/chain/parse.go +++ b/go-gost/x/config/parsing/chain/parse.go @@ -35,16 +35,18 @@ func ParseChain(cfg *config.ChainConfig, log logger.Logger) (chain.Chainer, erro for _, ch := range cfg.Hops { var hop hop.Hop var err error + owned := false if ch.Nodes != nil || ch.Plugin != nil { if hop, err = hop_parser.ParseHop(ch, log); err != nil { return nil, err } + owned = true } else { hop = registry.HopRegistry().Get(ch.Name) } if hop != nil { - c.AddHop(hop) + c.AddHop(hop, owned) } } diff --git a/go-gost/x/dialer/kcp/conn.go b/go-gost/x/dialer/kcp/conn.go index 41cf3d5..9e753ac 100644 --- a/go-gost/x/dialer/kcp/conn.go +++ b/go-gost/x/dialer/kcp/conn.go @@ -19,20 +19,23 @@ func (session *muxSession) Accept() (net.Conn, error) { } func (session *muxSession) Close() error { - if session.session == nil { + if session == nil || session.session == nil { return nil } return session.session.Close() } func (session *muxSession) IsClosed() bool { - if session.session == nil { + if session == nil || session.session == nil { return true } return session.session.IsClosed() } func (session *muxSession) NumStreams() int { + if session == nil || session.session == nil { + return 0 + } return session.session.NumStreams() } diff --git a/go-gost/x/dialer/kcp/dialer.go b/go-gost/x/dialer/kcp/dialer.go index 3651bea..0706cd0 100644 --- a/go-gost/x/dialer/kcp/dialer.go +++ b/go-gost/x/dialer/kcp/dialer.go @@ -11,6 +11,7 @@ import ( "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" kcp_util "github.com/go-gost/x/internal/util/kcp" + "github.com/go-gost/x/internal/util/sessionretire" mdutil "github.com/go-gost/x/metadata/util" "github.com/go-gost/x/registry" "github.com/xtaci/kcp-go/v5" @@ -25,6 +26,7 @@ func init() { type kcpDialer struct { sessions map[string]*muxSession sessionMutex sync.Mutex + retired bool logger logger.Logger md metadata options dialer.Options @@ -64,6 +66,9 @@ func (d *kcpDialer) Dial(ctx context.Context, addr string, opts ...dialer.DialOp d.sessionMutex.Lock() defer d.sessionMutex.Unlock() + if d.retired { + return nil, net.ErrClosed + } session, ok := d.sessions[addr] if session != nil && session.IsClosed() { @@ -171,3 +176,36 @@ func (d *kcpDialer) initSession(ctx context.Context, addr net.Addr, conn net.Pac func (d *kcpDialer) Multiplex() bool { return true } + +// Retire drains existing streams and closes their backing sessions once idle. +func (d *kcpDialer) Retire() { + for _, session := range d.detachSessions() { + sessionretire.Gracefully(session) + } +} + +// Close immediately releases all cached multiplex sessions. +func (d *kcpDialer) Close() error { + var errs []error + for _, session := range d.detachSessions() { + errs = append(errs, session.Close()) + } + return errors.Join(errs...) +} + +func (d *kcpDialer) detachSessions() []*muxSession { + if d == nil { + return nil + } + d.sessionMutex.Lock() + d.retired = true + sessions := make([]*muxSession, 0, len(d.sessions)) + for _, session := range d.sessions { + if session != nil { + sessions = append(sessions, session) + } + } + d.sessions = make(map[string]*muxSession) + d.sessionMutex.Unlock() + return sessions +} diff --git a/go-gost/x/dialer/kcp/dialer_test.go b/go-gost/x/dialer/kcp/dialer_test.go new file mode 100644 index 0000000..7e9db3e --- /dev/null +++ b/go-gost/x/dialer/kcp/dialer_test.go @@ -0,0 +1,17 @@ +package kcp + +import ( + "context" + "errors" + "net" + "testing" +) + +func TestRetiredDialerRejectsNewConnections(t *testing.T) { + dialer := NewDialer().(*kcpDialer) + dialer.Retire() + + if _, err := dialer.Dial(context.Background(), "127.0.0.1:1"); !errors.Is(err, net.ErrClosed) { + t.Fatalf("Dial error = %v, want net.ErrClosed", err) + } +} diff --git a/go-gost/x/dialer/mtcp/conn.go b/go-gost/x/dialer/mtcp/conn.go index 491449b..9d2d27e 100644 --- a/go-gost/x/dialer/mtcp/conn.go +++ b/go-gost/x/dialer/mtcp/conn.go @@ -20,13 +20,24 @@ func (session *muxSession) Accept() (net.Conn, error) { } func (session *muxSession) Close() error { + if session == nil { + return nil + } if session.session == nil { + if session.conn != nil { + conn := session.conn + session.conn = nil + return conn.Close() + } return nil } return session.session.Close() } func (session *muxSession) IsClosed() bool { + if session == nil { + return true + } if session.session == nil { return true } @@ -34,5 +45,8 @@ func (session *muxSession) IsClosed() bool { } func (session *muxSession) NumStreams() int { + if session == nil || session.session == nil { + return 0 + } return session.session.NumStreams() } diff --git a/go-gost/x/dialer/mtcp/dialer.go b/go-gost/x/dialer/mtcp/dialer.go index 7a1c5eb..8b74244 100644 --- a/go-gost/x/dialer/mtcp/dialer.go +++ b/go-gost/x/dialer/mtcp/dialer.go @@ -11,6 +11,7 @@ import ( "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" "github.com/go-gost/x/internal/util/mux" + "github.com/go-gost/x/internal/util/sessionretire" "github.com/go-gost/x/registry" ) @@ -21,6 +22,7 @@ func init() { type mtcpDialer struct { sessions map[string]*muxSession sessionMutex sync.Mutex + retired bool logger logger.Logger md metadata options dialer.Options @@ -55,6 +57,9 @@ func (d *mtcpDialer) Multiplex() bool { func (d *mtcpDialer) Dial(ctx context.Context, addr string, opts ...dialer.DialOption) (conn net.Conn, err error) { d.sessionMutex.Lock() defer d.sessionMutex.Unlock() + if d.retired { + return nil, net.ErrClosed + } session, ok := d.sessions[addr] if session != nil && session.IsClosed() { @@ -88,6 +93,10 @@ func (d *mtcpDialer) Handshake(ctx context.Context, conn net.Conn, options ...di d.sessionMutex.Lock() defer d.sessionMutex.Unlock() + if d.retired { + conn.Close() + return nil, net.ErrClosed + } if d.md.handshakeTimeout > 0 { conn.SetDeadline(time.Now().Add(d.md.handshakeTimeout)) @@ -129,3 +138,34 @@ func (d *mtcpDialer) initSession(ctx context.Context, conn net.Conn) (*muxSessio } return &muxSession{conn: conn, session: session}, nil } + +func (d *mtcpDialer) Retire() { + for _, session := range d.detachSessions() { + sessionretire.Gracefully(session) + } +} + +func (d *mtcpDialer) Close() error { + var errs []error + for _, session := range d.detachSessions() { + errs = append(errs, session.Close()) + } + return errors.Join(errs...) +} + +func (d *mtcpDialer) detachSessions() []*muxSession { + if d == nil { + return nil + } + d.sessionMutex.Lock() + d.retired = true + sessions := make([]*muxSession, 0, len(d.sessions)) + for _, session := range d.sessions { + if session != nil { + sessions = append(sessions, session) + } + } + d.sessions = make(map[string]*muxSession) + d.sessionMutex.Unlock() + return sessions +} diff --git a/go-gost/x/dialer/mtcp/dialer_test.go b/go-gost/x/dialer/mtcp/dialer_test.go new file mode 100644 index 0000000..a9a3575 --- /dev/null +++ b/go-gost/x/dialer/mtcp/dialer_test.go @@ -0,0 +1,30 @@ +package mtcp + +import ( + "context" + "errors" + "net" + "testing" +) + +func TestRetiredDialerRejectsNewConnections(t *testing.T) { + dialer := NewDialer().(*mtcpDialer) + dialer.Retire() + + if _, err := dialer.Dial(context.Background(), "127.0.0.1:1"); !errors.Is(err, net.ErrClosed) { + t.Fatalf("Dial error = %v, want net.ErrClosed", err) + } +} + +func TestSessionCloseReleasesPreHandshakeConnection(t *testing.T) { + conn, peer := net.Pipe() + defer peer.Close() + session := &muxSession{conn: conn} + + if err := session.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if !session.IsClosed() { + t.Fatal("pre-handshake session still reports open after Close") + } +} diff --git a/go-gost/x/dialer/mtls/conn.go b/go-gost/x/dialer/mtls/conn.go index 78d6cde..7413b90 100644 --- a/go-gost/x/dialer/mtls/conn.go +++ b/go-gost/x/dialer/mtls/conn.go @@ -20,13 +20,24 @@ func (session *muxSession) Accept() (net.Conn, error) { } func (session *muxSession) Close() error { + if session == nil { + return nil + } if session.session == nil { + if session.conn != nil { + conn := session.conn + session.conn = nil + return conn.Close() + } return nil } return session.session.Close() } func (session *muxSession) IsClosed() bool { + if session == nil { + return true + } if session.session == nil { return true } @@ -34,5 +45,8 @@ func (session *muxSession) IsClosed() bool { } func (session *muxSession) NumStreams() int { + if session == nil || session.session == nil { + return 0 + } return session.session.NumStreams() } diff --git a/go-gost/x/dialer/mtls/dialer.go b/go-gost/x/dialer/mtls/dialer.go index 2c7b902..b8ceed7 100644 --- a/go-gost/x/dialer/mtls/dialer.go +++ b/go-gost/x/dialer/mtls/dialer.go @@ -12,6 +12,7 @@ import ( "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" "github.com/go-gost/x/internal/util/mux" + "github.com/go-gost/x/internal/util/sessionretire" "github.com/go-gost/x/registry" ) @@ -22,6 +23,7 @@ func init() { type mtlsDialer struct { sessions map[string]*muxSession sessionMutex sync.Mutex + retired bool logger logger.Logger md metadata options dialer.Options @@ -56,6 +58,9 @@ func (d *mtlsDialer) Multiplex() bool { func (d *mtlsDialer) Dial(ctx context.Context, addr string, opts ...dialer.DialOption) (conn net.Conn, err error) { d.sessionMutex.Lock() defer d.sessionMutex.Unlock() + if d.retired { + return nil, net.ErrClosed + } session, ok := d.sessions[addr] if session != nil && session.IsClosed() { @@ -89,6 +94,10 @@ func (d *mtlsDialer) Handshake(ctx context.Context, conn net.Conn, options ...di d.sessionMutex.Lock() defer d.sessionMutex.Unlock() + if d.retired { + conn.Close() + return nil, net.ErrClosed + } if d.md.handshakeTimeout > 0 { conn.SetDeadline(time.Now().Add(d.md.handshakeTimeout)) @@ -136,3 +145,34 @@ func (d *mtlsDialer) initSession(ctx context.Context, conn net.Conn) (*muxSessio } return &muxSession{conn: conn, session: session}, nil } + +func (d *mtlsDialer) Retire() { + for _, session := range d.detachSessions() { + sessionretire.Gracefully(session) + } +} + +func (d *mtlsDialer) Close() error { + var errs []error + for _, session := range d.detachSessions() { + errs = append(errs, session.Close()) + } + return errors.Join(errs...) +} + +func (d *mtlsDialer) detachSessions() []*muxSession { + if d == nil { + return nil + } + d.sessionMutex.Lock() + d.retired = true + sessions := make([]*muxSession, 0, len(d.sessions)) + for _, session := range d.sessions { + if session != nil { + sessions = append(sessions, session) + } + } + d.sessions = make(map[string]*muxSession) + d.sessionMutex.Unlock() + return sessions +} diff --git a/go-gost/x/dialer/mtls/dialer_test.go b/go-gost/x/dialer/mtls/dialer_test.go new file mode 100644 index 0000000..df4ead1 --- /dev/null +++ b/go-gost/x/dialer/mtls/dialer_test.go @@ -0,0 +1,30 @@ +package mtls + +import ( + "context" + "errors" + "net" + "testing" +) + +func TestRetiredDialerRejectsNewConnections(t *testing.T) { + dialer := NewDialer().(*mtlsDialer) + dialer.Retire() + + if _, err := dialer.Dial(context.Background(), "127.0.0.1:1"); !errors.Is(err, net.ErrClosed) { + t.Fatalf("Dial error = %v, want net.ErrClosed", err) + } +} + +func TestSessionCloseReleasesPreHandshakeConnection(t *testing.T) { + conn, peer := net.Pipe() + defer peer.Close() + session := &muxSession{conn: conn} + + if err := session.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if !session.IsClosed() { + t.Fatal("pre-handshake session still reports open after Close") + } +} diff --git a/go-gost/x/dialer/mws/conn.go b/go-gost/x/dialer/mws/conn.go index 6b76ca5..e8c05b2 100644 --- a/go-gost/x/dialer/mws/conn.go +++ b/go-gost/x/dialer/mws/conn.go @@ -20,13 +20,24 @@ func (session *muxSession) Accept() (net.Conn, error) { } func (session *muxSession) Close() error { + if session == nil { + return nil + } if session.session == nil { + if session.conn != nil { + conn := session.conn + session.conn = nil + return conn.Close() + } return nil } return session.session.Close() } func (session *muxSession) IsClosed() bool { + if session == nil { + return true + } if session.session == nil { return true } @@ -34,5 +45,8 @@ func (session *muxSession) IsClosed() bool { } func (session *muxSession) NumStreams() int { + if session == nil || session.session == nil { + return 0 + } return session.session.NumStreams() } diff --git a/go-gost/x/dialer/mws/dialer.go b/go-gost/x/dialer/mws/dialer.go index 9ffe7a7..9ee461d 100644 --- a/go-gost/x/dialer/mws/dialer.go +++ b/go-gost/x/dialer/mws/dialer.go @@ -12,6 +12,7 @@ import ( "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" "github.com/go-gost/x/internal/util/mux" + "github.com/go-gost/x/internal/util/sessionretire" ws_util "github.com/go-gost/x/internal/util/ws" "github.com/go-gost/x/registry" "github.com/gorilla/websocket" @@ -25,6 +26,7 @@ func init() { type mwsDialer struct { sessions map[string]*muxSession sessionMutex sync.Mutex + retired bool tlsEnabled bool md metadata options dialer.Options @@ -70,6 +72,9 @@ func (d *mwsDialer) Multiplex() bool { func (d *mwsDialer) Dial(ctx context.Context, addr string, opts ...dialer.DialOption) (conn net.Conn, err error) { d.sessionMutex.Lock() defer d.sessionMutex.Unlock() + if d.retired { + return nil, net.ErrClosed + } session, ok := d.sessions[addr] if session != nil && session.IsClosed() { @@ -108,6 +113,10 @@ func (d *mwsDialer) Handshake(ctx context.Context, conn net.Conn, options ...dia d.sessionMutex.Lock() defer d.sessionMutex.Unlock() + if d.retired { + conn.Close() + return nil, net.ErrClosed + } session, ok := d.sessions[opts.Addr] if session != nil && session.conn != conn { @@ -208,3 +217,34 @@ func (d *mwsDialer) keepAlive(conn ws_util.WebsocketConn) { conn.SetWriteDeadline(time.Time{}) } } + +func (d *mwsDialer) Retire() { + for _, session := range d.detachSessions() { + sessionretire.Gracefully(session) + } +} + +func (d *mwsDialer) Close() error { + var errs []error + for _, session := range d.detachSessions() { + errs = append(errs, session.Close()) + } + return errors.Join(errs...) +} + +func (d *mwsDialer) detachSessions() []*muxSession { + if d == nil { + return nil + } + d.sessionMutex.Lock() + d.retired = true + sessions := make([]*muxSession, 0, len(d.sessions)) + for _, session := range d.sessions { + if session != nil { + sessions = append(sessions, session) + } + } + d.sessions = make(map[string]*muxSession) + d.sessionMutex.Unlock() + return sessions +} diff --git a/go-gost/x/dialer/mws/dialer_test.go b/go-gost/x/dialer/mws/dialer_test.go new file mode 100644 index 0000000..5c1d812 --- /dev/null +++ b/go-gost/x/dialer/mws/dialer_test.go @@ -0,0 +1,30 @@ +package mws + +import ( + "context" + "errors" + "net" + "testing" +) + +func TestRetiredDialerRejectsNewConnections(t *testing.T) { + dialer := NewDialer().(*mwsDialer) + dialer.Retire() + + if _, err := dialer.Dial(context.Background(), "127.0.0.1:1"); !errors.Is(err, net.ErrClosed) { + t.Fatalf("Dial error = %v, want net.ErrClosed", err) + } +} + +func TestSessionCloseReleasesPreHandshakeConnection(t *testing.T) { + conn, peer := net.Pipe() + defer peer.Close() + session := &muxSession{conn: conn} + + if err := session.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if !session.IsClosed() { + t.Fatal("pre-handshake session still reports open after Close") + } +} diff --git a/go-gost/x/hop/hop.go b/go-gost/x/hop/hop.go index c2b8296..66b88d2 100644 --- a/go-gost/x/hop/hop.go +++ b/go-gost/x/hop/hop.go @@ -3,6 +3,7 @@ package hop import ( "context" "encoding/json" + "errors" "io" "net" "sort" @@ -92,6 +93,7 @@ type chainHop struct { nodes []*chain.Node mu sync.RWMutex cancelFunc context.CancelFunc + stopOnce sync.Once options options } @@ -383,13 +385,52 @@ func (p *chainHop) parseNode(r io.Reader) ([]*chain.Node, error) { return nodes, nil } -func (p *chainHop) Close() error { - p.cancelFunc() - if p.options.fileLoader != nil { - p.options.fileLoader.Close() +func (p *chainHop) stopReload() { + if p == nil { + return } - if p.options.redisLoader != nil { - p.options.redisLoader.Close() - } - return nil + p.stopOnce.Do(func() { + p.cancelFunc() + if p.options.fileLoader != nil { + p.options.fileLoader.Close() + } + if p.options.redisLoader != nil { + p.options.redisLoader.Close() + } + if p.options.httpLoader != nil { + p.options.httpLoader.Close() + } + }) +} + +func (p *chainHop) Retire() { + if p == nil { + return + } + p.stopReload() + for _, node := range p.Nodes() { + if node == nil || node.Options().Transport == nil { + continue + } + if retirer, ok := node.Options().Transport.(interface{ Retire() }); ok { + retirer.Retire() + } + } +} + +func (p *chainHop) Close() error { + if p == nil { + return nil + } + p.stopReload() + var errs []error + for _, node := range p.Nodes() { + if node == nil || node.Options().Transport == nil { + continue + } + if closer, ok := node.Options().Transport.(io.Closer); ok { + errs = append(errs, closer.Close()) + } + } + return errors.Join(errs...) } diff --git a/go-gost/x/internal/util/sessionretire/sessionretire.go b/go-gost/x/internal/util/sessionretire/sessionretire.go new file mode 100644 index 0000000..88d982b --- /dev/null +++ b/go-gost/x/internal/util/sessionretire/sessionretire.go @@ -0,0 +1,58 @@ +package sessionretire + +import "time" + +const ( + defaultIdleGrace = time.Second + defaultPollPeriod = 100 * time.Millisecond +) + +// Session is the lifecycle surface shared by the multiplexed dialers. +type Session interface { + Close() error + IsClosed() bool + NumStreams() int +} + +// Gracefully closes a retired session after all existing streams have drained. +// A short idle grace covers the Dial/Handshake hand-off used by several dialers. +func Gracefully(session Session) { + if session == nil { + return + } + go waitUntilIdle(session, defaultIdleGrace, defaultPollPeriod) +} + +func waitUntilIdle(session Session, idleGrace, pollPeriod time.Duration) { + if session == nil { + return + } + if idleGrace <= 0 { + idleGrace = defaultIdleGrace + } + if pollPeriod <= 0 { + pollPeriod = defaultPollPeriod + } + + ticker := time.NewTicker(pollPeriod) + defer ticker.Stop() + + var idleSince time.Time + for { + if session.IsClosed() { + _ = session.Close() + return + } + if session.NumStreams() == 0 { + if idleSince.IsZero() { + idleSince = time.Now() + } else if time.Since(idleSince) >= idleGrace { + _ = session.Close() + return + } + } else { + idleSince = time.Time{} + } + <-ticker.C + } +} diff --git a/go-gost/x/internal/util/sessionretire/sessionretire_test.go b/go-gost/x/internal/util/sessionretire/sessionretire_test.go new file mode 100644 index 0000000..27b1d09 --- /dev/null +++ b/go-gost/x/internal/util/sessionretire/sessionretire_test.go @@ -0,0 +1,59 @@ +package sessionretire + +import ( + "sync" + "testing" + "time" +) + +type testSession struct { + mu sync.Mutex + streams int + closed bool +} + +func (s *testSession) Close() error { + s.mu.Lock() + s.closed = true + s.mu.Unlock() + return nil +} + +func (s *testSession) IsClosed() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.closed +} + +func (s *testSession) NumStreams() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.streams +} + +func TestWaitUntilIdlePreservesActiveStreams(t *testing.T) { + session := &testSession{streams: 1} + done := make(chan struct{}) + go func() { + waitUntilIdle(session, 20*time.Millisecond, time.Millisecond) + close(done) + }() + + time.Sleep(30 * time.Millisecond) + if session.IsClosed() { + t.Fatal("active session was closed") + } + + session.mu.Lock() + session.streams = 0 + session.mu.Unlock() + + select { + case <-done: + case <-time.After(250 * time.Millisecond): + t.Fatal("idle session was not closed") + } + if !session.IsClosed() { + t.Fatal("retired session did not close after becoming idle") + } +} diff --git a/go-gost/x/registry/chain.go b/go-gost/x/registry/chain.go index 900dcf2..8a6511d 100644 --- a/go-gost/x/registry/chain.go +++ b/go-gost/x/registry/chain.go @@ -28,7 +28,13 @@ func (r *chainRegistry) Register(name string, v chain.Chainer) error { } func (r *chainRegistry) replace(name string, v chain.Chainer) { - r.m.Store(name, v) + old, loaded := r.m.Swap(name, v) + if !loaded { + return + } + if retirer, ok := old.(interface{ Retire() }); ok { + retirer.Retire() + } } func (r *chainRegistry) Get(name string) chain.Chainer { diff --git a/go-gost/x/registry/chain_test.go b/go-gost/x/registry/chain_test.go index 5142e87..10fc584 100644 --- a/go-gost/x/registry/chain_test.go +++ b/go-gost/x/registry/chain_test.go @@ -16,6 +16,15 @@ func (c testChainer) Route(context.Context, string, string, ...chain.RouteOption return c.route } +type retiringTestChainer struct { + testChainer + retired bool +} + +func (c *retiringTestChainer) Retire() { + c.retired = true +} + type testRoute struct { nodes []*chain.Node } @@ -49,3 +58,20 @@ func TestReplaceChainOverwritesExistingRegistration(t *testing.T) { t.Fatalf("expected replacement chain route, got %#v", route) } } + +func TestReplaceChainRetiresPreviousRegistration(t *testing.T) { + name := "replace_chain_retire_tdd" + ChainRegistry().Unregister(name) + defer ChainRegistry().Unregister(name) + + old := &retiringTestChainer{} + if err := ChainRegistry().Register(name, old); err != nil { + t.Fatalf("register old chain: %v", err) + } + if err := ReplaceChain(name, testChainer{}); err != nil { + t.Fatalf("replace chain: %v", err) + } + if !old.retired { + t.Fatal("previous chain was not retired") + } +}