From 190047c8c0e11a22e800dd8dd3bc2cf74cd6990f Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Fri, 4 Apr 2025 21:05:54 +0800
Subject: [PATCH] fix: grpc transport not apply dial timeout

---
 transport/gun/gun.go             | 13 ++++++++-----
 transport/gun/transport.go       | 16 +++++++++++++---
 transport/gun/transport_close.go |  8 ++------
 3 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/transport/gun/gun.go b/transport/gun/gun.go
index b9665c97..b9f03ce7 100644
--- a/transport/gun/gun.go
+++ b/transport/gun/gun.go
@@ -22,6 +22,7 @@ import (
 	"github.com/metacubex/mihomo/common/buf"
 	"github.com/metacubex/mihomo/common/pool"
 	tlsC "github.com/metacubex/mihomo/component/tls"
+	C "github.com/metacubex/mihomo/constant"
 
 	"golang.org/x/net/http2"
 )
@@ -224,11 +225,9 @@ 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")
-		}
+		ctx, cancel := context.WithTimeout(ctx, C.DefaultTLSTimeout)
+		defer cancel()
 		pconn, err := dialFn(ctx, network, addr)
 		if err != nil {
 			return nil, err
@@ -291,9 +290,12 @@ func NewHTTP2Client(dialFn DialFn, tlsConfig *tls.Config, Fingerprint string, re
 		DisableCompression: true,
 		PingTimeout:        0,
 	}
+
+	ctx, cancel := context.WithCancel(context.Background())
 	wrap := &TransportWrap{
 		Transport: transport,
-		closed:    closed,
+		ctx:       ctx,
+		cancel:    cancel,
 	}
 	return wrap
 }
@@ -320,6 +322,7 @@ func StreamGunWithTransport(transport *TransportWrap, cfg *Config) (net.Conn, er
 		ProtoMinor: 0,
 		Header:     defaultHeader,
 	}
+	request = request.WithContext(transport.ctx)
 
 	conn := &Conn{
 		initFn: func() (io.ReadCloser, netAddr, error) {
diff --git a/transport/gun/transport.go b/transport/gun/transport.go
index 10798823..9c5c4375 100644
--- a/transport/gun/transport.go
+++ b/transport/gun/transport.go
@@ -1,16 +1,26 @@
 package gun
 
 import (
+	"context"
 	"net"
-
-	"github.com/metacubex/mihomo/common/atomic"
+	"sync"
 
 	"golang.org/x/net/http2"
 )
 
 type TransportWrap struct {
 	*http2.Transport
-	closed *atomic.Bool
+	ctx       context.Context
+	cancel    context.CancelFunc
+	closeOnce sync.Once
+}
+
+func (tw *TransportWrap) Close() error {
+	tw.closeOnce.Do(func() {
+		tw.cancel()
+		closeTransport(tw.Transport)
+	})
+	return nil
 }
 
 type netAddr struct {
diff --git a/transport/gun/transport_close.go b/transport/gun/transport_close.go
index 92beca4b..2b30ed7c 100644
--- a/transport/gun/transport_close.go
+++ b/transport/gun/transport_close.go
@@ -45,11 +45,8 @@ func closeClientConn(cc *http2.ClientConn) { // like forceCloseConn() in http2.C
 	_ = cc.Close()
 }
 
-func (tw *TransportWrap) Close() error {
-	if tw.closed.Swap(true) {
-		return nil // already closed
-	}
-	connPool := transportConnPool(tw.Transport)
+func closeTransport(tr *http2.Transport) {
+	connPool := transportConnPool(tr)
 	p := (*clientConnPool)((*efaceWords)(unsafe.Pointer(&connPool)).data)
 	p.mu.Lock()
 	defer p.mu.Unlock()
@@ -61,7 +58,6 @@ func (tw *TransportWrap) Close() error {
 	// cleanup
 	p.conns = make(map[string][]*http2.ClientConn)
 	p.keys = make(map[*http2.ClientConn][]string)
-	return nil
 }
 
 //go:linkname transportConnPool golang.org/x/net/http2.(*Transport).connPool