From 24a9ff6d03b89b77454d7ad9f8527651fa6c6f51 Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Fri, 4 Apr 2025 13:33:00 +0800
Subject: [PATCH] fix: disallow dialFunc be called after grpc transport has be
 closed

---
 transport/gun/gun.go             | 11 ++++++++---
 transport/gun/transport.go       |  6 +++++-
 transport/gun/transport_close.go | 16 +++++++++++++---
 3 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/transport/gun/gun.go b/transport/gun/gun.go
index 7738105a..b9665c97 100644
--- a/transport/gun/gun.go
+++ b/transport/gun/gun.go
@@ -46,7 +46,7 @@ type Conn struct {
 
 	reader io.ReadCloser
 	once   sync.Once
-	close  atomic.Bool
+	closed atomic.Bool
 	err    error
 	remain int
 	br     *bufio.Reader
@@ -71,7 +71,7 @@ func (g *Conn) initReader() {
 	}
 	g.netAddr = addr
 
-	if !g.close.Load() {
+	if !g.closed.Load() {
 		g.reader = reader
 		g.br = bufio.NewReader(reader)
 	} else {
@@ -184,7 +184,7 @@ func (g *Conn) FrontHeadroom() int {
 }
 
 func (g *Conn) Close() error {
-	g.close.Store(true)
+	g.closed.Store(true)
 	var errorArr []error
 
 	if reader := g.reader; reader != nil {
@@ -224,7 +224,11 @@ func (g *Conn) SetDeadline(t time.Time) error {
 }
 
 func NewHTTP2Client(dialFn DialFn, tlsConfig *tls.Config, Fingerprint string, realityConfig *tlsC.RealityConfig) *TransportWrap {
+	closed := &atomic.Bool{}
 	dialFunc := func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) {
+		if closed.Load() {
+			return nil, errors.New("transport closed")
+		}
 		pconn, err := dialFn(ctx, network, addr)
 		if err != nil {
 			return nil, err
@@ -289,6 +293,7 @@ func NewHTTP2Client(dialFn DialFn, tlsConfig *tls.Config, Fingerprint string, re
 	}
 	wrap := &TransportWrap{
 		Transport: transport,
+		closed:    closed,
 	}
 	return wrap
 }
diff --git a/transport/gun/transport.go b/transport/gun/transport.go
index 3f913cf9..10798823 100644
--- a/transport/gun/transport.go
+++ b/transport/gun/transport.go
@@ -1,12 +1,16 @@
 package gun
 
 import (
-	"golang.org/x/net/http2"
 	"net"
+
+	"github.com/metacubex/mihomo/common/atomic"
+
+	"golang.org/x/net/http2"
 )
 
 type TransportWrap struct {
 	*http2.Transport
+	closed *atomic.Bool
 }
 
 type netAddr struct {
diff --git a/transport/gun/transport_close.go b/transport/gun/transport_close.go
index 70a39b5b..92beca4b 100644
--- a/transport/gun/transport_close.go
+++ b/transport/gun/transport_close.go
@@ -11,9 +11,13 @@ import (
 )
 
 type clientConnPool struct {
-	t     *http2.Transport
-	mu    sync.Mutex
-	conns map[string][]*http2.ClientConn // key is host:port
+	t *http2.Transport
+
+	mu           sync.Mutex
+	conns        map[string][]*http2.ClientConn // key is host:port
+	dialing      map[string]unsafe.Pointer      // currently in-flight dials
+	keys         map[*http2.ClientConn][]string
+	addConnCalls map[string]unsafe.Pointer // in-flight addConnIfNeeded calls
 }
 
 type clientConn struct {
@@ -42,6 +46,9 @@ func closeClientConn(cc *http2.ClientConn) { // like forceCloseConn() in http2.C
 }
 
 func (tw *TransportWrap) Close() error {
+	if tw.closed.Swap(true) {
+		return nil // already closed
+	}
 	connPool := transportConnPool(tw.Transport)
 	p := (*clientConnPool)((*efaceWords)(unsafe.Pointer(&connPool)).data)
 	p.mu.Lock()
@@ -51,6 +58,9 @@ func (tw *TransportWrap) Close() error {
 			closeClientConn(cc)
 		}
 	}
+	// cleanup
+	p.conns = make(map[string][]*http2.ClientConn)
+	p.keys = make(map[*http2.ClientConn][]string)
 	return nil
 }