@@ -11,13 +11,15 @@ import (
1111 "sync/atomic"
1212 "time"
1313
14+ atomic2 "github.com/metacubex/mihomo/common/atomic"
1415 "github.com/metacubex/mihomo/common/httputils"
1516 "github.com/metacubex/mihomo/common/once"
1617 "github.com/metacubex/mihomo/component/dialer"
1718 C "github.com/metacubex/mihomo/constant"
1819 "github.com/metacubex/mihomo/transport/vmess"
1920
2021 "github.com/metacubex/http"
22+ "github.com/metacubex/http/httptrace"
2123 "golang.org/x/exp/slices"
2224)
2325
@@ -139,7 +141,7 @@ func (c *Client) resetHealthCheckTimer() {
139141 c .healthCheckTimer .Reset (DefaultHealthCheckTimeout )
140142}
141143
142- func (c * Client ) roundTrip (request * http.Request , conn * httpConn ) {
144+ func (c * Client ) roundTrip (ctx context. Context , request * http.Request , conn * httpConn ) error {
143145 c .startOnce .Do (c .start )
144146 pipeReader , pipeWriter := io .Pipe ()
145147 request .Body = pipeReader
@@ -151,20 +153,36 @@ func (c *Client) roundTrip(request *http.Request, conn *httpConn) {
151153 conn .closeFn = once .OnceFunc (func () {
152154 c .count .Add (- 1 )
153155 })
154- ctx , cancel := context .WithCancel (c .ctx ) // requestCtx must alive during conn not closed
155- conn .cancelFn = cancel // cancel ctx when conn closed
156+ requestCtx , cancel := context .WithCancel (c .ctx ) // requestCtx must alive during conn not closed
157+ conn .cancelFn = cancel // cancel ctx when conn closed
158+
159+ // Use gotConn to detect when TCP connection is established, so we can
160+ // return the conn immediately without waiting for the HTTP response.
161+ gotConn := make (chan bool , 1 )
162+ addrCtx := httputils .NewAddrContext (& conn .NetAddr , requestCtx )
163+ streamCtx := httptrace .WithClientTrace (addrCtx , & httptrace.ClientTrace {
164+ GotConn : func (info httptrace.GotConnInfo ) {
165+ select {
166+ case gotConn <- true :
167+ default : // GotConn maybe called multiple times, ignore the second and later calls
168+ }
169+ },
170+ })
171+
172+ var requestErr atomic2.TypedValue [error ]
156173 go func () {
157- timeout := time .AfterFunc (C .DefaultTCPTimeout , cancel ) // only cancel when RoundTrip timeout
158- defer timeout .Stop () // RoundTrip already returned, stop the timer
159- request = request .WithContext (httputils .NewAddrContext (& conn .NetAddr , ctx ))
174+ request = request .WithContext (streamCtx )
160175 response , err := c .roundTripper .RoundTrip (request )
161176 if err != nil {
177+ requestErr .Store (err )
178+ close (gotConn )
162179 _ = pipeWriter .CloseWithError (err )
163180 _ = pipeReader .CloseWithError (err )
164181 conn .setup (nil , err )
165182 } else if response .StatusCode != http .StatusOK {
166183 _ = response .Body .Close ()
167184 err = fmt .Errorf ("unexpected status code: %d" , response .StatusCode )
185+ requestErr .Store (err )
168186 _ = pipeWriter .CloseWithError (err )
169187 _ = pipeReader .CloseWithError (err )
170188 conn .setup (nil , err )
@@ -173,6 +191,13 @@ func (c *Client) roundTrip(request *http.Request, conn *httpConn) {
173191 conn .setup (response .Body , nil )
174192 }
175193 }()
194+
195+ select {
196+ case <- ctx .Done ():
197+ return ctx .Err ()
198+ case <- gotConn :
199+ return requestErr .Load ()
200+ }
176201}
177202
178203func (c * Client ) newConnectRequest (host , userAgent string ) * http.Request {
@@ -193,21 +218,33 @@ func (c *Client) newConnectRequest(host, userAgent string) *http.Request {
193218func (c * Client ) Dial (ctx context.Context , host string ) (net.Conn , error ) {
194219 request := c .newConnectRequest (host , TCPUserAgent )
195220 conn := & tcpConn {}
196- c .roundTrip (request , & conn .httpConn )
221+ err := c .roundTrip (ctx , request , & conn .httpConn )
222+ if err != nil {
223+ _ = conn .Close ()
224+ return nil , err
225+ }
197226 return conn , nil
198227}
199228
200229func (c * Client ) ListenPacket (ctx context.Context ) (net.PacketConn , error ) {
201230 request := c .newConnectRequest (UDPMagicAddress , UDPUserAgent )
202231 conn := & clientPacketConn {}
203- c .roundTrip (request , & conn .httpConn )
232+ err := c .roundTrip (ctx , request , & conn .httpConn )
233+ if err != nil {
234+ _ = conn .Close ()
235+ return nil , err
236+ }
204237 return conn , nil
205238}
206239
207240func (c * Client ) ListenICMP (ctx context.Context ) (* IcmpConn , error ) {
208241 request := c .newConnectRequest (ICMPMagicAddress , ICMPUserAgent )
209242 conn := & IcmpConn {}
210- c .roundTrip (request , & conn .httpConn )
243+ err := c .roundTrip (ctx , request , & conn .httpConn )
244+ if err != nil {
245+ _ = conn .Close ()
246+ return nil , err
247+ }
211248 return conn , nil
212249}
213250
0 commit comments