plugin/forward: bound DoT connection setup for retries (#8543)

Signed-off-by: houyuwushang <liuluoqianqiu@outlook.com>

plugin/forward/README.md

@@ -144,6 +144,12 @@ On each endpoint, the timeouts for communication are set as follows:
 
 * The DNS and DoT dial timeout defaults to 30s and can decrease automatically down to 1s based on early results.
   The DoQ handshake timeout is 5s.
+* DoT connection setup (TCP dial plus TLS handshake) is additionally bounded by the remaining
+  5s forwarding retry window, or an earlier request deadline. When retries are enabled, each
+  setup attempt is limited to half of the window available at the start of forwarding (at most
+  2.5s), so a stalled handshake leaves time to try a fresh connection. With `max_connect_attempts 1`,
+  setup may use the full remaining window. Failed handshakes close the connection; successful
+  connections remain reusable. These setup limits do not change the DNS exchange read timeout.
 * The read timeout is static at 2s.
 
 ## Metadata

plugin/forward/dot_test.go

@@ -0,0 +1,199 @@
+package forward
+
+import (
+	"context"
+	"crypto/tls"
+	"fmt"
+	"io"
+	"net"
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/coredns/caddy"
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
+	"github.com/coredns/coredns/plugin/test"
+
+	"github.com/miekg/dns"
+)
+
+// stalledTLSListener accepts TCP but withholds the TLS handshake on the first
+// connections. Later connections are served normally by dns.Server.
+type stalledTLSListener struct {
+	net.Listener
+	config        *tls.Config
+	stalls        int64
+	accepted      atomic.Int64
+	firstAccepted chan struct{}
+	closed        chan error
+	wg            sync.WaitGroup
+}
+
+func (l *stalledTLSListener) Accept() (net.Conn, error) {
+	for {
+		conn, err := l.Listener.Accept()
+		if err != nil {
+			return nil, err
+		}
+		n := l.accepted.Add(1)
+		if n == 1 {
+			close(l.firstAccepted)
+		}
+		if n > l.stalls {
+			return tls.Server(conn, l.config), nil
+		}
+		l.wg.Go(func() {
+			defer conn.Close()
+			// Release an unbounded old client after the forwarding window has elapsed.
+			conn.SetReadDeadline(time.Now().Add(2 * time.Second))
+			_, err := io.Copy(io.Discard, conn)
+			l.closed <- err
+		})
+	}
+}
+
+func TestForwardTLSHandshakeRetry(t *testing.T) {
+	originalTimeout := defaultTimeout
+	defaultTimeout = 800 * time.Millisecond
+	t.Cleanup(func() { defaultTimeout = originalTimeout })
+
+	for _, tc := range []struct {
+		name      string
+		setting   string
+		stalls    int64
+		wantConns int64
+		wantError bool
+		tcp       bool
+		deadline  time.Duration
+		cancel    bool
+		health    bool
+	}{
+		{name: "default", stalls: 1, wantConns: 2, health: true},
+		{name: "TCP downstream", stalls: 1, wantConns: 2, tcp: true},
+		{name: "explicit two attempts", setting: "max_connect_attempts 2", stalls: 1, wantConns: 2},
+		{name: "unlimited attempts", setting: "max_connect_attempts 0", stalls: 1, wantConns: 2},
+		{name: "single attempt", setting: "max_connect_attempts 1", stalls: 1, wantConns: 1, wantError: true},
+		{name: "all handshakes stall", stalls: 2, wantConns: 2, wantError: true},
+		{name: "caller deadline", stalls: 1, wantConns: 2, deadline: 500 * time.Millisecond},
+		{name: "caller canceled", stalls: 1, wantConns: 1, wantError: true, cancel: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			serverTLS, roots := makeForwardDoQTestTLS(t)
+			serverTLS.NextProtos = nil
+			listener, err := net.Listen("tcp", "127.0.0.1:0")
+			if err != nil {
+				t.Fatal(err)
+			}
+			upstream := &stalledTLSListener{
+				Listener: listener, config: serverTLS, stalls: tc.stalls,
+				closed: make(chan error, tc.stalls), firstAccepted: make(chan struct{}),
+			}
+			started := make(chan struct{})
+			server := &dns.Server{
+				Listener: upstream,
+				Handler: dns.HandlerFunc(func(w dns.ResponseWriter, r *dns.Msg) {
+					response := new(dns.Msg)
+					response.SetReply(r)
+					response.Answer = []dns.RR{test.A("example.org. 60 IN A 192.0.2.53")}
+					w.WriteMsg(response)
+				}),
+				NotifyStartedFunc: func() { close(started) },
+			}
+			serverDone := make(chan error, 1)
+			go func() { serverDone <- server.ActivateAndServe() }()
+			<-started
+			t.Cleanup(func() {
+				server.Shutdown()
+				if err := <-serverDone; err != nil {
+					t.Error(err)
+				}
+				upstream.wg.Wait()
+			})
+
+			controller := caddy.NewTestController("dns", fmt.Sprintf(`forward . tls://%s {
+				tls_servername doq.test
+				max_fails 0
+				%s
+			}`, listener.Addr(), tc.setting))
+			fs, err := parseForward(controller)
+			if err != nil {
+				t.Fatal(err)
+			}
+			f := fs[0]
+			p := f.proxies[0]
+			config := p.GetTransport().GetTLSConfig().Clone()
+			config.RootCAs = roots
+			p.SetTLSConfig(config)
+			// Manage the cache explicitly without background health-check connections.
+			runtime.SetFinalizer(p, nil)
+			p.GetTransport().Start()
+			t.Cleanup(p.GetTransport().Stop)
+
+			ctx := t.Context()
+			if tc.deadline != 0 {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithTimeout(ctx, tc.deadline)
+				defer cancel()
+			} else if tc.cancel {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithCancel(ctx)
+				defer cancel()
+				go func() {
+					select {
+					case <-upstream.firstAccepted:
+						cancel()
+					case <-ctx.Done():
+					}
+				}()
+			}
+			msg := new(dns.Msg)
+			msg.SetQuestion("example.org.", dns.TypeA)
+			msg.Id = 1234
+			recorder := dnstest.NewRecorder(&test.ResponseWriter{TCP: tc.tcp})
+			rcode, err := f.ServeDNS(ctx, recorder, msg)
+			if tc.wantError {
+				if err == nil || rcode != dns.RcodeServerFailure {
+					t.Fatalf("rcode = %d, error = %v, want SERVFAIL and an error", rcode, err)
+				}
+			} else {
+				if err != nil || rcode != 0 {
+					t.Fatalf("rcode = %d, error = %v, want a successful retry", rcode, err)
+				}
+				if recorder.Msg == nil || recorder.Msg.Id != 1234 || len(recorder.Msg.Answer) != 1 ||
+					recorder.Msg.Answer[0].String() != "example.org.\t60\tIN\tA\t192.0.2.53" {
+					t.Fatalf("unexpected response: %v", recorder.Msg)
+				}
+				// A successful TLS connection must survive cancellation of its dial context.
+				if _, err := f.ServeDNS(t.Context(), recorder, msg); err != nil {
+					t.Fatalf("cached TLS connection failed: %v", err)
+				}
+			}
+			if got := upstream.accepted.Load(); got != tc.wantConns {
+				t.Errorf("accepted %d connections, want %d", got, tc.wantConns)
+			}
+			if msg.Id != 1234 {
+				t.Errorf("request ID = %d, want 1234", msg.Id)
+			}
+			for range min(tc.stalls, upstream.accepted.Load()) {
+				select {
+				case err := <-upstream.closed:
+					if err != nil {
+						t.Errorf("stalled connection was not closed by the client: %v", err)
+					}
+				case <-time.After(time.Second):
+					t.Fatal("stalled connection was not closed")
+				}
+			}
+			if tc.health {
+				if err := p.GetHealthchecker().Check(p); err != nil {
+					t.Fatalf("TLS health check failed after recovery: %v", err)
+				}
+				if p.Fails() != 0 || upstream.accepted.Load() != tc.wantConns+1 {
+					t.Fatal("health check did not use a fresh, successful TLS connection")
+				}
+			}
+		})
+	}
+}

