fix(net): port transport resource cleanup

dot.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"net"
 	"runtime"
+	"sync"
 	"time"
 
 	"github.com/miekg/dns"
@@ -45,7 +46,8 @@ func (r *dotResolver) Resolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, erro
 
 const dotPoolSize = 16
 
-// dotConnPool manages a pool of TCP/TLS connections for DoT queries using a buffered channel.
+// dotConnPool manages a pool of TCP/TLS connections for DoT queries.
+// mu serializes retirement, ownership and idle queue operations.
 type dotConnPool struct {
 	uc        *UpstreamConfig
 	addrs     []string
@@ -53,6 +55,12 @@ type dotConnPool struct {
 	tlsConfig *tls.Config
 	dialer    *net.Dialer
 	conns     chan *dotConn
+
+	mu     sync.Mutex
+	closed bool
+	owned  map[*tls.Conn]struct{} // handshaking, checked-out and idle connections
+	ctx    context.Context
+	cancel context.CancelFunc
 }
 
 type dotConn struct {
@@ -81,6 +89,7 @@ func newDOTClientPool(_ context.Context, uc *UpstreamConfig, addrs []string) *do
 		tlsConfig.ServerName = uc.Domain
 	}
 
+	ctx, cancel := context.WithCancel(context.Background())
 	pool := &dotConnPool{
 		uc:        uc,
 		addrs:     addrs,
@@ -88,13 +97,16 @@ func newDOTClientPool(_ context.Context, uc *UpstreamConfig, addrs []string) *do
 		tlsConfig: tlsConfig,
 		dialer:    dialer,
 		conns:     make(chan *dotConn, dotPoolSize),
+		owned:     make(map[*tls.Conn]struct{}),
+		ctx:       ctx,
+		cancel:    cancel,
 	}
 
 	// Use SetFinalizer here because we need to call a method on the pool itself.
 	// AddCleanup would require passing the pool as arg (which panics) or capturing
 	// it in a closure (which prevents GC). SetFinalizer is appropriate for this case.
 	runtime.SetFinalizer(pool, func(p *dotConnPool) {
-		p.CloseIdleConnections()
+		p.Close()
 	})
 
 	return pool
@@ -124,18 +136,28 @@ func (p *dotConnPool) Resolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, erro
 }
 
 // getConn gets a TCP/TLS connection from the pool or creates a new one.
