fix(net): close losing dials and retire replaced transports

Co-authored-by: Ginder Singh <ginder@windscribe.com>

config.go

@@ -534,11 +534,36 @@ func (uc *UpstreamConfig) ForceReBootstrap() {
 	uc.rebootstrap.Store(rebootstrapNotStarted)
 }
 
+// CloseTransports retires the upstream's transports, including active DoT and
+// HTTP/3 connections. Callers must stop publishing this upstream before calling
+// it and must not reuse it without calling SetupTransport.
+func (uc *UpstreamConfig) CloseTransports() {
+	uc.retireTransports()
+}
+
+// retireTransports permanently closes DoT pools and HTTP/3 transports before
+// their owner discards them. Like closeTransports, it does not acquire an
+// upstream lock: callers that serialize slot access must keep that serialization
+// across retirement and replacement.
+func (uc *UpstreamConfig) retireTransports() {
+	uc.closeTransports()
+	for _, p := range []*dotConnPool{uc.dotClientPool, uc.dotClientPool4, uc.dotClientPool6} {
+		if p != nil {
+			p.Close()
+		}
+	}
+	// Unlike CloseIdleConnections, Close also cancels pending dials and active
+	// HTTP/3 requests, and prevents late requests from reviving the transport.
+	for _, rt := range []http.RoundTripper{uc.http3RoundTripper, uc.http3RoundTripper4, uc.http3RoundTripper6} {
+		if c, ok := rt.(interface{ Close() error }); ok {
+			_ = c.Close()
+		}
+	}
+}
+
 // closeTransports closes idle connections on all existing transports.
