Improve network reset

dns/transport/https_transport.go

@@ -5,7 +5,7 @@ import (
 	"errors"
 	"net"
 	"net/http"
-	"net/url"
+	"sync"
 	"sync/atomic"
 
 	"github.com/sagernet/sing-box/common/tls"
@@ -22,42 +22,50 @@ type HTTPSTransportWrapper struct {
 	http2Transport *http2.Transport
 	httpTransport  *http.Transport
 	fallback       *atomic.Bool
+	connAccess     sync.Mutex
+	connections    map[*httpsTrackedConn]struct{}
+	closed         bool
 }
 
-func NewHTTPSTransportWrapper(dialer N.Dialer, serverAddr M.Socksaddr, destination *url.URL) *HTTPSTransportWrapper {
-	var fallback atomic.Bool
-	if destination.Scheme == "http" {
-		// plain HTTP DoH used by Tailscale
-		fallback.Store(true)
+func NewHTTPSTransportWrapper(dialer N.Dialer, serverAddr M.Socksaddr, fallback *atomic.Bool) *HTTPSTransportWrapper {
+	wrapper := &HTTPSTransportWrapper{
+		fallback:    fallback,
+		connections: make(map[*httpsTrackedConn]struct{}),
 	}
-	return &HTTPSTransportWrapper{
-		http2Transport: &http2.Transport{
-			DialTLSContext: func(ctx context.Context, _, _ string, _ *tls.STDConfig) (net.Conn, error) {
-				resultConn, err := dialer.DialContext(ctx, N.NetworkTCP, serverAddr)
-				if err != nil {
-					return nil, err
+	wrapper.http2Transport = &http2.Transport{
+		DialTLSContext: func(ctx context.Context, _, _ string, _ *tls.STDConfig) (net.Conn, error) {
+			resultConn, err := dialer.DialContext(ctx, N.NetworkTCP, serverAddr)
+			if err != nil {
+				return nil, err
+			}
+			if tlsConn, isTLSConn := resultConn.(tls.Conn); isTLSConn {
+				state := tlsConn.ConnectionState()
+				if state.NegotiatedProtocol != http2.NextProtoTLS {
+					tlsConn.Close()
+					fallback.Store(true)
+					return nil, errFallback
 				}
-				if tlsConn, isTLSConn := resultConn.(tls.Conn); isTLSConn {
-					state := tlsConn.ConnectionState()
-					if state.NegotiatedProtocol != http2.NextProtoTLS {
-						tlsConn.Close()
-						fallback.Store(true)
-						return nil, errFallback
-					}
-				}
-				return resultConn, nil
-			},
+			}
+			return wrapper.trackConn(resultConn)
 		},
-		httpTransport: &http.Transport{
-			DialContext: func(ctx context.Context, _, addr string) (net.Conn, error) {
-				return dialer.DialContext(ctx, N.NetworkTCP, serverAddr)
-			},
-			DialTLSContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
-				return dialer.DialContext(ctx, N.NetworkTCP, serverAddr)
-			},
-		},
-		fallback: &fallback,
 	}
+	wrapper.httpTransport = &http.Transport{
+		DialContext: func(ctx context.Context, _, addr string) (net.Conn, error) {
+			resultConn, err := dialer.DialContext(ctx, N.NetworkTCP, serverAddr)
+			if err != nil {
+				return nil, err
+			}
+			return wrapper.trackConn(resultConn)
+		},
+		DialTLSContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
+			resultConn, err := dialer.DialContext(ctx, N.NetworkTCP, serverAddr)
+			if err != nil {
+				return nil, err
+			}
+			return wrapper.trackConn(resultConn)
+		},
+	}
+	return wrapper
 }
 
 func (h *HTTPSTransportWrapper) RoundTrip(request *http.Request) (*http.Response, error) {
@@ -74,17 +82,47 @@ func (h *HTTPSTransportWrapper) RoundTrip(request *http.Request) (*http.Response
 	return response, nil
 }
 
-func (h *HTTPSTransportWrapper) CloseIdleConnections() {
+func (h *HTTPSTransportWrapper) trackConn(conn net.Conn) (net.Conn, error) {
+	trackedConn := &httpsTrackedConn{Conn: conn, wrapper: h}
+	h.connAccess.Lock()
+	if h.closed {
+		h.connAccess.Unlock()
+		conn.Close()
+		return nil, net.ErrClosed
+	}
+	h.connections[trackedConn] = struct{}{}
+	h.connAccess.Unlock()
+	return trackedConn, nil
+}
+
+func (h *HTTPSTransportWrapper) Close() {
+	h.connAccess.Lock()
+	if h.closed {
+		h.connAccess.Unlock()
+		return
+	}
+	h.closed = true
+	connections := make([]*httpsTrackedConn, 0, len(h.connections))
+	for trackedConn := range h.connections {
+		connections = append(connections, trackedConn)
+	}
+	h.connections = nil
+	h.connAccess.Unlock()
+	for _, trackedConn := range connections {
+		trackedConn.Conn.Close()
+	}
 	h.http2Transport.CloseIdleConnections()
 	h.httpTransport.CloseIdleConnections()
 }
 
-func (h *HTTPSTransportWrapper) Clone() *HTTPSTransportWrapper {
-	return &HTTPSTransportWrapper{
-		httpTransport: h.httpTransport,
-		http2Transport: &http2.Transport{
-			DialTLSContext: h.http2Transport.DialTLSContext,
-		},
-		fallback: h.fallback,
-	}
+type httpsTrackedConn struct {
+	net.Conn
+	wrapper *HTTPSTransportWrapper
+}
+
+func (c *httpsTrackedConn) Close() error {
+	c.wrapper.connAccess.Lock()
+	delete(c.wrapper.connections, c)
+	c.wrapper.connAccess.Unlock()
+	return c.Conn.Close()
 }