-// A connection is taken from the channel while in use; putConn returns it.
+// Checked-out connections remain owned until putConn discards them or Close runs.
 func (p *dotConnPool) getConn(ctx context.Context) (net.Conn, error) {
 	for {
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
+		p.mu.Lock()
+		if p.closed {
+			p.mu.Unlock()
+			return nil, net.ErrClosed
+		}
 		select {
 		case dc := <-p.conns:
+			p.mu.Unlock()
 			if dc.conn != nil && isAlive(dc.conn) {
 				return dc.conn, nil
 			}
 			if dc.conn != nil {
-				dc.conn.Close()
+				p.putConn(dc.conn, false)
 			}
 		default:
+			p.mu.Unlock()
 			_, conn, err := p.dialConn(ctx)
 			if err != nil {
 				return nil, err
@@ -145,87 +167,66 @@ func (p *dotConnPool) getConn(ctx context.Context) (net.Conn, error) {
 	}
 }
 
-// putConn returns a connection to the pool for reuse by other goroutines.
+// putConn cannot repopulate a retired pool, even for a successful late response.
 func (p *dotConnPool) putConn(conn net.Conn, isGood bool) {
-	if !isGood || conn == nil {
-		if conn != nil {
-			conn.Close()
-		}
+	if conn == nil {
 		return
 	}
 	dc := &dotConn{conn: conn.(*tls.Conn)}
-	select {
-	case p.conns <- dc:
-	default:
-		// Channel full, close the connection
-		dc.conn.Close()
+	p.mu.Lock()
+	if isGood && !p.closed {
+		select {
+		case p.conns <- dc:
+			p.mu.Unlock()
+			return
+		default:
+		}
 	}
+	delete(p.owned, dc.conn)
+	p.mu.Unlock()
+	_ = dc.conn.Close()
 }
 
-// dialConn creates a new TCP/TLS connection.
+// dialConn creates a new TCP/TLS connection. Retirement cancels pending TCP
+// dials and handshakes; finishing a parallel race cancels and closes its losers.
 func (p *dotConnPool) dialConn(ctx context.Context) (string, *tls.Conn, error) {
+	ctx, cancel := context.WithCancel(ctx)
+	stop := context.AfterFunc(p.ctx, cancel)
+	defer stop()
+	defer cancel()
+	if p.ctx.Err() != nil {
+		return "", nil, net.ErrClosed
+	}
 	logger := LoggerFromCtx(ctx)
-	var endpoint string
-
 	if p.uc.BootstrapIP != "" {
-		endpoint = net.JoinHostPort(p.uc.BootstrapIP, p.port)
+		endpoint := net.JoinHostPort(p.uc.BootstrapIP, p.port)
 		Log(ctx, logger.Debug(), "Sending DoT request to: %s", endpoint)
-		conn, err := p.dialer.DialContext(ctx, "tcp", endpoint)
-		if err != nil {
-			return "", nil, err
-		}
-		tlsConn := tls.Client(conn, p.tlsConfig)
-		if err := tlsConn.HandshakeContext(ctx); err != nil {
-			conn.Close()
-			return "", nil, err
-		}
-		return endpoint, tlsConn, nil
+		conn, err := p.dialTLSConn(ctx, endpoint, p.tlsConfig)
+		return endpoint, conn, err
 	}
-
-	// Try bootstrap IPs in parallel
 	if len(p.addrs) > 0 {
 		type result struct {
 			conn *tls.Conn
 			addr string
 			err  error
 		}
-
-		ch := make(chan result, len(p.addrs))
-		done := make(chan struct{})
-		defer close(done)
-
+		// An unbuffered handoff leaves every losing connection with its dialer.
+		ch := make(chan result)
 		for _, addr := range p.addrs {
 			go func(addr string) {
 				endpoint := net.JoinHostPort(addr, p.port)
-				conn, err := p.dialer.DialContext(ctx, "tcp", endpoint)
-				if err != nil {
-					select {
-					case ch <- result{conn: nil, addr: endpoint, err: err}:
-					case <-done:
-					}
-					return
-				}
 				tlsConfig := p.tlsConfig.Clone()
 				tlsConfig.ServerName = p.uc.Domain
-				tlsConn := tls.Client(conn, tlsConfig)
-				if err := tlsConn.HandshakeContext(ctx); err != nil {
-					conn.Close()
-					select {
-					case ch <- result{conn: nil, addr: endpoint, err: err}:
-					case <-done:
-					}
-					return
-				}
+				conn, err := p.dialTLSConn(ctx, endpoint, tlsConfig)
 				select {
-				case ch <- result{conn: tlsConn, addr: endpoint, err: nil}:
-				case <-done:
+				case ch <- result{conn: conn, addr: endpoint, err: err}:
+				case <-ctx.Done():
 					if conn != nil {
-						conn.Close()
+						p.putConn(conn, false)
 					}
 				}
 			}(addr)
 		}
-
 		errs := make([]error, 0, len(p.addrs))
 		for range len(p.addrs) {
 			select {
@@ -241,32 +242,47 @@ func (p *dotConnPool) dialConn(ctx context.Context) (string, *tls.Conn, error) {
 				return "", nil, ctx.Err()
 			}
 		}
-
 		return "", nil, errors.Join(errs...)
 	}
-
-	// Fallback to endpoint resolution
-	endpoint = p.uc.Endpoint
+	endpoint := p.uc.Endpoint
 	Log(ctx, logger.Debug(), "Sending DoT request to: %s", endpoint)
-	conn, err := p.dialer.DialContext(ctx, "tcp", endpoint)
-	if err != nil {
-		return "", nil, err
-	}
-	tlsConn := tls.Client(conn, p.tlsConfig)
-	if err := tlsConn.HandshakeContext(ctx); err != nil {
-		conn.Close()
-		return "", nil, err
-	}
-	return endpoint, tlsConn, nil
+	conn, err := p.dialTLSConn(ctx, endpoint, p.tlsConfig)
+	return endpoint, conn, err
 }
 
-// CloseIdleConnections closes all connections in the pool.
-// Connections currently checked out (in use) are not closed.
+// Register ownership before the TLS handshake so Close interrupts stalled peers.
+// A TCP dial racing retirement cannot publish a socket into the closed pool.
+func (p *dotConnPool) dialTLSConn(ctx context.Context, endpoint string, cfg *tls.Config) (*tls.Conn, error) {
+	conn, err := p.dialer.DialContext(ctx, "tcp", endpoint)
+	if err != nil {
+		return nil, err
+	}
+	tlsConn := tls.Client(conn, cfg)
+	p.mu.Lock()
+	if p.closed {
+		p.mu.Unlock()
+		_ = conn.Close()
+		return nil, net.ErrClosed
+	}
+	p.owned[tlsConn] = struct{}{}
+	p.mu.Unlock()
+	if err := tlsConn.HandshakeContext(ctx); err != nil {
+		p.putConn(tlsConn, false)
+		return nil, err
+	}
+	return tlsConn, nil
+}
+
+// CloseIdleConnections closes only queued connections. The pool remains usable
+// and outstanding queries may return their connections. Replacement uses Close.
 func (p *dotConnPool) CloseIdleConnections() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
 	for {
 		select {
 		case dc := <-p.conns:
 			if dc.conn != nil {
+				delete(p.owned, dc.conn)
 				dc.conn.Close()
 			}
 		default:
@@ -275,6 +291,26 @@ func (p *dotConnPool) CloseIdleConnections() {
 	}
 }
 
+// Close permanently retires this pool. Closing the underlying TCP sockets also
+// interrupts TLS reads/writes without waiting to send a close_notify.
+// Keep the channel open: outstanding borrowers may still call putConn.
+func (p *dotConnPool) Close() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.closed {
+		return
+	}
+	p.closed = true
+	p.cancel()
+	for conn := range p.owned {
+		_ = conn.NetConn().Close()
+		delete(p.owned, conn)
+	}
+	for len(p.conns) > 0 {
+		<-p.conns
+	}
+}
+
 func isAlive(c *tls.Conn) bool {
 	// Set a very short deadline for the read
 	c.SetReadDeadline(time.Now().Add(1 * time.Millisecond))