fix(agent): retire replaced tunnel sessions (#541)

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

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

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