plugin/forward/forward.go

@@ -141,6 +141,15 @@ func (f *Forward) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg
 		maxConnectAttempts = uint64(defaultConnectAttemptsPerUpstream) * uint64(len(list))
 	}
 	connectAttempts := uint64(0)
+	tlsDeadline := deadline
+	if d, ok := ctx.Deadline(); ok && d.Before(tlsDeadline) {
+		tlsDeadline = d
+	}
+	tlsConnectTimeout := time.Until(tlsDeadline)
+	if maxConnectAttempts != 1 {
+		// Reserve time for a fresh connection if the first TLS handshake stalls.
+		tlsConnectTimeout /= 2
+	}
 
 	for time.Now().Before(deadline) && ctx.Err() == nil && (maxConnectAttempts == 0 || connectAttempts < maxConnectAttempts) {
 		if i >= len(list) {
@@ -187,6 +196,10 @@ func (f *Forward) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg
 		opts := f.opts
 
 		for {
+			opts.TLSConnectDeadline = time.Now().Add(tlsConnectTimeout)
+			if opts.TLSConnectDeadline.After(tlsDeadline) {
+				opts.TLSConnectDeadline = tlsDeadline
+			}
 			ret, localAddr, upstreamProto, err = proxy.Connect(ctx, state, opts)
 
 			if err == proxyPkg.ErrCachedClosed { // The peer closed a cached TCP or QUIC connection before the query was sent.

plugin/pkg/proxy/connect.go

@@ -57,6 +57,20 @@ func (t *Transport) updateDialTimeout(newDialTime time.Duration) {
 
 // Dial dials the address configured in transport, potentially reusing a connection or creating a new one.
 func (t *Transport) Dial(proto string) (*persistConn, bool, error) {
+	return t.DialContext(context.Background(), proto)
+}
+
+// DialContext is like Dial but honors ctx while establishing a new connection,
+// including the TLS handshake. The context does not affect the connection after dialing.
+func (t *Transport) DialContext(ctx context.Context, proto string) (*persistConn, bool, error) {
+	return t.dial(ctx, proto, time.Time{})
+}
+
+func (t *Transport) dial(ctx context.Context, proto string, tlsDeadline time.Time) (*persistConn, bool, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, false, err
+	}
+
 	// If tls has been configured; use it.
 	if t.tlsConfig != nil {
 		proto = "tcp-tls"
@@ -97,6 +111,11 @@ func (t *Transport) Dial(proto string) (*persistConn, bool, error) {
 
 	connCacheMissesCount.WithLabelValues(t.proxyName, t.addr, proto).Add(1)
 
+	if proto == "tcp-tls" && !tlsDeadline.IsZero() {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithDeadline(ctx, tlsDeadline)
+		defer cancel()
+	}
 	reqTime := time.Now()
 	timeout := t.dialTimeout()
 	dialer := &net.Dialer{Timeout: timeout}
@@ -112,13 +131,13 @@ func (t *Transport) Dial(proto string) (*persistConn, bool, error) {
 	// pass nil tlsConfig to use system default
 	client := dns.Client{Net: proto, Dialer: dialer, TLSConfig: t.tlsConfig}
 
-	conn, err := client.Dial(t.addr)
+	conn, err := client.DialContext(ctx, t.addr)
 
 	t.updateDialTimeout(time.Since(reqTime))
 	return &persistConn{c: conn, created: time.Now()}, false, err
 }
 
-func (p *Proxy) lookupDNS(_ctx context.Context, state request.Request, opts Options) (*dns.Msg, net.Addr, string, error) {
+func (p *Proxy) lookupDNS(ctx context.Context, state request.Request, opts Options) (*dns.Msg, net.Addr, string, error) {
 	var proto string
 	switch {
 	case opts.ForceTCP: // TCP flag has precedence over UDP flag
@@ -144,7 +163,7 @@ func (p *Proxy) lookupDNS(_ctx context.Context, state request.Request, opts Opti
 		}
 	}
 
-	pc, cached, err := p.transport.Dial(proto)
+	pc, cached, err := p.transport.dial(ctx, proto, opts.TLSConnectDeadline)
 	if err != nil {
 		return nil, nil, proto, err
 	}

plugin/pkg/proxy/dial_context_test.go

@@ -0,0 +1,157 @@
+package proxy
+
+import (
+	"context"
+	"crypto/tls"
+	"errors"
+	"io"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
+	"github.com/coredns/coredns/plugin/pkg/transport"
+	"github.com/coredns/coredns/plugin/test"
+	"github.com/coredns/coredns/request"
+
+	"github.com/miekg/dns"
+)
+
+func TestConnectTLSHandshakeContext(t *testing.T) {
+	for _, deadline := range []bool{false, true} {
+		name := "cancel"
+		if deadline {
+			name = "deadline"
+		}
+		t.Run(name, func(t *testing.T) {
+			listener, err := net.Listen("tcp", "127.0.0.1:0")
+			if err != nil {
+				t.Fatal(err)
+			}
+			t.Cleanup(func() { listener.Close() })
+			started := make(chan struct{})
+			closed := make(chan error, 1)
+			go func() {
+				conn, err := listener.Accept()
+				if err != nil {
+					closed <- err
+					return
+				}
+				defer conn.Close()
+				// A watchdog bounds the old implementation without waiting for its 30s dial timeout.
+				conn.SetReadDeadline(time.Now().Add(time.Second))
+				var first [1]byte
+				if _, err := io.ReadFull(conn, first[:]); err != nil {
+					closed <- err
+					return
+				}
+				close(started)
+				_, err = io.Copy(io.Discard, conn)
+				closed <- err
+			}()
+
+			ctx, cancel := context.WithCancel(t.Context())
+			wantErr := context.Canceled
+			if deadline {
+				cancel()
+				ctx, cancel = context.WithTimeout(t.Context(), 100*time.Millisecond)
+				wantErr = context.DeadlineExceeded
+			}
+			defer cancel()
+			p := NewProxy("forward", listener.Addr().String(), transport.TLS)
+			p.SetTLSConfig(&tls.Config{})
+			msg := new(dns.Msg)
+			msg.SetQuestion("example.org.", dns.TypeA)
+			originalID := msg.Id
+			result := make(chan error, 1)
+			go func() {
+				_, _, _, err := p.Connect(ctx, request.Request{Req: msg, W: &test.ResponseWriter{}}, Options{})
+				result <- err
+			}()
+			select {
+			case <-started:
+			case <-time.After(2 * time.Second):
+				t.Fatal("TLS client did not start the handshake")
+			}
+			if !deadline {
+				cancel()
+			}
+			select {
+			case err := <-result:
+				if !errors.Is(err, wantErr) {
+					t.Errorf("Connect error = %v, want %v", err, wantErr)
+				}
+			case <-time.After(2 * time.Second):
+				t.Fatal("TLS connect did not finish")
+			}
+			if err := <-closed; err != nil {
+				t.Errorf("stalled connection was not closed by the client: %v", err)
+			}
+			if msg.Id != originalID {
+				t.Errorf("request ID = %d, want %d", msg.Id, originalID)
+			}
+			for _, conns := range p.transport.conns {
+				if len(conns) != 0 {
+					t.Fatal("failed TLS connection was cached")
+				}
+			}
+		})
+	}
+}
+
+func TestDialContextCanceledKeepsCachedConnection(t *testing.T) {
+	server := dnstest.NewServer(func(_ dns.ResponseWriter, _ *dns.Msg) {})
+	t.Cleanup(server.Close)
+	for _, proto := range []string{"udp", "tcp"} {
+		t.Run(proto, func(t *testing.T) {
+			tr := newTransport("forward", server.Addr)
+			pc, _, err := tr.Dial(proto)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer pc.c.Close()
+			tr.Yield(pc)
+			ctx, cancel := context.WithCancel(t.Context())
+			cancel()
+			if _, cached, err := tr.DialContext(ctx, proto); !errors.Is(err, context.Canceled) || cached {
+				t.Fatalf("canceled DialContext: cached = %v, error = %v", cached, err)
+			}
+			reused, cached, err := tr.DialContext(t.Context(), proto)
+			if err != nil || !cached || reused != pc {
+				t.Fatalf("cached connection was lost: cached = %v, error = %v", cached, err)
+			}
+		})
+	}
+}
+
+func TestConnectDNSIgnoresTLSConnectDeadline(t *testing.T) {
+	server := dnstest.NewServer(func(w dns.ResponseWriter, r *dns.Msg) {
+		response := new(dns.Msg)
+		response.SetReply(r)
+		response.Answer = []dns.RR{test.A("example.org. 60 IN A 192.0.2.53")}
+		w.WriteMsg(response)
+	})
+	t.Cleanup(server.Close)
+	for _, tcp := range []bool{false, true} {
+		name := "udp"
+		if tcp {
+			name = "tcp"
+		}
+		t.Run(name, func(t *testing.T) {
+			p := NewProxy("forward", server.Addr, transport.DNS)
+			t.Cleanup(func() { p.transport.cleanup(true) })
+			msg := new(dns.Msg)
+			msg.SetQuestion("example.org.", dns.TypeA)
+			opts := Options{TLSConnectDeadline: time.Now().Add(-time.Second)}
+			response, _, proto, err := p.Connect(t.Context(), request.Request{
+				Req: msg, W: &test.ResponseWriter{TCP: tcp},
+			}, opts)
+			if err != nil {
+				t.Fatalf("plain DNS was limited by the TLS deadline: %v", err)
+			}
+			if proto != name || len(response.Answer) != 1 {
+				t.Fatalf("unexpected %s response: %v", proto, response)
+			}
+		})
+	}
+}

plugin/pkg/proxy/errors.go

@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"errors"
+	"time"
 )
 
 var (
@@ -25,4 +26,7 @@ type Options struct {
 	HCRecursionDesired bool
 	// HCDomain sets domain for Proxy healthcheck requests
 	HCDomain string
+	// TLSConnectDeadline bounds the TCP dial and TLS handshake for DoT, not the DNS exchange.
+	// A zero value leaves the adaptive dial timeout and caller context as the only limits.
+	TLSConnectDeadline time.Time
 }