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.