-// This is called before creating new transports during re-bootstrap to
-// force in-flight requests on stale connections to fail quickly, rather
-// than waiting for the full context deadline (e.g. 5s) after a firewall
-// state table flush kills the underlying TCP/QUIC connections.
+// It does not retire DoT pools or HTTP/3 transports: active work may continue
+// and the resources remain usable. Use retireTransports before replacing them.
 func (uc *UpstreamConfig) closeTransports() {
 	if t := uc.transport; t != nil {
 		t.CloseIdleConnections()
@@ -588,11 +613,10 @@ func (uc *UpstreamConfig) SetupTransport() {
 		return
 	}
 
-	// Close existing transport connections before creating new ones.
-	// This forces in-flight requests on stale connections (e.g. after a
-	// firewall state table flush) to fail fast instead of waiting for
-	// the full context deadline timeout.
-	uc.closeTransports()
+	// Retire old DoT/HTTP3 resources before replacing their slots. Idle-only
+	// cleanup would leave active queries and late work owning orphaned resources.
+	// DoH and DoQ retain their existing cleanup semantics.
+	uc.retireTransports()
 
 	ips := uc.bootstrapIPs
 	switch uc.IPStack {

config_quic.go

@@ -31,9 +31,16 @@ func (uc *UpstreamConfig) newDOH3Transport(addrs []string) http.RoundTripper {
 			}
 			remoteAddr, err := net.ResolveUDPAddr("udp", addr)
 			if err != nil {
+				udpConn.Close()
 				return nil, err
 			}
-			return quic.DialEarly(ctx, udpConn, remoteAddr, tlsCfg, cfg)
+			conn, err := quic.DialEarly(ctx, udpConn, remoteAddr, tlsCfg, cfg)
+			if err != nil {
+				udpConn.Close()
+				return nil, err
+			}
+			closeUDPConnWhenDone(conn, udpConn)
+			return conn, nil
 		}
 		dialAddrs := make([]string, len(addrs))
 		for i := range addrs {
@@ -73,18 +80,25 @@ func (uc *UpstreamConfig) dotTransport(dnsType uint16) *dotConnPool {
 //   - quic dialer is different with net.Dialer
 //   - simplification for quic free version
 type parallelDialerResult struct {
-	conn *quic.Conn
-	err  error
+	conn    *quic.Conn
+	udpConn *net.UDPConn
+	err     error
+}
+
+// closeUDPConnWhenDone closes udpConn once conn terminates. quic.DialEarly does
+// not take ownership of the socket it is handed, so the caller must.
+func closeUDPConnWhenDone(conn *quic.Conn, udpConn *net.UDPConn) {
+	go func() {
+		<-conn.Context().Done()
+		_ = udpConn.Close()
+	}()
 }
 
 // quicParallelDialer races DialEarly across a list of remote addresses and
-// returns the first successful connection. When transport is non-nil, all
-// dials share that transport's UDP socket, which removes both the per-dial
-// socket allocation and the winner-path socket leak that an owner-of-the-conn
-// receiver cannot clean up. When transport is nil, the dialer falls back to a
-// fresh UDP socket per attempt (compat path used where no shared transport is
-// available yet); the loser paths close their sockets, and the winner path's
-// socket is owned by quic.DialEarly's internal transport.
+// returns the first successful connection. When transport is non-nil, all dials
+// share that transport's UDP socket. When transport is nil, the dialer falls
+// back to a fresh UDP socket per attempt; losing sockets are closed as their
+// dials settle, and the winning one is closed with the connection it carries.
 type quicParallelDialer struct {
 	transport *quic.Transport
 }
@@ -97,8 +111,6 @@ func (d *quicParallelDialer) Dial(ctx context.Context, addrs []string, tlsCfg *t
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	done := make(chan struct{})
-	defer close(done)
 	ch := make(chan *parallelDialerResult, len(addrs))
 	var wg sync.WaitGroup
 	wg.Add(len(addrs))
@@ -133,16 +145,7 @@ func (d *quicParallelDialer) Dial(ctx context.Context, addrs []string, tlsCfg *t
 					udpConn = nil
 				}
 			}
-			select {
-			case ch <- &parallelDialerResult{conn: conn, err: err}:
-			case <-done:
-				if conn != nil {
-					conn.CloseWithError(quic.ApplicationErrorCode(http3.ErrCodeNoError), "")
-				}
-				if udpConn != nil {
-					udpConn.Close()
-				}
-			}
+			ch <- &parallelDialerResult{conn: conn, udpConn: udpConn, err: err}
 		}(addr)
 	}
 
@@ -150,7 +153,11 @@ func (d *quicParallelDialer) Dial(ctx context.Context, addrs []string, tlsCfg *t
 	for res := range ch {
 		if res.err == nil {
 			cancel()
-			return res.conn, res.err
+			if res.udpConn != nil {
+				closeUDPConnWhenDone(res.conn, res.udpConn)
+			}
+			go closeLosingQuicConns(ch)
+			return res.conn, nil
 		}
 		errs = append(errs, res.err)
 	}
@@ -158,6 +165,18 @@ func (d *quicParallelDialer) Dial(ctx context.Context, addrs []string, tlsCfg *t
 	return nil, errors.Join(errs...)
 }
 
+// closeLosingQuicConns drains ch and releases everything left in it.
+func closeLosingQuicConns(ch <-chan *parallelDialerResult) {
+	for res := range ch {
+		if res.conn != nil {
+			res.conn.CloseWithError(quic.ApplicationErrorCode(http3.ErrCodeNoError), "")
+		}
+		if res.udpConn != nil {
+			_ = res.udpConn.Close()
+		}
+	}
+}
+
 func (uc *UpstreamConfig) newDOQConnPool(addrs []string) *doqConnPool {
 	if uc.Type != ResolverTypeDOQ {
 		return nil

dot.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"net"
 	"runtime"
+	"sync"
 	"time"
 
 	"github.com/miekg/dns"
@@ -36,7 +37,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
@@ -44,6 +46,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 {
@@ -72,6 +80,7 @@ func newDOTClientPool(uc *UpstreamConfig, addrs []string) *dotConnPool {
 		tlsConfig.ServerName = uc.Domain
 	}
 
+	ctx, cancel := context.WithCancel(context.Background())
 	pool := &dotConnPool{
 		uc:        uc,
 		addrs:     addrs,
@@ -79,13 +88,16 @@ func newDOTClientPool(uc *UpstreamConfig, addrs []string) *dotConnPool {
 		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
@@ -115,18 +127,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
@@ -136,87 +158,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 := ProxyLogger.Load()
-	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 {
@@ -232,32 +233,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:
@@ -266,6 +282,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))

internal/net/net.go

@@ -295,8 +295,6 @@ func (d *ParallelDialer) DialContext(ctx context.Context, network string, addrs
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	done := make(chan struct{})
-	defer close(done)
 	ch := make(chan *parallelDialerResult, len(live))
 	var wg sync.WaitGroup
 	wg.Add(len(live))
@@ -318,25 +316,31 @@ func (d *ParallelDialer) DialContext(ctx context.Context, network string, addrs
 			} else {
 				unreachable.markReachable(addr)
 			}
-			select {
-			case ch <- &parallelDialerResult{conn: conn, err: err}:
-			case <-done:
-				if conn != nil {
-					logger.Debug().Msgf("connection closed: %s", conn.RemoteAddr())
-					conn.Close()
-				}
-			}
+			ch <- &parallelDialerResult{conn: conn, err: err}
 		}(addr)
 	}
 
-	errs := make([]error, 0, len(addrs))
+	errs := make([]error, 0, len(live))
 	for res := range ch {
 		if res.err == nil {
 			cancel()
 			logger.Debug().Msgf("connected to %s", res.conn.RemoteAddr())
-			return res.conn, res.err
+			// Losing dials may already have connected, or may still connect
+			// after cancel; close whatever they hand back.
+			go closeLosingConns(ch, logger)
+			return res.conn, nil
 		}
 		errs = append(errs, res.err)
 	}
 	return nil, errors.Join(errs...)
 }
+
+// closeLosingConns drains ch and closes every connection left in it.
+func closeLosingConns(ch <-chan *parallelDialerResult, logger *zerolog.Logger) {
+	for res := range ch {
+		if res.conn != nil {
+			logger.Debug().Msgf("connection closed: %s", res.conn.RemoteAddr())
+			_ = res.conn.Close()
+		}
+	}
+}

transport_replacement_test.go

@@ -0,0 +1,329 @@
+package ctrld
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"encoding/base64"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+	"net/url"
+	"sync"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/miekg/dns"
+	"github.com/quic-go/quic-go"
+	"github.com/quic-go/quic-go/http3"
+)
+
+var transportReplacements = []struct {
+	name string
+	run  func(*UpstreamConfig)
+}{
+	{"SetupTransport", (*UpstreamConfig).SetupTransport},
+	{"ForceReBootstrap", (*UpstreamConfig).ForceReBootstrap},
+}
+
+// Accept multiple connections so a replacement must answer on the same real
+// endpoint, while the original query is still held by the server.
+func replacementDOTServer(t *testing.T) (*UpstreamConfig, <-chan struct{}, func(), <-chan struct{}) {
+	t.Helper()
+	cert := generateTestCertificate(t)
+	ln, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{Certificates: []tls.Certificate{cert.tlsCert}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	entered, release, peerClosed, done := make(chan struct{}), make(chan struct{}), make(chan struct{}), make(chan struct{})
+	var once sync.Once
+	unblock := func() { once.Do(func() { close(release) }) }
+	var mu sync.Mutex
+	var conns []net.Conn
+	var wg sync.WaitGroup
+	go func() {
+		defer close(done)
+		for {
+			c, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			mu.Lock()
+			conns = append(conns, c)
+			mu.Unlock()
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				defer c.Close()
+				dc := &dns.Conn{Conn: c}
+				for {
+					q, err := dc.ReadMsg()
+					if err != nil {
+						return
+					}
+					held := q.Question[0].Name == "held.example."
+					if held {
+						close(entered)
+						defer close(peerClosed)
+						<-release
+					}
+					a := new(dns.Msg)
+					a.SetReply(q)
+					if err := dc.WriteMsg(a); err != nil {
+						return
+					}
+				}
+			}()
+		}
+	}()
+	roots := x509.NewCertPool()
+	roots.AddCert(cert.cert)
+	uc := &UpstreamConfig{Type: ResolverTypeDOT, Endpoint: ln.Addr().String(), BootstrapIP: "127.0.0.1", Domain: "localhost", certPool: roots}
+	uc.ensureSetupTransport()
+	t.Cleanup(func() {
+		unblock()
+		uc.CloseTransports()
+		_ = ln.Close()
+		retirementWait(t, done, "DoT accept loop cleanup") // No more wg.Add after this.
+		mu.Lock()
+		for _, c := range conns {
+			_ = c.Close()
+		}
+		mu.Unlock()
+		wg.Wait()
+	})
+	return uc, entered, unblock, peerClosed
+}
+
+func replacementAnswer(t *testing.T, resolve func(context.Context, *dns.Msg) (*dns.Msg, error)) {
+	t.Helper()
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	q := retirementQuery()
+	a, err := resolve(ctx, q)
+	if err != nil {
+		t.Fatalf("replacement query failed: %v", err)
+	}
+	if a == nil || !a.Response || a.Id != q.Id || len(a.Question) != 1 || a.Question[0] != q.Question[0] {
+		t.Fatalf("replacement returned wrong DNS answer: %v", a)
+	}
+}
+
+func replacementResult(t *testing.T, result <-chan error) error {
+	t.Helper()
+	select {
+	case err := <-result:
+		return err
+	case <-time.After(3 * time.Second):
+		t.Fatal("old query survived replacement")
+		return nil
+	}
+}
+
+func TestReplaceTransportsDoTActive(t *testing.T) {
+	for _, replace := range transportReplacements {
+		t.Run(replace.name, func(t *testing.T) {
+			uc, entered, release, peerClosed := replacementDOTServer(t)
+			old := uc.dotClientPool
+			selected := make(chan struct{})
+			old.dialer.ControlContext = func(context.Context, string, string, syscall.RawConn) error {
+				close(selected)
+				return nil
+			}
+			r := &dotResolver{uc: uc}
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+			q := new(dns.Msg)
+			q.SetQuestion("held.example.", dns.TypeA)
+			result := make(chan error, 1)
+			go func() { _, err := r.Resolve(ctx, q); result <- err }()
+			retirementWait(t, selected, "DoT transport selection")
+			retirementWait(t, entered, "active DoT query")
+			replace.run(uc)
+			if uc.dotClientPool == old {
+				t.Fatal("DoT pool was not replaced")
+			}
+			if err := replacementResult(t, result); !errors.Is(err, net.ErrClosed) {
+				t.Fatalf("old query did not fail from socket retirement: %v", err)
+			}
+			// This answer must arrive before releasing the old server handler.
+			replacementAnswer(t, r.Resolve)
+			release()
+			retirementWait(t, peerClosed, "old DoT peer closure")
+			if _, err := old.Resolve(ctx, retirementQuery()); !errors.Is(err, net.ErrClosed) {
+				t.Fatalf("late query revived replaced DoT pool: %v", err)
+			}
+			replacementAnswer(t, r.Resolve)
+		})
+	}
+}
+
+func TestReplaceTransportsDoTLateReturn(t *testing.T) {
+	for _, replace := range transportReplacements {
+		for _, slot := range []string{"default", "v4", "v6", "aliased"} {
+			t.Run(replace.name+"/"+slot, func(t *testing.T) {
+				uc, _, _, _ := replacementDOTServer(t)
+				old := uc.dotClientPool
+				r := &dotResolver{uc: uc}
+				replacementAnswer(t, r.Resolve)
+				ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+				defer cancel()
+				conn, err := old.getConn(ctx)
+				if err != nil {
+					t.Fatal(err)
+				}
+				// Model each old slot, including the IPv6-to-IPv4 alias, without
+				// requiring an external IPv6 capability probe in this fixture.
+				switch slot {
+				case "v4":
+					uc.dotClientPool, uc.dotClientPool4 = nil, old
+				case "v6":
+					uc.dotClientPool, uc.dotClientPool6 = nil, old
+				case "aliased":
+					uc.dotClientPool4, uc.dotClientPool6 = old, old
+				}
+				replace.run(uc)
+				// A successful borrower returning after the slots change must not
+				// put its connection back into the now-orphaned pool.
+				old.putConn(conn, true)
+				if len(old.conns) != 0 {
+					t.Fatal("late return repopulated replaced DoT pool")
+				}
+				if _, err := old.getConn(ctx); !errors.Is(err, net.ErrClosed) {
+					t.Fatalf("replaced DoT pool accepted late work: %v", err)
+				}
+				replacementAnswer(t, r.Resolve)
+			})
+		}
+	}
+}
+
+func TestReplaceTransportsDoH3IdleAndLateWork(t *testing.T) {
+	for _, replace := range transportReplacements {
+		t.Run(replace.name, func(t *testing.T) {
+			srv := newTestHTTP3Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusOK)
+			}))
+			roots := x509.NewCertPool()
+			roots.AddCert(srv.cert)
+			uc := &UpstreamConfig{Type: ResolverTypeDOH3, Endpoint: "https://" + srv.addr, BootstrapIP: "127.0.0.1", certPool: roots}
+			defer uc.CloseTransports()
+			// The real warmup path initializes the transport and leaves it idle.
+			if err := uc.ErrorPing(); err != nil {
+				t.Fatal(err)
+			}
+			old := uc.http3RoundTripper
+			defer old.(io.Closer).Close()
+			replace.run(uc)
+			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+			defer cancel()
+			req, _ := http.NewRequestWithContext(ctx, "HEAD", uc.Endpoint, nil)
+			resp, err := old.RoundTrip(req)
+			if resp != nil {
+				resp.Body.Close()
+			}
+			if err == nil || ctx.Err() != nil {
+				t.Fatalf("late work revived idle replaced HTTP/3 transport: %v (context: %v)", err, ctx.Err())
+			}
+			if err := uc.ErrorPing(); err != nil {
+				t.Fatalf("replacement HTTP/3 transport is not usable: %v", err)
+			}
+		})
+	}
+}
+
+func TestReplaceTransportsDoH3ActiveAndLateWork(t *testing.T) {
+	for _, replace := range transportReplacements {
+		for _, slot := range []string{"default", "v4", "v6", "aliased"} {
+			t.Run(replace.name+"/"+slot, func(t *testing.T) {
+				entered, release, peerCanceled := make(chan struct{}), make(chan struct{}), make(chan struct{})
+				srv := newTestHTTP3Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					if r.URL.Path == "/late" {
+						w.WriteHeader(http.StatusOK)
+						return
+					}
+					data, err := base64.RawURLEncoding.DecodeString(r.URL.Query().Get("dns"))
+					q := new(dns.Msg)
+					if err != nil || q.Unpack(data) != nil || len(q.Question) != 1 {
+						http.Error(w, "bad DNS query", http.StatusBadRequest)
+						return
+					}
+					if q.Question[0].Name == "held.example." {
+						close(entered)
+						select {
+						case <-r.Context().Done():
+							close(peerCanceled)
+						case <-release:
+						}
+						return
+					}
+					a := new(dns.Msg)
+					a.SetReply(q)
+					data, _ = a.Pack()
+					w.Header().Set("Content-Type", headerApplicationDNS)
+					_, _ = w.Write(data)
+				}))
+				defer close(release)
+				roots := x509.NewCertPool()
+				roots.AddCert(srv.cert)
+				u, err := url.Parse("https://" + srv.addr + "/dns-query")
+				if err != nil {
+					t.Fatal(err)
+				}
+				uc := &UpstreamConfig{Type: ResolverTypeDOH3, Endpoint: u.String(), u: u, BootstrapIP: "127.0.0.1", certPool: roots}
+				uc.ensureSetupTransport()
+				defer uc.CloseTransports()
+				old := uc.http3RoundTripper
+				// Clean the old transport even if a regression leaves it orphaned.
+				defer old.(io.Closer).Close()
+				selected := make(chan struct{})
+				rt := old.(*http3.Transport)
+				dial := rt.Dial
+				var selectedOnce sync.Once
+				rt.Dial = func(ctx context.Context, addr string, tlsCfg *tls.Config, cfg *quic.Config) (*quic.Conn, error) {
+					selectedOnce.Do(func() { close(selected) })
+					return dial(ctx, addr, tlsCfg, cfg)
+				}
+				r := newDohResolver(uc)
+				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+				defer cancel()
+				q := new(dns.Msg)
+				q.SetQuestion("held.example.", dns.TypeA)
+				result := make(chan error, 1)
+				go func() { _, err := r.Resolve(ctx, q); result <- err }()
+				// A network exchange alone isn't a Go happens-before edge. Observe
+				// selection before writing slots; still execute the production dial.
+				retirementWait(t, selected, "HTTP/3 transport selection")
+				retirementWait(t, entered, "active HTTP/3 DNS query")
+				switch slot {
+				case "v4":
+					uc.http3RoundTripper, uc.http3RoundTripper4 = nil, old
+				case "v6":
+					uc.http3RoundTripper, uc.http3RoundTripper6 = nil, old
+				case "aliased":
+					uc.http3RoundTripper4, uc.http3RoundTripper6 = old, old
+				}
+				replace.run(uc)
+				if uc.http3RoundTripper == old {
+					t.Fatal("HTTP/3 transport was not replaced")
+				}
+				retirementWait(t, peerCanceled, "old HTTP/3 request cancellation")
+				if err := replacementResult(t, result); err == nil || ctx.Err() != nil {
+					t.Fatalf("old HTTP/3 query not aborted by replacement: %v (context: %v)", err, ctx.Err())
+				}
+				replacementAnswer(t, r.Resolve)
+				req, _ := http.NewRequestWithContext(ctx, "GET", "https://"+srv.addr+"/late", nil)
+				resp, err := old.RoundTrip(req)
+				if resp != nil {
+					resp.Body.Close()
+				}
+				if err == nil || ctx.Err() != nil {
+					t.Fatalf("late request revived replaced HTTP/3 transport: %v (context: %v)", err, ctx.Err())
+				}
+				replacementAnswer(t, r.Resolve)
+			})
+		}
+	}
+}

