Merge 25d29f5aa11850c021efa9ad98b34402d073c416 into e52d13c59fbe6bb905db149e411564b8596b6f02

commit: 78705f3500b9842c30351109009546aba1d713ba
author: ermite · 31h ago · 2026-10-21 12:38 UTC
merge of: e52d13c + 25d29f5 · diff against the first parent
signature: unsigned

3 files · +515 -51

M docs/sandboxes.md +1 -0


@@ -78,6 +78,7 @@ From a sandbox:
78	78
79	79		- Source address and destination checks run before any "established" accept. A guest cannot forge another's flow.
80	80		- The DNS forwarder allows 32 in-flight queries per listener (extra UDP is dropped), 16 TCP connections, and 30 s per TCP connection.
	81	+	- Upstream queries use UDP for 1 s, then TCP on the same resolver (also when UDP is truncated). One query takes at most 6 s. A UDP answer larger than the client accepts (512 bytes without EDNS) is replaced by an empty reply with TC=1, so the client retries over TCP.
81	82		- Golden rulesets: `internal/netpolicy/testdata`.
82	83
83	84		Host side effects:


@@ -78,6 +78,7 @@ From a sandbox:
78		78
79	- Source address and destination checks run before any "established" accept. A guest cannot forge another's flow.	79		- Source address and destination checks run before any "established" accept. A guest cannot forge another's flow.
80	- The DNS forwarder allows 32 in-flight queries per listener (extra UDP is dropped), 16 TCP connections, and 30 s per TCP connection.	80		- The DNS forwarder allows 32 in-flight queries per listener (extra UDP is dropped), 16 TCP connections, and 30 s per TCP connection.
		81	+	- Upstream queries use UDP for 1 s, then TCP on the same resolver (also when UDP is truncated). One query takes at most 6 s. A UDP answer larger than the client accepts (512 bytes without EDNS) is replaced by an empty reply with TC=1, so the client retries over TCP.
81	- Golden rulesets: `internal/netpolicy/testdata`.	82		- Golden rulesets: `internal/netpolicy/testdata`.
82		83
83	Host side effects:	84		Host side effects:

M internal/dnsfwd/dnsfwd.go +140 -20