transport_retirement_test.go

@@ -0,0 +1,420 @@
+package ctrld
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+	"sync"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/miekg/dns"
+	"github.com/quic-go/quic-go/http3"
+)
+
+func retirementWait(t *testing.T, ch <-chan struct{}, what string) {
+	t.Helper()
+	select {
+	case <-ch:
+	case <-time.After(3 * time.Second):
+		t.Fatalf("timed out waiting for %s", what)
+	}
+}
+
+// One real TLS connection, with the response optionally held until release.
+// Cleanup joins the server goroutine and never touches host DNS or global state.
+func retirementDOTServer(t *testing.T, hold bool) (*UpstreamConfig, <-chan struct{}, func(), <-chan struct{}) {
+	t.Helper()
+	cert := generateTestCertificate(t)
+	ln, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{Certificates: []tls.Certificate{cert.tlsCert}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	entered, release, peerClosed, done := make(chan struct{}), make(chan struct{}), make(chan struct{}), make(chan struct{})
+	var once sync.Once
+	unblock := func() { once.Do(func() { close(release) }) }
+	if !hold {
+		unblock()
+	}
+	var mu sync.Mutex
+	var accepted net.Conn
+	go func() {
+		defer close(done)
+		c, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		mu.Lock()
+		accepted = c
+		mu.Unlock()
+		defer c.Close()
+		dc := &dns.Conn{Conn: c}
+		q, err := dc.ReadMsg()
+		if err != nil {
+			return
+		}
+		close(entered)
+		<-release
+		a := new(dns.Msg)
+		a.SetReply(q)
+		_ = dc.WriteMsg(a)
+		// No more queries are sent on this fixture; EOF proves peer retirement.
+		_, _ = dc.ReadMsg()
+		close(peerClosed)
+	}()
+	roots := x509.NewCertPool()
+	roots.AddCert(cert.cert)
+	uc := &UpstreamConfig{Type: ResolverTypeDOT, Endpoint: ln.Addr().String(), BootstrapIP: "127.0.0.1", Domain: "localhost", certPool: roots}
+	uc.SetupTransport()
+	uc.transportOnce.Do(func() {}) // SetupTransport above is the resolver's initial setup.
+	t.Cleanup(func() {
+		unblock()
+		uc.CloseTransports()
+		_ = ln.Close()
+		mu.Lock()
+		if accepted != nil {
+			_ = accepted.Close()
+		}
+		mu.Unlock()
+		retirementWait(t, done, "DoT server cleanup")
+	})
+	return uc, entered, unblock, peerClosed
+}
+
+func retirementQuery() *dns.Msg {
+	q := new(dns.Msg)
+	q.SetQuestion("retirement.example.", dns.TypeA)
+	return q
+}
+
+func TestCloseTransportsDoTActive(t *testing.T) {
+	uc, entered, release, peerClosed := retirementDOTServer(t, true)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	result := make(chan error, 1)
+	go func() { _, err := (&dotResolver{uc: uc}).Resolve(ctx, retirementQuery()); result <- err }()
+	retirementWait(t, entered, "active DoT query")
+	uc.CloseTransports()
+	select {
+	case err := <-result:
+		if !errors.Is(err, net.ErrClosed) {
+			// miekg/dns has its own two-second read timeout, independent of
+			// the longer context deadline. A timeout is not retirement proof.
+			t.Fatalf("active DoT query was not closed by retirement: %v", err)
+		}
+	case <-time.After(3 * time.Second):
+		t.Fatal("active DoT query survived retirement until its own deadline")
+	}
+	release()
+	retirementWait(t, peerClosed, "DoT peer socket closure")
+	if _, err := (&dotResolver{uc: uc}).Resolve(ctx, retirementQuery()); !errors.Is(err, net.ErrClosed) {
+		t.Fatalf("retired resolver accepted another query: %v", err)
+	}
+	if n := len(uc.dotClientPool.conns); n != 0 {
+		t.Fatalf("late return repopulated pool: %d", n)
+	}
+}
+
+func TestCloseTransportsDoTIdleAndLateReturn(t *testing.T) {
+	for _, slot := range []string{"default", "v4", "v6", "aliased"} {
+		for _, checkedOut := range []bool{false, true} {
+			t.Run(slot+map[bool]string{false: "/idle", true: "/checked-out"}[checkedOut], func(t *testing.T) {
+				uc, _, _, peerClosed := retirementDOTServer(t, false)
+				p := uc.dotClientPool
+				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+				defer cancel()
+				if _, err := p.Resolve(ctx, retirementQuery()); err != nil {
+					t.Fatal(err)
+				}
+				var conn net.Conn
+				if checkedOut {
+					var err error
+					conn, err = p.getConn(ctx)
+					if err != nil {
+						t.Fatal(err)
+					}
+				}
+				switch slot {
+				case "v4":
+					uc.dotClientPool = nil
+					uc.dotClientPool4 = p
+				case "v6":
+					uc.dotClientPool = nil
+					uc.dotClientPool6 = p
+				case "aliased":
+					uc.dotClientPool4 = p
+					uc.dotClientPool6 = p
+				}
+				uc.CloseTransports()
+				retirementWait(t, peerClosed, "idle/checked-out DoT socket closure")
+				// Simulate a response that completed just before retirement and returns late.
+				if conn != nil {
+					p.putConn(conn, true)
+				}
+				uc.CloseTransports() // aliased slots and repeated retirement are safe.
+				if n := len(p.conns); n != 0 {
+					t.Fatalf("retired pool has %d connections", n)
+				}
+				if _, err := p.getConn(ctx); !errors.Is(err, net.ErrClosed) {
+					t.Fatalf("get after close: %v", err)
+				}
+			})
+		}
+	}
+}
+
+func TestCloseTransportsDoTHandshake(t *testing.T) {
+	for _, mode := range []string{"bootstrap", "parallel", "fallback"} {
+		t.Run(mode, func(t *testing.T) {
+			ln, err := net.Listen("tcp", "127.0.0.1:0")
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer ln.Close()
+			entered, peerClosed := make(chan struct{}), make(chan struct{})
+			go func() {
+				c, err := ln.Accept()
+				if err != nil {
+					return
+				}
+				defer c.Close()
+				var b [1]byte
+				if _, err := c.Read(b[:]); err != nil {
+					return
+				}
+				close(entered) // ClientHello was received, but no TLS response is sent.
+				_, _ = io.Copy(io.Discard, c)
+				close(peerClosed)
+			}()
+			uc := &UpstreamConfig{Type: ResolverTypeDOT, Endpoint: ln.Addr().String(), BootstrapIP: "127.0.0.1", Domain: "localhost"}
+			var addrs []string
+			if mode != "bootstrap" {
+				uc.BootstrapIP = ""
+			}
+			if mode == "parallel" {
+				addrs = []string{"127.0.0.1"}
+			}
+			p := newDOTClientPool(uc, addrs)
+			// The raw fallback fixture has no bootstrap hostname to infer SNI from.
+			p.tlsConfig.ServerName = "localhost"
+			uc.dotClientPool = p
+			defer uc.CloseTransports()
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+			result := make(chan error, 1)
+			go func() { _, err := p.Resolve(ctx, retirementQuery()); result <- err }()
+			retirementWait(t, entered, "stalled TLS handshake")
+			uc.CloseTransports()
+			retirementWait(t, peerClosed, "handshaking socket closure")
+			select {
+			case err := <-result:
+				if err == nil {
+					t.Fatal("handshake succeeded after retirement")
+				}
+			case <-time.After(3 * time.Second):
+				t.Fatal("handshake was not canceled")
+			}
+		})
+	}
+}
+
+func TestCloseTransportsDoTPendingDial(t *testing.T) {
+	// Per-pool dialer seam holds a real TCP dial before connect; no global hooks.
+	uc := &UpstreamConfig{Type: ResolverTypeDOT, Endpoint: "127.0.0.1:853", BootstrapIP: "127.0.0.1"}
+	p := newDOTClientPool(uc, nil)
+	uc.dotClientPool = p
+	defer uc.CloseTransports()
+	entered, canceled := make(chan struct{}), make(chan struct{})
+	p.dialer.ControlContext = func(ctx context.Context, _, _ string, _ syscall.RawConn) error {
+		close(entered)
+		<-ctx.Done()
+		close(canceled)
+		return ctx.Err()
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	result := make(chan error, 1)
+	go func() { _, err := p.Resolve(ctx, retirementQuery()); result <- err }()
+	retirementWait(t, entered, "pending TCP dial")
+	uc.CloseTransports()
+	retirementWait(t, canceled, "TCP dial cancellation")
+	select {
+	case err := <-result:
+		if err == nil {
+			t.Fatal("dial succeeded after close")
+		}
+	case <-time.After(3 * time.Second):
+		t.Fatal("dial did not return")
+	}
+}
+
+func TestDoTParallelDialClosesLosers(t *testing.T) {
+	const attempts = 4
+	cert := generateTestCertificate(t)
+	ln, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{Certificates: []tls.Certificate{cert.tlsCert}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+	allAccepted, release := make(chan struct{}), make(chan struct{})
+	closed, received := make(chan struct{}, attempts), make(chan struct{}, attempts)
+	var unblock sync.Once
+	releaseAll := func() { unblock.Do(func() { close(release) }) }
+	defer releaseAll()
+	var wg sync.WaitGroup
+	go func() {
+		for i := 0; i < attempts; i++ {
+			c, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				defer c.Close()
+				defer func() { closed <- struct{}{} }()
+				_ = c.SetDeadline(time.Now().Add(5 * time.Second))
+				<-release
+				var b [1]byte
+				if _, err := c.Read(b[:]); err == nil {
+					received <- struct{}{}
+					_, _ = io.Copy(io.Discard, c)
+				}
+			}()
+		}
+		close(allAccepted)
+	}()
+	roots := x509.NewCertPool()
+	roots.AddCert(cert.cert)
+	uc := &UpstreamConfig{Type: ResolverTypeDOT, Endpoint: ln.Addr().String(), Domain: "localhost", certPool: roots}
+	addrs := make([]string, attempts)
+	for i := range addrs {
+		addrs[i] = "127.0.0.1"
+	}
+	p := newDOTClientPool(uc, addrs)
+	uc.dotClientPool = p
+	defer uc.CloseTransports()
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	type result struct {
+		conn net.Conn
+		err  error
+	}
+	resultCh := make(chan result, 1)
+	go func() { conn, err := p.getConn(ctx); resultCh <- result{conn, err} }()
+	retirementWait(t, allAccepted, "all parallel TCP connections")
+	releaseAll()
+	var conn net.Conn
+	select {
+	case res := <-resultCh:
+		if res.err != nil {
+			t.Fatal(res.err)
+		}
+		conn = res.conn
+	case <-time.After(3 * time.Second):
+		t.Fatal("parallel dial failed to return a winner")
+	}
+	// A dial-context cancellation must close losers, not the winning connection.
+	if _, err := conn.Write([]byte{42}); err != nil {
+		t.Fatal(err)
+	}
+	retirementWait(t, received, "winner data")
+	for i := 0; i < attempts-1; i++ {
+		retirementWait(t, closed, "losing parallel socket closure")
+	}
+	uc.CloseTransports()
+	retirementWait(t, closed, "winning socket retirement")
+	p.putConn(conn, true)
+	wg.Wait()
+}
+
+func TestDoTIdleCleanupPreservesActive(t *testing.T) {
+	uc, entered, release, _ := retirementDOTServer(t, true)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	result := make(chan error, 1)
+	go func() { _, err := uc.dotClientPool.Resolve(ctx, retirementQuery()); result <- err }()
+	retirementWait(t, entered, "active query")
+	uc.closeTransports() // Idle-only cleanup, not transport replacement.
+	release()
+	select {
+	case err := <-result:
+		if err != nil {
+			t.Fatalf("idle cleanup aborted active query: %v", err)
+		}
+	case <-time.After(3 * time.Second):
+		t.Fatal("query did not finish")
+	}
+	conn, err := uc.dotClientPool.getConn(ctx)
+	if err != nil {
+		t.Fatalf("idle cleanup retired pool: %v", err)
+	}
+	uc.dotClientPool.putConn(conn, true)
+}
+
+func TestCloseTransportsDoH3Active(t *testing.T) {
+	for _, slot := range []string{"default", "v4", "v6", "aliased"} {
+		t.Run(slot, func(t *testing.T) {
+			entered, release, peerCanceled := make(chan struct{}), make(chan struct{}), make(chan struct{})
+			srv := newTestHTTP3Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				close(entered)
+				select {
+				case <-r.Context().Done():
+					close(peerCanceled)
+				case <-release:
+				}
+			}))
+			defer close(release)
+			roots := x509.NewCertPool()
+			roots.AddCert(srv.cert)
+			uc := &UpstreamConfig{Type: ResolverTypeDOH3, BootstrapIP: "127.0.0.1", certPool: roots}
+			rt := uc.newDOH3Transport(nil).(*http3.Transport)
+			defer rt.Close()
+			switch slot {
+			case "default":
+				uc.http3RoundTripper = rt
+			case "v4":
+				uc.http3RoundTripper4 = rt
+			case "v6":
+				uc.http3RoundTripper6 = rt
+			case "aliased":
+				uc.http3RoundTripper = rt
+				uc.http3RoundTripper4 = rt
+				uc.http3RoundTripper6 = rt
+			}
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+			request := func() error {
+				req, _ := http.NewRequestWithContext(ctx, "GET", "https://"+srv.addr, nil)
+				resp, err := rt.RoundTrip(req)
+				if err == nil {
+					defer resp.Body.Close()
+					_, err = io.ReadAll(resp.Body)
+				}
+				return err
+			}
+			result := make(chan error, 1)
+			go func() { result <- request() }()
+			retirementWait(t, entered, "active HTTP/3 request")
+			uc.CloseTransports()
+			retirementWait(t, peerCanceled, "HTTP/3 server request cancellation")
+			select {
+			case err := <-result:
+				if err == nil {
+					t.Fatal("active HTTP/3 request survived close")
+				}
+			case <-time.After(3 * time.Second):
+				t.Fatal("HTTP/3 request was not aborted")
+			}
+			if err := request(); err == nil {
+				t.Fatal("retired HTTP/3 transport accepted another request")
+			}
+			uc.CloseTransports()
+		})
+	}
+}