@@ -39,8 +39,15 @@ import (
39	39		const (
40	40		// TCPLifetime is the absolute lifetime of a TCP connection, set at accept.
41	41		TCPLifetime = 30 * time.Second
42		-	// UpstreamTimeout is the deadline of one upstream attempt.
	42	+	// UpstreamTimeout is the deadline of one upstream TCP attempt.
43	43		UpstreamTimeout = 3 * time.Second
	44	+	// UDPAttemptTimeout is the deadline of one upstream UDP attempt. It is short
	45	+	// on purpose: glibc guests use "timeout:2 attempts:2", so the reply (or its
	46	+	// TC=1 / TCP fallback) must reach them well inside the first 2 s window.
	47	+	UDPAttemptTimeout = time.Second
	48	+	// MaxExchange caps one whole exchange over all upstreams and transports, so
	49	+	// a long upstream list cannot hold an in-flight slot for long.
	50	+	MaxExchange = 6 * time.Second
44	51		// TCPIdleTimeout closes TCP client connections that stay silent.
45	52		TCPIdleTimeout = 10 * time.Second
46	53

@@ -53,6 +60,9 @@ const (
53	60
54	61		maxUDP = 4096 // largest datagram relayed (EDNS0 sizes up to 4096)
55	62		hdrLen = 12
	63	+
	64	+	// plainUDPSize is the largest UDP reply for a client without EDNS (RFC 1035).
	65	+	plainUDPSize = 512
56	66		)
57	67
58	68		// Forwarder relays DNS queries. The zero value is not usable; use New.

@@ -61,6 +71,11 @@ type Forwarder struct {
61	71		port int
62	72		// Lifetime overrides TCPLifetime; set before Listen (tests).
63	73		Lifetime time.Duration
	74	+	// Upstream deadlines, set to the constants above by New; tests shorten
	75	+	// them before Listen.
	76	+	udpTimeout, tcpTimeout, maxExchange time.Duration
	77	+	// dial opens upstream connections; tests replace it with in-memory fakes.
	78	+	dial func(ctx context.Context, network, addr string) (net.Conn, error)
64	79		// FreeBind binds with IP_FREEBIND, so that the gateway address may be
65	80		// configured after the listener. Off by default: with Docker the address
66	81		// exists when the network is created and a missing one is an error to

@@ -97,7 +112,12 @@ func New(upstreams func() []string, port int) *Forwarder {
97	112		upstreams: upstreams,
98	113		port: port,
99	114		Lifetime: TCPLifetime,
100		-	listeners: map[netip.Addr]*listener{},
	115	+
	116	+	udpTimeout: UDPAttemptTimeout,
	117	+	tcpTimeout: UpstreamTimeout,
	118	+	maxExchange: MaxExchange,
	119	+	dial: (&net.Dialer{}).DialContext,
	120	+	listeners: map[netip.Addr]*listener{},
101	121		}
102	122		}
103	123

@@ -232,6 +252,11 @@ func (f Forwarder) serveUDP(ctx context.Context, l listener, log zerowrap.Logg
232	252		go func() {
233	253		defer func() { <-l.inflight }()
234	254		resp := f.exchange(ctx, l, q, false)
	255	+	// A reply larger than the client accepts would be dropped or cut
	256	+	// by its stack: send TC=1 so that it retries over TCP at once.
	257	+	if len(resp) > clientUDPSize(q) {
	258	+	resp = truncated(resp)
	259	+	}
235	260		if len(resp) > 0 {
236	261		_, _ = pc.WriteTo(resp, from)
237	262		}

@@ -311,34 +336,130 @@ func (f Forwarder) serveConn(ctx context.Context, l listener, c net.Conn, dead
311	336		}
312	337		}
313	338
314		-	// exchange relays q to the upstreams in order and returns the first answer, or
315		-	// a SERVFAIL when none answers. The transport towards the upstream matches the
316		-	// client's: a client that fell back to TCP gets a complete answer. The
	339	+	// exchange relays q to the upstreams in order and returns the first answer,
	340	+	// SERVFAIL when none answers, or nil when ctx is cancelled. A TCP client is
	341	+	// served over TCP. A UDP client is served over UDP for at most udpTimeout; if
	342	+	// that fails, times out or returns TC=1, the same upstream is retried over TCP
	343	+	// (at most tcpTimeout) before the next upstream is tried, because some
	344	+	// resolvers or network paths lose large plain-UDP answers. Each upstream gets
	345	+	// its own budget, and the whole exchange is capped at maxExchange. The
317	346		// overridden name, if any, is answered without asking the upstreams.
318	347		func (f Forwarder) exchange(ctx context.Context, l listener, q []byte, overTCP bool) []byte {
319	348		if resp := f.overrideAnswer(l.addr, q); resp != nil {
320	349		return resp
321	350		}
	351	+	cctx, cancel := context.WithTimeout(ctx, f.maxExchange)
	352	+	defer cancel()
	353	+	valid := func(resp []byte, err error) bool {
	354	+	return err == nil && len(resp) >= hdrLen && resp[0] == q[0] && resp[1] == q[1]
	355	+	}
322	356		for _, up := range f.upstreams() {
323		-	var resp []byte
324		-	var err error
	357	+	if cctx.Err() != nil {
	358	+	break
	359	+	}
325	360		if overTCP {
326		-	resp, err = exchangeTCP(ctx, up, q)
327		-	} else {
328		-	resp, err = exchangeUDP(ctx, up, q)
	361	+	if resp, err := f.exchangeTCP(cctx, up, q, f.tcpTimeout); valid(resp, err) {
	362	+	return resp
	363	+	}
	364	+	continue
	365	+	}
	366	+	resp, err := f.exchangeUDP(cctx, up, q, f.udpTimeout)
	367	+	var tc []byte
	368	+	if valid(resp, err) {
	369	+	if resp[2]&0x02 == 0 {
	370	+	return resp
	371	+	}
	372	+	tc = resp
	373	+	}
	374	+	if full, err := f.exchangeTCP(cctx, up, q, f.tcpTimeout); valid(full, err) {
	375	+	return full
329	376		}
330		-	if err == nil && len(resp) >= hdrLen && resp[0] == q[0] && resp[1] == q[1] {
331		-	return resp
	377	+	if tc != nil {
	378	+	return tc // the client retries over TCP itself
332	379		}
333	380		}
	381	+	if ctx.Err() != nil {
	382	+	return nil
	383	+	}
334	384		return servfail(q)
335	385		}
336	386
337		-	func exchangeUDP(ctx context.Context, up string, q []byte) ([]byte, error) {
338		-	ctx, cancel := context.WithTimeout(ctx, UpstreamTimeout)
	387	+	// clientUDPSize returns the largest UDP reply the client of q accepts: 512
	388	+	// without EDNS, else its advertised OPT size (at least 512, at most maxUDP).
	389	+	func clientUDPSize(q []byte) int {
	390	+	size := plainUDPSize
	391	+	if len(q) < hdrLen {
	392	+	return size
	393	+	}
	394	+	off := hdrLen
	395	+	for i := 0; i < int(binary.BigEndian.Uint16(q[4:])); i++ { // questions
	396	+	if off = skipName(q, off); off < 0 \|\| off+4 > len(q) {
	397	+	return size
	398	+	}
	399	+	off += 4
	400	+	}
	401	+	rrs := int(binary.BigEndian.Uint16(q[6:])) + int(binary.BigEndian.Uint16(q[8:])) + int(binary.BigEndian.Uint16(q[10:]))
	402	+	for i := 0; i < rrs; i++ {
	403	+	if off = skipName(q, off); off < 0 \|\| off+10 > len(q) {
	404	+	return size
	405	+	}
	406	+	if binary.BigEndian.Uint16(q[off:]) == 41 { // OPT: class is the UDP size
	407	+	size = max(size, int(binary.BigEndian.Uint16(q[off+2:])))
	408	+	break
	409	+	}
	410	+	off += 10 + int(binary.BigEndian.Uint16(q[off+8:]))
	411	+	}
	412	+	return min(size, maxUDP)
	413	+	}
	414	+
	415	+	// skipName returns the offset after the name at off, or -1 if malformed.
	416	+	func skipName(m []byte, off int) int {
	417	+	for off < len(m) {
	418	+	switch c := int(m[off]); {
	419	+	case c == 0:
	420	+	return off + 1
	421	+	case c&0xc0 == 0xc0:
	422	+	if off+2 > len(m) {
	423	+	return -1
	424	+	}
	425	+	return off + 2
	426	+	case c&0xc0 != 0:
	427	+	return -1
	428	+	default:
	429	+	off += 1 + c
	430	+	}
	431	+	}
	432	+	return -1
	433	+	}
	434	+
	435	+	// truncated returns the header and question of resp with TC=1 and no records.
	436	+	func truncated(resp []byte) []byte {
	437	+	if len(resp) < hdrLen {
	438	+	return nil
	439	+	}
	440	+	end := hdrLen
	441	+	for i := 0; i < int(binary.BigEndian.Uint16(resp[4:])); i++ {
	442	+	if end = skipName(resp, end); end < 0 \|\| end+4 > len(resp) {
	443	+	end = hdrLen
	444	+	break
	445	+	}
	446	+	end += 4
	447	+	}
	448	+	r := append([]byte(nil), resp[:end]...)
	449	+	r[2] \|= 0x02 // TC
	450	+	if end == hdrLen {
	451	+	r[4], r[5] = 0, 0
	452	+	}
	453	+	for i := 6; i < hdrLen; i++ { // ANCOUNT, NSCOUNT, ARCOUNT
	454	+	r[i] = 0
	455	+	}
	456	+	return r
	457	+	}
	458	+
	459	+	func (f *Forwarder) exchangeUDP(ctx context.Context, up string, q []byte, timeout time.Duration) ([]byte, error) {
	460	+	ctx, cancel := context.WithTimeout(ctx, timeout)
339	461		defer cancel()
340		-	var d net.Dialer
341		-	c, err := d.DialContext(ctx, "udp", up)
	462	+	c, err := f.dial(ctx, "udp", up)
342	463		if err != nil {
343	464		return nil, err
344	465		}

@@ -361,11 +482,10 @@ func exchangeUDP(ctx context.Context, up string, q []byte) ([]byte, error) {
361	482		}
362	483		}
363	484
364		-	func exchangeTCP(ctx context.Context, up string, q []byte) ([]byte, error) {
365		-	ctx, cancel := context.WithTimeout(ctx, UpstreamTimeout)
	485	+	func (f *Forwarder) exchangeTCP(ctx context.Context, up string, q []byte, timeout time.Duration) ([]byte, error) {
	486	+	ctx, cancel := context.WithTimeout(ctx, timeout)
366	487		defer cancel()
367		-	var d net.Dialer
368		-	c, err := d.DialContext(ctx, "tcp", up)
	488	+	c, err := f.dial(ctx, "tcp", up)
369	489		if err != nil {
370	490		return nil, err
371	491		}


@@ -39,8 +39,15 @@ import (
39		const (	39		const (
40		// TCPLifetime is the absolute lifetime of a TCP connection, set at accept.	40		// TCPLifetime is the absolute lifetime of a TCP connection, set at accept.
41		TCPLifetime = 30 * time.Second	41		TCPLifetime = 30 * time.Second
42	-	// UpstreamTimeout is the deadline of one upstream attempt.	42	+	// UpstreamTimeout is the deadline of one upstream TCP attempt.
43		UpstreamTimeout = 3 * time.Second	43		UpstreamTimeout = 3 * time.Second
			44	+	// UDPAttemptTimeout is the deadline of one upstream UDP attempt. It is short
			45	+	// on purpose: glibc guests use "timeout:2 attempts:2", so the reply (or its
			46	+	// TC=1 / TCP fallback) must reach them well inside the first 2 s window.
			47	+	UDPAttemptTimeout = time.Second
			48	+	// MaxExchange caps one whole exchange over all upstreams and transports, so
			49	+	// a long upstream list cannot hold an in-flight slot for long.
			50	+	MaxExchange = 6 * time.Second
44		// TCPIdleTimeout closes TCP client connections that stay silent.	51		// TCPIdleTimeout closes TCP client connections that stay silent.
45		TCPIdleTimeout = 10 * time.Second	52		TCPIdleTimeout = 10 * time.Second
46			53

@@ -53,6 +60,9 @@ const (
53			60
54		maxUDP = 4096 // largest datagram relayed (EDNS0 sizes up to 4096)	61		maxUDP = 4096 // largest datagram relayed (EDNS0 sizes up to 4096)
55		hdrLen = 12	62		hdrLen = 12
			63	+
			64	+	// plainUDPSize is the largest UDP reply for a client without EDNS (RFC 1035).
			65	+	plainUDPSize = 512
56		)	66		)
57			67
58		// Forwarder relays DNS queries. The zero value is not usable; use New.	68		// Forwarder relays DNS queries. The zero value is not usable; use New.

@@ -61,6 +71,11 @@ type Forwarder struct {
61		port int	71		port int
62		// Lifetime overrides TCPLifetime; set before Listen (tests).	72		// Lifetime overrides TCPLifetime; set before Listen (tests).
63		Lifetime time.Duration	73		Lifetime time.Duration
			74	+	// Upstream deadlines, set to the constants above by New; tests shorten
			75	+	// them before Listen.
			76	+	udpTimeout, tcpTimeout, maxExchange time.Duration
			77	+	// dial opens upstream connections; tests replace it with in-memory fakes.
			78	+	dial func(ctx context.Context, network, addr string) (net.Conn, error)
64		// FreeBind binds with IP_FREEBIND, so that the gateway address may be	79		// FreeBind binds with IP_FREEBIND, so that the gateway address may be
65		// configured after the listener. Off by default: with Docker the address	80		// configured after the listener. Off by default: with Docker the address
66		// exists when the network is created and a missing one is an error to	81		// exists when the network is created and a missing one is an error to

@@ -97,7 +112,12 @@ func New(upstreams func() []string, port int) *Forwarder {
97		upstreams: upstreams,	112		upstreams: upstreams,
98		port: port,	113		port: port,
99		Lifetime: TCPLifetime,	114		Lifetime: TCPLifetime,
100	-	listeners: map[netip.Addr]*listener{},	115	+
			116	+	udpTimeout: UDPAttemptTimeout,
			117	+	tcpTimeout: UpstreamTimeout,
			118	+	maxExchange: MaxExchange,
			119	+	dial: (&net.Dialer{}).DialContext,
			120	+	listeners: map[netip.Addr]*listener{},
101		}	121		}
102		}	122		}
103			123

@@ -232,6 +252,11 @@ func (f Forwarder) serveUDP(ctx context.Context, l listener, log zerowrap.Logg
232		go func() {	252		go func() {
233		defer func() { <-l.inflight }()	253		defer func() { <-l.inflight }()
234		resp := f.exchange(ctx, l, q, false)	254		resp := f.exchange(ctx, l, q, false)
			255	+	// A reply larger than the client accepts would be dropped or cut
			256	+	// by its stack: send TC=1 so that it retries over TCP at once.
			257	+	if len(resp) > clientUDPSize(q) {
			258	+	resp = truncated(resp)
			259	+	}
235		if len(resp) > 0 {	260		if len(resp) > 0 {
236		_, _ = pc.WriteTo(resp, from)	261		_, _ = pc.WriteTo(resp, from)
237		}	262		}

@@ -311,34 +336,130 @@ func (f Forwarder) serveConn(ctx context.Context, l listener, c net.Conn, dead
311		}	336		}
312		}	337		}
313			338
314	-	// exchange relays q to the upstreams in order and returns the first answer, or	339	+	// exchange relays q to the upstreams in order and returns the first answer,
315	-	// a SERVFAIL when none answers. The transport towards the upstream matches the	340	+	// SERVFAIL when none answers, or nil when ctx is cancelled. A TCP client is
316	-	// client's: a client that fell back to TCP gets a complete answer. The	341	+	// served over TCP. A UDP client is served over UDP for at most udpTimeout; if
			342	+	// that fails, times out or returns TC=1, the same upstream is retried over TCP
			343	+	// (at most tcpTimeout) before the next upstream is tried, because some
			344	+	// resolvers or network paths lose large plain-UDP answers. Each upstream gets
			345	+	// its own budget, and the whole exchange is capped at maxExchange. The
317		// overridden name, if any, is answered without asking the upstreams.	346		// overridden name, if any, is answered without asking the upstreams.
318		func (f Forwarder) exchange(ctx context.Context, l listener, q []byte, overTCP bool) []byte {	347		func (f Forwarder) exchange(ctx context.Context, l listener, q []byte, overTCP bool) []byte {
319		if resp := f.overrideAnswer(l.addr, q); resp != nil {	348		if resp := f.overrideAnswer(l.addr, q); resp != nil {
320		return resp	349		return resp
321		}	350		}
			351	+	cctx, cancel := context.WithTimeout(ctx, f.maxExchange)
			352	+	defer cancel()
			353	+	valid := func(resp []byte, err error) bool {
			354	+	return err == nil && len(resp) >= hdrLen && resp[0] == q[0] && resp[1] == q[1]
			355	+	}
322		for _, up := range f.upstreams() {	356		for _, up := range f.upstreams() {
323	-	var resp []byte	357	+	if cctx.Err() != nil {
324	-	var err error	358	+	break
			359	+	}
325		if overTCP {	360		if overTCP {
326	-	resp, err = exchangeTCP(ctx, up, q)	361	+	if resp, err := f.exchangeTCP(cctx, up, q, f.tcpTimeout); valid(resp, err) {
327	-	} else {	362	+	return resp
328	-	resp, err = exchangeUDP(ctx, up, q)	363	+	}
			364	+	continue
			365	+	}
			366	+	resp, err := f.exchangeUDP(cctx, up, q, f.udpTimeout)
			367	+	var tc []byte
			368	+	if valid(resp, err) {
			369	+	if resp[2]&0x02 == 0 {
			370	+	return resp
			371	+	}
			372	+	tc = resp
			373	+	}
			374	+	if full, err := f.exchangeTCP(cctx, up, q, f.tcpTimeout); valid(full, err) {
			375	+	return full
329		}	376		}
330	-	if err == nil && len(resp) >= hdrLen && resp[0] == q[0] && resp[1] == q[1] {	377	+	if tc != nil {
331	-	return resp	378	+	return tc // the client retries over TCP itself
332		}	379		}
333		}	380		}
			381	+	if ctx.Err() != nil {
			382	+	return nil
			383	+	}
334		return servfail(q)	384		return servfail(q)
335		}	385		}
336			386
337	-	func exchangeUDP(ctx context.Context, up string, q []byte) ([]byte, error) {	387	+	// clientUDPSize returns the largest UDP reply the client of q accepts: 512
338	-	ctx, cancel := context.WithTimeout(ctx, UpstreamTimeout)	388	+	// without EDNS, else its advertised OPT size (at least 512, at most maxUDP).
			389	+	func clientUDPSize(q []byte) int {
			390	+	size := plainUDPSize
			391	+	if len(q) < hdrLen {
			392	+	return size
			393	+	}
			394	+	off := hdrLen
			395	+	for i := 0; i < int(binary.BigEndian.Uint16(q[4:])); i++ { // questions
			396	+	if off = skipName(q, off); off < 0 \|\| off+4 > len(q) {
			397	+	return size
			398	+	}
			399	+	off += 4
			400	+	}
			401	+	rrs := int(binary.BigEndian.Uint16(q[6:])) + int(binary.BigEndian.Uint16(q[8:])) + int(binary.BigEndian.Uint16(q[10:]))
			402	+	for i := 0; i < rrs; i++ {
			403	+	if off = skipName(q, off); off < 0 \|\| off+10 > len(q) {
			404	+	return size
			405	+	}
			406	+	if binary.BigEndian.Uint16(q[off:]) == 41 { // OPT: class is the UDP size
			407	+	size = max(size, int(binary.BigEndian.Uint16(q[off+2:])))
			408	+	break
			409	+	}
			410	+	off += 10 + int(binary.BigEndian.Uint16(q[off+8:]))
			411	+	}
			412	+	return min(size, maxUDP)
			413	+	}
			414	+
			415	+	// skipName returns the offset after the name at off, or -1 if malformed.
			416	+	func skipName(m []byte, off int) int {
			417	+	for off < len(m) {
			418	+	switch c := int(m[off]); {
			419	+	case c == 0:
			420	+	return off + 1
			421	+	case c&0xc0 == 0xc0:
			422	+	if off+2 > len(m) {
			423	+	return -1
			424	+	}
			425	+	return off + 2
			426	+	case c&0xc0 != 0:
			427	+	return -1
			428	+	default:
			429	+	off += 1 + c
			430	+	}
			431	+	}
			432	+	return -1
			433	+	}
			434	+
			435	+	// truncated returns the header and question of resp with TC=1 and no records.
			436	+	func truncated(resp []byte) []byte {
			437	+	if len(resp) < hdrLen {
			438	+	return nil
			439	+	}
			440	+	end := hdrLen
			441	+	for i := 0; i < int(binary.BigEndian.Uint16(resp[4:])); i++ {
			442	+	if end = skipName(resp, end); end < 0 \|\| end+4 > len(resp) {
			443	+	end = hdrLen
			444	+	break
			445	+	}
			446	+	end += 4
			447	+	}
			448	+	r := append([]byte(nil), resp[:end]...)
			449	+	r[2] \|= 0x02 // TC
			450	+	if end == hdrLen {
			451	+	r[4], r[5] = 0, 0
			452	+	}
			453	+	for i := 6; i < hdrLen; i++ { // ANCOUNT, NSCOUNT, ARCOUNT
			454	+	r[i] = 0
			455	+	}
			456	+	return r
			457	+	}
			458	+
			459	+	func (f *Forwarder) exchangeUDP(ctx context.Context, up string, q []byte, timeout time.Duration) ([]byte, error) {
			460	+	ctx, cancel := context.WithTimeout(ctx, timeout)
339		defer cancel()	461		defer cancel()
340	-	var d net.Dialer	462	+	c, err := f.dial(ctx, "udp", up)
341	-	c, err := d.DialContext(ctx, "udp", up)
342		if err != nil {	463		if err != nil {
343		return nil, err	464		return nil, err
344		}	465		}

@@ -361,11 +482,10 @@ func exchangeUDP(ctx context.Context, up string, q []byte) ([]byte, error) {
361		}	482		}
362		}	483		}
363			484
364	-	func exchangeTCP(ctx context.Context, up string, q []byte) ([]byte, error) {	485	+	func (f *Forwarder) exchangeTCP(ctx context.Context, up string, q []byte, timeout time.Duration) ([]byte, error) {
365	-	ctx, cancel := context.WithTimeout(ctx, UpstreamTimeout)	486	+	ctx, cancel := context.WithTimeout(ctx, timeout)
366		defer cancel()	487		defer cancel()
367	-	var d net.Dialer	488	+	c, err := f.dial(ctx, "tcp", up)
368	-	c, err := d.DialContext(ctx, "tcp", up)
369		if err != nil {	489		if err != nil {
370		return nil, err	490		return nil, err
371		}	491		}

M internal/dnsfwd/dnsfwd_test.go +374 -31


@@ -16,6 +16,7 @@ import (
16	16		"sync/atomic"
17	17		"syscall"
18	18		"testing"
	19	+	"testing/synctest"
19	20		"time"
20	21		)
21	22

@@ -39,18 +40,32 @@ func answer(q []byte) []byte {
39	40		return append(r, "ANSWER"...)
40	41		}
41	42
	43	+	// listenUDPAndTCP binds UDP on a free loopback port and TCP on the same port.
	44	+	// The TCP port may already be taken (outbound connections use the same
	45	+	// ephemeral range), so it picks again a few times, like Listen does.
	46	+	func listenUDPAndTCP(t *testing.T) (net.PacketConn, net.Listener) {
	47	+	t.Helper()
	48	+	for attempt := 0; ; attempt++ {
	49	+	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	50	+	if err != nil {
	51	+	t.Skip("no loopback:", err)
	52	+	}
	53	+	tl, err := net.Listen("tcp", pc.LocalAddr().String())
	54	+	if err == nil {
	55	+	return pc, tl
	56	+	}
	57	+	_ = pc.Close()
	58	+	if attempt == 49 \|\| !errors.Is(err, syscall.EADDRINUSE) {
	59	+	t.Fatal(err)
	60	+	}
	61	+	}
	62	+	}
	63	+
42	64		// fakeUpstream serves UDP and TCP on one port of 127.0.0.1 and returns it.
43	65		func fakeUpstream(t *testing.T, handle func(q []byte) []byte) string {
44	66		t.Helper()
45		-	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
46		-	if err != nil {
47		-	t.Skip("no loopback:", err)
48		-	}
	67	+	pc, tl := listenUDPAndTCP(t)
49	68		addr := pc.LocalAddr().String()
50		-	tl, err := net.Listen("tcp", addr)
51		-	if err != nil {
52		-	t.Fatal(err)
53		-	}
54	69		t.Cleanup(func() { _ = pc.Close(); _ = tl.Close() })
55	70		go func() {
56	71		buf := make([]byte, 4096)

@@ -92,8 +107,17 @@ func fakeUpstream(t *testing.T, handle func(q []byte) []byte) string {
92	107		}
93	108
94	109		func start(t testing.T, ups ...string) (Forwarder, string) {
	110	+	t.Helper()
	111	+	return startWith(t, nil, ups...)
	112	+	}
	113	+
	114	+	// startWith is start with a hook to adjust the Forwarder before Listen.
	115	+	func startWith(t testing.T, cfg func(Forwarder), ups ...string) (*Forwarder, string) {
95	116		t.Helper()
96	117		f := New(func() []string { return ups }, 0)
	118	+	if cfg != nil {
	119	+	cfg(f)
	120	+	}
97	121		ip := netip.MustParseAddr("127.0.0.1")
98	122		if err := f.Listen(context.Background(), ip); err != nil {
99	123		t.Fatal(err)

@@ -180,29 +204,6 @@ func TestServfailWhenNoUpstreamAnswers(t *testing.T) {
180	204		}
181	205		}
182	206
183		-	func TestMismatchedUpstreamIDIgnored(t *testing.T) {
184		-	// Upstream answers with a wrong ID: treated as no answer, client gets SERVFAIL.
185		-	up := fakeUpstream(t, func(q []byte) []byte {
186		-	r := answer(q)
187		-	r[0] ^= 0xff
188		-	return r
189		-	})
190		-	f, addr := start(t, up)
191		-	_ = f
192		-	c, _ := net.Dial("udp", addr)
193		-	defer func() { _ = c.Close() }()
194		-	_ = c.SetDeadline(time.Now().Add(2UpstreamTimeout + 2time.Second))
195		-	_, _ = c.Write(query(5, "a.test"))
196		-	buf := make([]byte, 512)
197		-	n, err := c.Read(buf)
198		-	if err != nil {
199		-	t.Fatal(err)
200		-	}
201		-	if buf[3]&0x0f != 2 {
202		-	t.Fatalf("want SERVFAIL, got %x", buf[:n])
203		-	}
204		-	}
205		-
206	207		func TestShortPacketsIgnored(t *testing.T) {
207	208		up := fakeUpstream(t, answer)
208	209		_, addr := start(t, up)

@@ -612,3 +613,345 @@ func TestOverrideIgnoresMalformedAndNonQueries(t *testing.T) {
612	613		}
613	614		}
614	615		}
	616	+
	617	+	// bigAnswer returns an answer to q padded to size bytes, with TC set if tc.
	618	+	func bigAnswer(q []byte, size int, tc bool) []byte {
	619	+	r := answer(q)
	620	+	if tc {
	621	+	r[2] \|= 0x02
	622	+	}
	623	+	for len(r) < size {
	624	+	r = append(r, 'x')
	625	+	}
	626	+	return r
	627	+	}
	628	+
	629	+	func TestLargeAnswerNonEDNSClientGetsTC(t *testing.T) {
	630	+	up := fakeUpstream(t, func(q []byte) []byte { return bigAnswer(q, 700, false) })
	631	+	_, addr := start(t, up)
	632	+	q := queryType(0x3333, "storage.googleapis.com", 1, false)
	633	+	resp := udpExchange(t, addr, q)
	634	+	if resp[2]&0x02 == 0 {
	635	+	t.Fatalf("TC not set: %x", resp)
	636	+	}
	637	+	if len(resp) != len(q) \|\| binary.BigEndian.Uint16(resp[4:]) != 1 \|\| binary.BigEndian.Uint16(resp[6:]) != 0 {
	638	+	t.Fatalf("want header+question only, got %x", resp)
	639	+	}
	640	+	if resp[0] != q[0] \|\| resp[1] != q[1] {
	641	+	t.Fatal("ID changed")
	642	+	}
	643	+	}
	644	+
	645	+	func TestLargeAnswerEDNSClientGetsFullAnswer(t *testing.T) {
	646	+	up := fakeUpstream(t, func(q []byte) []byte { return bigAnswer(q, 700, false) })
	647	+	_, addr := start(t, up)
	648	+	// queryType advertises 0x1000 (4096) bytes.
	649	+	resp := udpExchange(t, addr, queryType(0x4444, "storage.googleapis.com", 1, true))
	650	+	if len(resp) != 700 \|\| resp[2]&0x02 != 0 {
	651	+	t.Fatalf("want the full 700-byte answer, got len %d tc=%d", len(resp), resp[2]&0x02)
	652	+	}
	653	+	}
	654	+
	655	+	// withOPT appends an OPT record (root name, type 41) with the given class and
	656	+	// RDLEN to a query and bumps ARCOUNT.
	657	+	func withOPT(q []byte, size, rdlen uint16, body int) []byte {
	658	+	q = append([]byte(nil), q...)
	659	+	q[11] = 1
	660	+	q = append(q, 0, 0, 41)
	661	+	q = binary.BigEndian.AppendUint16(q, size)
	662	+	q = append(q, 0, 0, 0, 0)
	663	+	q = binary.BigEndian.AppendUint16(q, rdlen)
	664	+	return append(q, make([]byte, body)...)
	665	+	}
	666	+
	667	+	func TestClientUDPSize(t *testing.T) {
	668	+	plain := query(1, "a.b")
	669	+	bigQD := append([]byte(nil), plain[:hdrLen]...)
	670	+	bigQD[4], bigQD[5] = 0xff, 0xff
	671	+	bigQD = append(bigQD, 1, 'a', 0, 0, 1, 0, 1)
	672	+	ptrQ := append([]byte(nil), plain[:hdrLen]...)
	673	+	ptrQ = append(ptrQ, 0xc0, hdrLen, 0, 1, 0, 1) // compression pointer in question
	674	+	tests := []struct {
	675	+	name string
	676	+	q []byte
	677	+	want int
	678	+	}{
	679	+	{"no edns", plain, 512},
	680	+	{"edns 1232", withOPT(plain, 1232, 0, 0), 1232},
	681	+	{"edns below 512", withOPT(plain, 100, 0, 0), 512},
	682	+	{"edns above cap", withOPT(plain, 65535, 0, 0), maxUDP},
	683	+	{"edns 4096", withOPT(plain, 4096, 0, 0), 4096},
	684	+	{"qdcount 65535 short body", bigQD, 512},
	685	+	{"compression pointer in question", ptrQ, 512},
	686	+	{"OPT rdlen overruns buffer", withOPT(plain, 4096, 500, 3), 4096},
	687	+	{"truncated header", plain[:15], 512},
	688	+	{"short", plain[:5], 512},
	689	+	}
	690	+	for _, tt := range tests {
	691	+	t.Run(tt.name, func(t *testing.T) {
	692	+	if got := clientUDPSize(tt.q); got != tt.want {
	693	+	t.Fatalf("got %d, want %d", got, tt.want)
	694	+	}
	695	+	})
	696	+	}
	697	+	}
	698	+
	699	+	func TestTruncated(t *testing.T) {
	700	+	q := queryType(7, "a.b", 1, false)
	701	+	full := bigAnswer(q, 700, false)
	702	+	full[7] = 5 // ANCOUNT
	703	+	got := truncated(full)
	704	+	if len(got) != len(q) \|\| got[2]&0x02 == 0 \|\| got[6] != 0 \|\| got[7] != 0 \|\| got[4] != 0 \|\| got[5] != 1 {
	705	+	t.Fatalf("bad truncated reply %x", got)
	706	+	}
	707	+	bad := append([]byte(nil), full...)
	708	+	bad[4], bad[5] = 0xff, 0xff // QDCOUNT lies, question parse fails
	709	+	got = truncated(bad)
	710	+	if len(got) != hdrLen \|\| got[2]&0x02 == 0 \|\| got[4] != 0 \|\| got[5] != 0 {
	711	+	t.Fatalf("bad-question case: %x", got)
	712	+	}
	713	+	if truncated(full[:5]) != nil {
	714	+	t.Fatal("short input must give nil")
	715	+	}
	716	+	}
	717	+
	718	+	func TestServfailAndOverrideRepliesAreNotTruncated(t *testing.T) {
	719	+	dead := "127.0.0.1:1" // nothing listens: UDP and TCP both fail at once
	720	+	_, addr := startOverride(t, "forge.test", dead)
	721	+	resp := udpExchange(t, addr, queryType(9, "x.test", 1, false))
	722	+	if len(resp) != hdrLen \|\| resp[3]&0x0f != 2 \|\| resp[2]&0x02 != 0 {
	723	+	t.Fatalf("want a plain 12-byte SERVFAIL, got %x", resp)
	724	+	}
	725	+	resp = udpExchange(t, addr, queryType(10, "forge.test", 1, false))
	726	+	if resp[2]&0x02 != 0 \|\| resp[3]&0x0f != 0 \|\| binary.BigEndian.Uint16(resp[6:]) != 1 {
	727	+	t.Fatalf("want the override A answer, got %x", resp)
	728	+	}
	729	+	}
	730	+
	731	+	// fakeNet is an in-memory upstream network for Forwarder.dial. One Write on a
	732	+	// UDP conn is one datagram (net.Pipe delivers a Write to a single Read when the
	733	+	// reader's buffer is large enough). A handler returning nil never answers.
	734	+	type fakeNet struct {
	735	+	udp, tcp func(addr string, q []byte) []byte
	736	+	udpCalls, tcpCalls atomic.Int64
	737	+	}
	738	+
	739	+	func (n *fakeNet) dial(ctx context.Context, network, addr string) (net.Conn, error) {
	740	+	if err := ctx.Err(); err != nil { // like net.Dialer
	741	+	return nil, err
	742	+	}
	743	+	c, s := net.Pipe()
	744	+	switch network {
	745	+	case "udp":
	746	+	go func() {
	747	+	defer func() { _ = s.Close() }()
	748	+	buf := make([]byte, 4096)
	749	+	for {
	750	+	k, err := s.Read(buf)
	751	+	if err != nil {
	752	+	return
	753	+	}
	754	+	n.udpCalls.Add(1)
	755	+	if r := n.udp(addr, append([]byte(nil), buf[:k]...)); r != nil {
	756	+	if _, err := s.Write(r); err != nil {
	757	+	return
	758	+	}
	759	+	}
	760	+	}
	761	+	}()
	762	+	case "tcp":
	763	+	go func() {
	764	+	defer func() { _ = s.Close() }()
	765	+	var l [2]byte
	766	+	if _, err := io.ReadFull(s, l[:]); err != nil {
	767	+	return
	768	+	}
	769	+	q := make([]byte, binary.BigEndian.Uint16(l[:]))
	770	+	if _, err := io.ReadFull(s, q); err != nil {
	771	+	return
	772	+	}
	773	+	n.tcpCalls.Add(1)
	774	+	if r := n.tcp(addr, q); r != nil {
	775	+	_, _ = s.Write(append(binary.BigEndian.AppendUint16(nil, uint16(len(r))), r...))
	776	+	}
	777	+	_, _ = io.Copy(io.Discard, s) // wait for the client to hang up
	778	+	}()
	779	+	default:
	780	+	return nil, errors.New("unexpected network " + network)
	781	+	}
	782	+	return c, nil
	783	+	}
	784	+
	785	+	// silent is a handler that never answers.
	786	+	func silent(string, []byte) []byte { return nil }
	787	+
	788	+	// fakeForwarder returns a Forwarder with production timeouts whose upstreams
	789	+	// are served by n, and a listener for exchange.
	790	+	func fakeForwarder(n fakeNet, ups ...string) (Forwarder, *listener) {
	791	+	f := New(func() []string { return ups }, 0)
	792	+	f.dial = n.dial
	793	+	return f, &listener{addr: netip.MustParseAddr("127.0.0.1")}
	794	+	}
	795	+
	796	+	func isServfail(r []byte, id uint16) bool {
	797	+	return len(r) == hdrLen && binary.BigEndian.Uint16(r) == id && r[3]&0x0f == 2
	798	+	}
	799	+
	800	+	// The tests below run in a synctest bubble: the fake clock makes the real
	801	+	// production timeouts (1 s UDP, 3 s TCP, 6 s cap) instant and exact.
	802	+
	803	+	func TestExchangeSilentUDPFallsBackToTCPAfterUDPTimeout(t *testing.T) {
	804	+	synctest.Test(t, func(t *testing.T) {
	805	+	n := &fakeNet{udp: silent, tcp: func(_ string, q []byte) []byte { return answer(q) }}
	806	+	f, l := fakeForwarder(n, "up:53")
	807	+	start := time.Now()
	808	+	resp := f.exchange(t.Context(), l, query(1, "storage.googleapis.com"), false)
	809	+	if !strings.HasSuffix(string(resp), "ANSWER") {
	810	+	t.Fatalf("want the TCP answer, got %x", resp)
	811	+	}
	812	+	if got := time.Since(start); got != UDPAttemptTimeout {
	813	+	t.Fatalf("answered after %v, want %v", got, UDPAttemptTimeout)
	814	+	}
	815	+	if n.udpCalls.Load() != 1 \|\| n.tcpCalls.Load() != 1 {
	816	+	t.Fatalf("udp=%d tcp=%d, want 1 and 1", n.udpCalls.Load(), n.tcpCalls.Load())
	817	+	}
	818	+	})
	819	+	}
	820	+
	821	+	func TestExchangeTruncatedUDPRetriedOverTCP(t *testing.T) {
	822	+	synctest.Test(t, func(t *testing.T) {
	823	+	n := &fakeNet{
	824	+	udp: func(_ string, q []byte) []byte { return bigAnswer(q, 0, true) },
	825	+	tcp: func(_ string, q []byte) []byte { return bigAnswer(q, 700, false) },
	826	+	}
	827	+	f, l := fakeForwarder(n, "up:53")
	828	+	start := time.Now()
	829	+	resp := f.exchange(t.Context(), l, query(2, "example.com"), false)
	830	+	if len(resp) != 700 \|\| resp[2]&0x02 != 0 {
	831	+	t.Fatalf("want the full TCP answer, got len %d tc=%d", len(resp), resp[2]&0x02)
	832	+	}
	833	+	if time.Since(start) != 0 {
	834	+	t.Fatalf("took %v, want no waiting", time.Since(start))
	835	+	}
	836	+	})
	837	+	}
	838	+
	839	+	func TestExchangeKeepsTruncatedReplyWhenTCPFails(t *testing.T) {
	840	+	synctest.Test(t, func(t *testing.T) {
	841	+	n := &fakeNet{udp: func(_ string, q []byte) []byte { return bigAnswer(q, 0, true) }, tcp: silent}
	842	+	f, l := fakeForwarder(n, "up:53")
	843	+	resp := f.exchange(t.Context(), l, query(3, "example.com"), false)
	844	+	if len(resp) < hdrLen \|\| resp[2]&0x02 == 0 {
	845	+	t.Fatalf("want the TC=1 reply, got %x", resp)
	846	+	}
	847	+	})
	848	+	}
	849	+
	850	+	func TestExchangeSecondUpstreamNotStarved(t *testing.T) {
	851	+	synctest.Test(t, func(t *testing.T) {
	852	+	n := &fakeNet{
	853	+	udp: func(a string, q []byte) []byte {
	854	+	if a == "dead:53" {
	855	+	return nil
	856	+	}
	857	+	return answer(q)
	858	+	},
	859	+	tcp: func(a string, q []byte) []byte {
	860	+	if a == "dead:53" {
	861	+	return nil
	862	+	}
	863	+	return answer(q)
	864	+	},
	865	+	}
	866	+	f, l := fakeForwarder(n, "dead:53", "live:53")
	867	+	start := time.Now()
	868	+	resp := f.exchange(t.Context(), l, query(4, "a.test"), false)
	869	+	if !strings.HasSuffix(string(resp), "ANSWER") {
	870	+	t.Fatalf("want the second upstream's answer, got %x", resp)
	871	+	}
	872	+	if got, want := time.Since(start), UDPAttemptTimeout+UpstreamTimeout; got != want {
	873	+	t.Fatalf("took %v, want %v", got, want)
	874	+	}
	875	+	})
	876	+	}
	877	+
	878	+	func TestExchangeMismatchedUpstreamIDIgnored(t *testing.T) {
	879	+	synctest.Test(t, func(t *testing.T) {
	880	+	wrong := func(_ string, q []byte) []byte {
	881	+	r := answer(q)
	882	+	r[0] ^= 0xff
	883	+	return r
	884	+	}
	885	+	f, l := fakeForwarder(&fakeNet{udp: wrong, tcp: wrong}, "up:53")
	886	+	start := time.Now()
	887	+	resp := f.exchange(t.Context(), l, query(5, "a.test"), false)
	888	+	if !isServfail(resp, 5) {
	889	+	t.Fatalf("want SERVFAIL, got %x", resp)
	890	+	}
	891	+	// UDP keeps listening for the right ID until its timeout; the TCP
	892	+	// reply is read once and rejected.
	893	+	if got := time.Since(start); got != UDPAttemptTimeout {
	894	+	t.Fatalf("took %v, want %v", got, UDPAttemptTimeout)
	895	+	}
	896	+	})
	897	+	}
	898	+
	899	+	func TestExchangeCappedAtMaxExchange(t *testing.T) {
	900	+	synctest.Test(t, func(t *testing.T) {
	901	+	n := &fakeNet{udp: silent, tcp: silent}
	902	+	f, l := fakeForwarder(n, "a:53", "b:53", "c:53")
	903	+	start := time.Now()
	904	+	resp := f.exchange(t.Context(), l, query(6, "a.test"), false)
	905	+	if !isServfail(resp, 6) {
	906	+	t.Fatalf("want SERVFAIL, got %x", resp)
	907	+	}
	908	+	if got := time.Since(start); got != MaxExchange {
	909	+	t.Fatalf("took %v, want %v", got, MaxExchange)
	910	+	}
	911	+	// a: 1 s UDP + 3 s TCP; b: 1 s UDP + 1 s TCP (cut by the cap); c: skipped.
	912	+	if n.udpCalls.Load() != 2 \|\| n.tcpCalls.Load() != 2 {
	913	+	t.Fatalf("udp=%d tcp=%d, want 2 and 2", n.udpCalls.Load(), n.tcpCalls.Load())
	914	+	}
	915	+	})
	916	+	}
	917	+
	918	+	func TestExchangeTCPClientUsesTCPOnly(t *testing.T) {
	919	+	synctest.Test(t, func(t *testing.T) {
	920	+	n := &fakeNet{udp: silent, tcp: func(_ string, q []byte) []byte { return answer(q) }}
	921	+	f, l := fakeForwarder(n, "up:53")
	922	+	resp := f.exchange(t.Context(), l, query(7, "a.test"), true)
	923	+	if !strings.HasSuffix(string(resp), "ANSWER") \|\| n.udpCalls.Load() != 0 {
	924	+	t.Fatalf("resp=%x udp=%d", resp, n.udpCalls.Load())
	925	+	}
	926	+	})
	927	+	}
	928	+
	929	+	func TestExchangeCancelledReturnsNil(t *testing.T) {
	930	+	synctest.Test(t, func(t *testing.T) {
	931	+	f, l := fakeForwarder(&fakeNet{udp: silent, tcp: silent}, "up:53")
	932	+	ctx, cancel := context.WithCancel(t.Context())
	933	+	time.AfterFunc(500*time.Millisecond, cancel)
	934	+	start := time.Now()
	935	+	if resp := f.exchange(ctx, l, query(8, "a.test"), false); resp != nil {
	936	+	t.Fatalf("want nil on cancellation, got %x", resp)
	937	+	}
	938	+	// Cancellation is noticed at the next attempt boundary: the blocked UDP
	939	+	// read runs to its own deadline, then no TCP attempt is started.
	940	+	if got := time.Since(start); got != UDPAttemptTimeout {
	941	+	t.Fatalf("returned after %v, want %v", got, UDPAttemptTimeout)
	942	+	}
	943	+	})
	944	+	}
	945	+
	946	+	// End to end through real client-facing sockets; the upstream is in memory.
	947	+	func TestUDPClientGetsTCPAnswerWhenUpstreamUDPSilent(t *testing.T) {
	948	+	n := &fakeNet{udp: silent, tcp: func(_ string, q []byte) []byte { return answer(q) }}
	949	+	_, addr := startWith(t, func(f *Forwarder) {
	950	+	f.dial = n.dial
	951	+	f.udpTimeout = 50 * time.Millisecond
	952	+	}, "up:53")
	953	+	resp := udpExchange(t, addr, query(0x1111, "storage.googleapis.com"))
	954	+	if !strings.HasSuffix(string(resp), "ANSWER") \|\| resp[3]&0x0f != 0 {
	955	+	t.Fatalf("want the TCP answer, got %x", resp)
	956	+	}
	957	+	}


@@ -16,6 +16,7 @@ import (
16		"sync/atomic"	16		"sync/atomic"
17		"syscall"	17		"syscall"
18		"testing"	18		"testing"
			19	+	"testing/synctest"
19		"time"	20		"time"
20		)	21		)
21			22

@@ -39,18 +40,32 @@ func answer(q []byte) []byte {
39		return append(r, "ANSWER"...)	40		return append(r, "ANSWER"...)
40		}	41		}
41			42
			43	+	// listenUDPAndTCP binds UDP on a free loopback port and TCP on the same port.
			44	+	// The TCP port may already be taken (outbound connections use the same
			45	+	// ephemeral range), so it picks again a few times, like Listen does.
			46	+	func listenUDPAndTCP(t *testing.T) (net.PacketConn, net.Listener) {
			47	+	t.Helper()
			48	+	for attempt := 0; ; attempt++ {
			49	+	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
			50	+	if err != nil {
			51	+	t.Skip("no loopback:", err)
			52	+	}
			53	+	tl, err := net.Listen("tcp", pc.LocalAddr().String())
			54	+	if err == nil {
			55	+	return pc, tl
			56	+	}
			57	+	_ = pc.Close()
			58	+	if attempt == 49 \|\| !errors.Is(err, syscall.EADDRINUSE) {
			59	+	t.Fatal(err)
			60	+	}
			61	+	}
			62	+	}
			63	+
42		// fakeUpstream serves UDP and TCP on one port of 127.0.0.1 and returns it.	64		// fakeUpstream serves UDP and TCP on one port of 127.0.0.1 and returns it.
43		func fakeUpstream(t *testing.T, handle func(q []byte) []byte) string {	65		func fakeUpstream(t *testing.T, handle func(q []byte) []byte) string {
44		t.Helper()	66		t.Helper()
45	-	pc, err := net.ListenPacket("udp", "127.0.0.1:0")	67	+	pc, tl := listenUDPAndTCP(t)
46	-	if err != nil {
47	-	t.Skip("no loopback:", err)
48	-	}
49		addr := pc.LocalAddr().String()	68		addr := pc.LocalAddr().String()
50	-	tl, err := net.Listen("tcp", addr)
51	-	if err != nil {
52	-	t.Fatal(err)
53	-	}
54		t.Cleanup(func() { _ = pc.Close(); _ = tl.Close() })	69		t.Cleanup(func() { _ = pc.Close(); _ = tl.Close() })
55		go func() {	70		go func() {
56		buf := make([]byte, 4096)	71		buf := make([]byte, 4096)

@@ -92,8 +107,17 @@ func fakeUpstream(t *testing.T, handle func(q []byte) []byte) string {
92		}	107		}
93			108
94		func start(t testing.T, ups ...string) (Forwarder, string) {	109		func start(t testing.T, ups ...string) (Forwarder, string) {
			110	+	t.Helper()
			111	+	return startWith(t, nil, ups...)
			112	+	}
			113	+
			114	+	// startWith is start with a hook to adjust the Forwarder before Listen.
			115	+	func startWith(t testing.T, cfg func(Forwarder), ups ...string) (*Forwarder, string) {
95		t.Helper()	116		t.Helper()
96		f := New(func() []string { return ups }, 0)	117		f := New(func() []string { return ups }, 0)
			118	+	if cfg != nil {
			119	+	cfg(f)
			120	+	}
97		ip := netip.MustParseAddr("127.0.0.1")	121		ip := netip.MustParseAddr("127.0.0.1")
98		if err := f.Listen(context.Background(), ip); err != nil {	122		if err := f.Listen(context.Background(), ip); err != nil {
99		t.Fatal(err)	123		t.Fatal(err)

@@ -180,29 +204,6 @@ func TestServfailWhenNoUpstreamAnswers(t *testing.T) {
180		}	204		}
181		}	205		}
182			206
183	-	func TestMismatchedUpstreamIDIgnored(t *testing.T) {
184	-	// Upstream answers with a wrong ID: treated as no answer, client gets SERVFAIL.
185	-	up := fakeUpstream(t, func(q []byte) []byte {
186	-	r := answer(q)
187	-	r[0] ^= 0xff
188	-	return r
189	-	})
190	-	f, addr := start(t, up)
191	-	_ = f
192	-	c, _ := net.Dial("udp", addr)
193	-	defer func() { _ = c.Close() }()
194	-	_ = c.SetDeadline(time.Now().Add(2UpstreamTimeout + 2time.Second))
195	-	_, _ = c.Write(query(5, "a.test"))
196	-	buf := make([]byte, 512)
197	-	n, err := c.Read(buf)
198	-	if err != nil {
199	-	t.Fatal(err)
200	-	}
201	-	if buf[3]&0x0f != 2 {
202	-	t.Fatalf("want SERVFAIL, got %x", buf[:n])
203	-	}
204	-	}
205	-
206		func TestShortPacketsIgnored(t *testing.T) {	207		func TestShortPacketsIgnored(t *testing.T) {
207		up := fakeUpstream(t, answer)	208		up := fakeUpstream(t, answer)
208		_, addr := start(t, up)	209		_, addr := start(t, up)

@@ -612,3 +613,345 @@ func TestOverrideIgnoresMalformedAndNonQueries(t *testing.T) {
612		}	613		}
613		}	614		}
614		}	615		}
			616	+
			617	+	// bigAnswer returns an answer to q padded to size bytes, with TC set if tc.
			618	+	func bigAnswer(q []byte, size int, tc bool) []byte {
			619	+	r := answer(q)
			620	+	if tc {
			621	+	r[2] \|= 0x02
			622	+	}
			623	+	for len(r) < size {
			624	+	r = append(r, 'x')
			625	+	}
			626	+	return r
			627	+	}
			628	+
			629	+	func TestLargeAnswerNonEDNSClientGetsTC(t *testing.T) {
			630	+	up := fakeUpstream(t, func(q []byte) []byte { return bigAnswer(q, 700, false) })
			631	+	_, addr := start(t, up)
			632	+	q := queryType(0x3333, "storage.googleapis.com", 1, false)
			633	+	resp := udpExchange(t, addr, q)
			634	+	if resp[2]&0x02 == 0 {
			635	+	t.Fatalf("TC not set: %x", resp)
			636	+	}
			637	+	if len(resp) != len(q) \|\| binary.BigEndian.Uint16(resp[4:]) != 1 \|\| binary.BigEndian.Uint16(resp[6:]) != 0 {
			638	+	t.Fatalf("want header+question only, got %x", resp)
			639	+	}
			640	+	if resp[0] != q[0] \|\| resp[1] != q[1] {
			641	+	t.Fatal("ID changed")
			642	+	}
			643	+	}
			644	+
			645	+	func TestLargeAnswerEDNSClientGetsFullAnswer(t *testing.T) {
			646	+	up := fakeUpstream(t, func(q []byte) []byte { return bigAnswer(q, 700, false) })
			647	+	_, addr := start(t, up)
			648	+	// queryType advertises 0x1000 (4096) bytes.
			649	+	resp := udpExchange(t, addr, queryType(0x4444, "storage.googleapis.com", 1, true))
			650	+	if len(resp) != 700 \|\| resp[2]&0x02 != 0 {
			651	+	t.Fatalf("want the full 700-byte answer, got len %d tc=%d", len(resp), resp[2]&0x02)
			652	+	}
			653	+	}
			654	+
			655	+	// withOPT appends an OPT record (root name, type 41) with the given class and
			656	+	// RDLEN to a query and bumps ARCOUNT.
			657	+	func withOPT(q []byte, size, rdlen uint16, body int) []byte {
			658	+	q = append([]byte(nil), q...)
			659	+	q[11] = 1
			660	+	q = append(q, 0, 0, 41)
			661	+	q = binary.BigEndian.AppendUint16(q, size)
			662	+	q = append(q, 0, 0, 0, 0)
			663	+	q = binary.BigEndian.AppendUint16(q, rdlen)
			664	+	return append(q, make([]byte, body)...)
			665	+	}
			666	+
			667	+	func TestClientUDPSize(t *testing.T) {
			668	+	plain := query(1, "a.b")
			669	+	bigQD := append([]byte(nil), plain[:hdrLen]...)
			670	+	bigQD[4], bigQD[5] = 0xff, 0xff
			671	+	bigQD = append(bigQD, 1, 'a', 0, 0, 1, 0, 1)
			672	+	ptrQ := append([]byte(nil), plain[:hdrLen]...)
			673	+	ptrQ = append(ptrQ, 0xc0, hdrLen, 0, 1, 0, 1) // compression pointer in question
			674	+	tests := []struct {
			675	+	name string
			676	+	q []byte
			677	+	want int
			678	+	}{
			679	+	{"no edns", plain, 512},
			680	+	{"edns 1232", withOPT(plain, 1232, 0, 0), 1232},
			681	+	{"edns below 512", withOPT(plain, 100, 0, 0), 512},
			682	+	{"edns above cap", withOPT(plain, 65535, 0, 0), maxUDP},
			683	+	{"edns 4096", withOPT(plain, 4096, 0, 0), 4096},
			684	+	{"qdcount 65535 short body", bigQD, 512},
			685	+	{"compression pointer in question", ptrQ, 512},
			686	+	{"OPT rdlen overruns buffer", withOPT(plain, 4096, 500, 3), 4096},
			687	+	{"truncated header", plain[:15], 512},
			688	+	{"short", plain[:5], 512},
			689	+	}
			690	+	for _, tt := range tests {
			691	+	t.Run(tt.name, func(t *testing.T) {
			692	+	if got := clientUDPSize(tt.q); got != tt.want {
			693	+	t.Fatalf("got %d, want %d", got, tt.want)
			694	+	}
			695	+	})
			696	+	}
			697	+	}
			698	+
			699	+	func TestTruncated(t *testing.T) {
			700	+	q := queryType(7, "a.b", 1, false)
			701	+	full := bigAnswer(q, 700, false)
			702	+	full[7] = 5 // ANCOUNT
			703	+	got := truncated(full)
			704	+	if len(got) != len(q) \|\| got[2]&0x02 == 0 \|\| got[6] != 0 \|\| got[7] != 0 \|\| got[4] != 0 \|\| got[5] != 1 {
			705	+	t.Fatalf("bad truncated reply %x", got)
			706	+	}
			707	+	bad := append([]byte(nil), full...)
			708	+	bad[4], bad[5] = 0xff, 0xff // QDCOUNT lies, question parse fails
			709	+	got = truncated(bad)
			710	+	if len(got) != hdrLen \|\| got[2]&0x02 == 0 \|\| got[4] != 0 \|\| got[5] != 0 {
			711	+	t.Fatalf("bad-question case: %x", got)
			712	+	}
			713	+	if truncated(full[:5]) != nil {
			714	+	t.Fatal("short input must give nil")
			715	+	}
			716	+	}
			717	+
			718	+	func TestServfailAndOverrideRepliesAreNotTruncated(t *testing.T) {
			719	+	dead := "127.0.0.1:1" // nothing listens: UDP and TCP both fail at once
			720	+	_, addr := startOverride(t, "forge.test", dead)
			721	+	resp := udpExchange(t, addr, queryType(9, "x.test", 1, false))
			722	+	if len(resp) != hdrLen \|\| resp[3]&0x0f != 2 \|\| resp[2]&0x02 != 0 {
			723	+	t.Fatalf("want a plain 12-byte SERVFAIL, got %x", resp)
			724	+	}
			725	+	resp = udpExchange(t, addr, queryType(10, "forge.test", 1, false))
			726	+	if resp[2]&0x02 != 0 \|\| resp[3]&0x0f != 0 \|\| binary.BigEndian.Uint16(resp[6:]) != 1 {
			727	+	t.Fatalf("want the override A answer, got %x", resp)
			728	+	}
			729	+	}
			730	+
			731	+	// fakeNet is an in-memory upstream network for Forwarder.dial. One Write on a
			732	+	// UDP conn is one datagram (net.Pipe delivers a Write to a single Read when the
			733	+	// reader's buffer is large enough). A handler returning nil never answers.
			734	+	type fakeNet struct {
			735	+	udp, tcp func(addr string, q []byte) []byte
			736	+	udpCalls, tcpCalls atomic.Int64
			737	+	}
			738	+
			739	+	func (n *fakeNet) dial(ctx context.Context, network, addr string) (net.Conn, error) {
			740	+	if err := ctx.Err(); err != nil { // like net.Dialer
			741	+	return nil, err
			742	+	}
			743	+	c, s := net.Pipe()
			744	+	switch network {
			745	+	case "udp":
			746	+	go func() {
			747	+	defer func() { _ = s.Close() }()
			748	+	buf := make([]byte, 4096)
			749	+	for {
			750	+	k, err := s.Read(buf)
			751	+	if err != nil {
			752	+	return
			753	+	}
			754	+	n.udpCalls.Add(1)
			755	+	if r := n.udp(addr, append([]byte(nil), buf[:k]...)); r != nil {
			756	+	if _, err := s.Write(r); err != nil {
			757	+	return
			758	+	}
			759	+	}
			760	+	}
			761	+	}()
			762	+	case "tcp":
			763	+	go func() {
			764	+	defer func() { _ = s.Close() }()
			765	+	var l [2]byte
			766	+	if _, err := io.ReadFull(s, l[:]); err != nil {
			767	+	return
			768	+	}
			769	+	q := make([]byte, binary.BigEndian.Uint16(l[:]))
			770	+	if _, err := io.ReadFull(s, q); err != nil {
			771	+	return
			772	+	}
			773	+	n.tcpCalls.Add(1)
			774	+	if r := n.tcp(addr, q); r != nil {
			775	+	_, _ = s.Write(append(binary.BigEndian.AppendUint16(nil, uint16(len(r))), r...))
			776	+	}
			777	+	_, _ = io.Copy(io.Discard, s) // wait for the client to hang up
			778	+	}()
			779	+	default:
			780	+	return nil, errors.New("unexpected network " + network)
			781	+	}
			782	+	return c, nil
			783	+	}
			784	+
			785	+	// silent is a handler that never answers.
			786	+	func silent(string, []byte) []byte { return nil }
			787	+
			788	+	// fakeForwarder returns a Forwarder with production timeouts whose upstreams
			789	+	// are served by n, and a listener for exchange.
			790	+	func fakeForwarder(n fakeNet, ups ...string) (Forwarder, *listener) {
			791	+	f := New(func() []string { return ups }, 0)
			792	+	f.dial = n.dial
			793	+	return f, &listener{addr: netip.MustParseAddr("127.0.0.1")}
			794	+	}
			795	+
			796	+	func isServfail(r []byte, id uint16) bool {
			797	+	return len(r) == hdrLen && binary.BigEndian.Uint16(r) == id && r[3]&0x0f == 2
			798	+	}
			799	+
			800	+	// The tests below run in a synctest bubble: the fake clock makes the real
			801	+	// production timeouts (1 s UDP, 3 s TCP, 6 s cap) instant and exact.
			802	+
			803	+	func TestExchangeSilentUDPFallsBackToTCPAfterUDPTimeout(t *testing.T) {
			804	+	synctest.Test(t, func(t *testing.T) {
			805	+	n := &fakeNet{udp: silent, tcp: func(_ string, q []byte) []byte { return answer(q) }}
			806	+	f, l := fakeForwarder(n, "up:53")
			807	+	start := time.Now()
			808	+	resp := f.exchange(t.Context(), l, query(1, "storage.googleapis.com"), false)
			809	+	if !strings.HasSuffix(string(resp), "ANSWER") {
			810	+	t.Fatalf("want the TCP answer, got %x", resp)
			811	+	}
			812	+	if got := time.Since(start); got != UDPAttemptTimeout {
			813	+	t.Fatalf("answered after %v, want %v", got, UDPAttemptTimeout)
			814	+	}
			815	+	if n.udpCalls.Load() != 1 \|\| n.tcpCalls.Load() != 1 {
			816	+	t.Fatalf("udp=%d tcp=%d, want 1 and 1", n.udpCalls.Load(), n.tcpCalls.Load())
			817	+	}
			818	+	})
			819	+	}
			820	+
			821	+	func TestExchangeTruncatedUDPRetriedOverTCP(t *testing.T) {
			822	+	synctest.Test(t, func(t *testing.T) {
			823	+	n := &fakeNet{
			824	+	udp: func(_ string, q []byte) []byte { return bigAnswer(q, 0, true) },
			825	+	tcp: func(_ string, q []byte) []byte { return bigAnswer(q, 700, false) },
			826	+	}
			827	+	f, l := fakeForwarder(n, "up:53")
			828	+	start := time.Now()
			829	+	resp := f.exchange(t.Context(), l, query(2, "example.com"), false)
			830	+	if len(resp) != 700 \|\| resp[2]&0x02 != 0 {
			831	+	t.Fatalf("want the full TCP answer, got len %d tc=%d", len(resp), resp[2]&0x02)
			832	+	}
			833	+	if time.Since(start) != 0 {
			834	+	t.Fatalf("took %v, want no waiting", time.Since(start))
			835	+	}
			836	+	})
			837	+	}
			838	+
			839	+	func TestExchangeKeepsTruncatedReplyWhenTCPFails(t *testing.T) {
			840	+	synctest.Test(t, func(t *testing.T) {
			841	+	n := &fakeNet{udp: func(_ string, q []byte) []byte { return bigAnswer(q, 0, true) }, tcp: silent}
			842	+	f, l := fakeForwarder(n, "up:53")
			843	+	resp := f.exchange(t.Context(), l, query(3, "example.com"), false)
			844	+	if len(resp) < hdrLen \|\| resp[2]&0x02 == 0 {
			845	+	t.Fatalf("want the TC=1 reply, got %x", resp)
			846	+	}
			847	+	})
			848	+	}
			849	+
			850	+	func TestExchangeSecondUpstreamNotStarved(t *testing.T) {
			851	+	synctest.Test(t, func(t *testing.T) {
			852	+	n := &fakeNet{
			853	+	udp: func(a string, q []byte) []byte {
			854	+	if a == "dead:53" {
			855	+	return nil
			856	+	}
			857	+	return answer(q)
			858	+	},
			859	+	tcp: func(a string, q []byte) []byte {
			860	+	if a == "dead:53" {
			861	+	return nil
			862	+	}
			863	+	return answer(q)
			864	+	},
			865	+	}
			866	+	f, l := fakeForwarder(n, "dead:53", "live:53")
			867	+	start := time.Now()
			868	+	resp := f.exchange(t.Context(), l, query(4, "a.test"), false)
			869	+	if !strings.HasSuffix(string(resp), "ANSWER") {
			870	+	t.Fatalf("want the second upstream's answer, got %x", resp)
			871	+	}
			872	+	if got, want := time.Since(start), UDPAttemptTimeout+UpstreamTimeout; got != want {
			873	+	t.Fatalf("took %v, want %v", got, want)
			874	+	}
			875	+	})
			876	+	}
			877	+
			878	+	func TestExchangeMismatchedUpstreamIDIgnored(t *testing.T) {
			879	+	synctest.Test(t, func(t *testing.T) {
			880	+	wrong := func(_ string, q []byte) []byte {
			881	+	r := answer(q)
			882	+	r[0] ^= 0xff
			883	+	return r
			884	+	}
			885	+	f, l := fakeForwarder(&fakeNet{udp: wrong, tcp: wrong}, "up:53")
			886	+	start := time.Now()
			887	+	resp := f.exchange(t.Context(), l, query(5, "a.test"), false)
			888	+	if !isServfail(resp, 5) {
			889	+	t.Fatalf("want SERVFAIL, got %x", resp)
			890	+	}
			891	+	// UDP keeps listening for the right ID until its timeout; the TCP
			892	+	// reply is read once and rejected.
			893	+	if got := time.Since(start); got != UDPAttemptTimeout {
			894	+	t.Fatalf("took %v, want %v", got, UDPAttemptTimeout)
			895	+	}
			896	+	})
			897	+	}
			898	+
			899	+	func TestExchangeCappedAtMaxExchange(t *testing.T) {
			900	+	synctest.Test(t, func(t *testing.T) {
			901	+	n := &fakeNet{udp: silent, tcp: silent}
			902	+	f, l := fakeForwarder(n, "a:53", "b:53", "c:53")
			903	+	start := time.Now()
			904	+	resp := f.exchange(t.Context(), l, query(6, "a.test"), false)
			905	+	if !isServfail(resp, 6) {
			906	+	t.Fatalf("want SERVFAIL, got %x", resp)
			907	+	}
			908	+	if got := time.Since(start); got != MaxExchange {
			909	+	t.Fatalf("took %v, want %v", got, MaxExchange)
			910	+	}
			911	+	// a: 1 s UDP + 3 s TCP; b: 1 s UDP + 1 s TCP (cut by the cap); c: skipped.
			912	+	if n.udpCalls.Load() != 2 \|\| n.tcpCalls.Load() != 2 {
			913	+	t.Fatalf("udp=%d tcp=%d, want 2 and 2", n.udpCalls.Load(), n.tcpCalls.Load())
			914	+	}
			915	+	})
			916	+	}
			917	+
			918	+	func TestExchangeTCPClientUsesTCPOnly(t *testing.T) {
			919	+	synctest.Test(t, func(t *testing.T) {
			920	+	n := &fakeNet{udp: silent, tcp: func(_ string, q []byte) []byte { return answer(q) }}
			921	+	f, l := fakeForwarder(n, "up:53")
			922	+	resp := f.exchange(t.Context(), l, query(7, "a.test"), true)
			923	+	if !strings.HasSuffix(string(resp), "ANSWER") \|\| n.udpCalls.Load() != 0 {
			924	+	t.Fatalf("resp=%x udp=%d", resp, n.udpCalls.Load())
			925	+	}
			926	+	})
			927	+	}
			928	+
			929	+	func TestExchangeCancelledReturnsNil(t *testing.T) {
			930	+	synctest.Test(t, func(t *testing.T) {
			931	+	f, l := fakeForwarder(&fakeNet{udp: silent, tcp: silent}, "up:53")
			932	+	ctx, cancel := context.WithCancel(t.Context())
			933	+	time.AfterFunc(500*time.Millisecond, cancel)
			934	+	start := time.Now()
			935	+	if resp := f.exchange(ctx, l, query(8, "a.test"), false); resp != nil {
			936	+	t.Fatalf("want nil on cancellation, got %x", resp)
			937	+	}
			938	+	// Cancellation is noticed at the next attempt boundary: the blocked UDP
			939	+	// read runs to its own deadline, then no TCP attempt is started.
			940	+	if got := time.Since(start); got != UDPAttemptTimeout {
			941	+	t.Fatalf("returned after %v, want %v", got, UDPAttemptTimeout)
			942	+	}
			943	+	})
			944	+	}
			945	+
			946	+	// End to end through real client-facing sockets; the upstream is in memory.
			947	+	func TestUDPClientGetsTCPAnswerWhenUpstreamUDPSilent(t *testing.T) {
			948	+	n := &fakeNet{udp: silent, tcp: func(_ string, q []byte) []byte { return answer(q) }}
			949	+	_, addr := startWith(t, func(f *Forwarder) {
			950	+	f.dial = n.dial
			951	+	f.udpTimeout = 50 * time.Millisecond
			952	+	}, "up:53")
			953	+	resp := udpExchange(t, addr, query(0x1111, "storage.googleapis.com"))
			954	+	if !strings.HasSuffix(string(resp), "ANSWER") \|\| resp[3]&0x0f != 0 {
			955	+	t.Fatalf("want the TCP answer, got %x", resp)
			956	+	}
			957	+	}