feat(dig): implement +udp and +udp=wait-duplicates (#31)

This commit adds support for explicitly requesting the DNS-over-UDP
protocol, as well as for waiting for duplicate responses. Waiting for
duplicate responses allows to identify censorship in places such as
China and Iran.

internal/qa/registry.go

@@ -45,6 +45,27 @@ var Registry = []ScenarioDescriptor{
 		},
 	},
 
+	{
+		Name: "dnsOverUdpCensorshipWithDuplicates",
+		Editors: []ScenarioEditor{
+			CensorDNSLikeIran("www.example.com"),
+		},
+		Argv: []string{
+			"rbmk", "dig", "+udp=wait-duplicates", "+noall", "+logs", "@8.8.8.8", "A", "www.example.com",
+		},
+		ExpectedErr: nil,
+		ExpectedSeq: []ExpectedEvent{
+			{Msg: "connectStart"},
+			{Msg: "connectDone"},
+			{Msg: "dnsQuery"},
+			{Pattern: MatchAnyRead | MatchAnyWrite},
+			{Msg: "dnsResponse"},
+			{Pattern: MatchAnyRead | MatchAnyWrite},
+			{Msg: "dnsResponse"},
+			{Pattern: MatchAnyClose},
+		},
+	},
+
 	//
 	// DNS over TCP
 	//

pkg/cli/dig/README.md

@@ -119,6 +119,18 @@ address to use. The implied port is `53/tcp`.
 Uses DNS-over-TLS. The @server argument is the hostname or IP
 address to use. The implied port is `853/tcp`.
 
+### `+udp`
+
+Use DNS-over-UDP (default behavior).
+
+### `+udp=wait-duplicates`
+
+Use DNS-over-UDP and wait for the full query timeout to collect
+duplicate responses. Only the first (i.e., non-duplicate) response
+is printed to the stdout. All responses (including duplicates)
+are included in the structured logs. This option is useful
+for detecting DNS-based censorship in China and Iran.
+
 ## Examples
 
 The following invocation resolves `www.example.com` IPv6 address

pkg/cli/dig/dig.go

@@ -56,6 +56,7 @@ func (cmd command) Main(ctx context.Context, env cliutils.Environment, argv ...s
 		ServerAddr:     "8.8.8.8",
 		ServerPort:     "53",
 		URLPath:        "/dns-query",
+		WaitDuplicates: false,
 	}
 
 	// 3. create command line parser
@@ -105,6 +106,7 @@ func (cmd command) Main(ctx context.Context, env cliutils.Environment, argv ...s
 			case arg == "+https":
 				task.Protocol = "doh"
 				task.ServerPort = "443"
+				task.WaitDuplicates = false
 				continue
 
 			case arg == "+logs":
@@ -131,11 +133,19 @@ func (cmd command) Main(ctx context.Context, env cliutils.Environment, argv ...s
 			case arg == "+tcp":
 				task.Protocol = "tcp"
 				task.ServerPort = "53"
+				task.WaitDuplicates = false
 				continue
 
 			case arg == "+tls":
 				task.Protocol = "dot"
 				task.ServerPort = "853"
+				task.WaitDuplicates = false
+				continue
+
+			case arg == "+udp" || arg == "+udp=wait-duplicates":
+				task.Protocol = "udp"
+				task.ServerPort = "53"
+				task.WaitDuplicates = arg == "+udp=wait-duplicates"
 				continue
 
 			default:

pkg/cli/dig/task.go

@@ -4,13 +4,15 @@ package dig
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
 	"net"
 	"net/http"
 	"net/url"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/miekg/dns"
@@ -68,6 +70,11 @@ type Task struct {
 
 	// URLPath is the MANDATORY URL path when using DoH.
 	URLPath string
+
+	// WaitDuplicates is the OPTIONAL flag indicating
+	// whether we should wait for duplicate DNS-over-UDP
+	// responses (for detecting censorship).
+	WaitDuplicates bool
 }
 
 // queryTypeMap maps query types strings to DNS query types.
@@ -179,7 +186,7 @@ func (task *Task) Run(ctx context.Context) error {
 	fmt.Fprintf(task.QueryWriter, ";; Query:\n%s\n", query.String())
 
 	// Perform the DNS query
-	response, err := transport.Query(ctx, server, query)
+	response, err := task.query(ctx, transport, server, query)
 	if err != nil {
 		return fmt.Errorf("query round-trip failed: %w", err)
 	}
@@ -201,6 +208,46 @@ func (task *Task) Run(ctx context.Context) error {
 	return nil
 }
 
+// query performs the query and returns response or error.
+//
+// If the WaitDuplicates flag is set, this function will wait
+// for duplicate responses, emit all the related structured logs,
+// and return the first response received. This function blocks
+// until the timeout configured in the context expires. Note that
+// all responses (including duplicates) are automatically
+// logged through the transport's logger.
+func (task *Task) query(
+	ctx context.Context,
+	txp *dnscore.Transport,
+	addr *dnscore.ServerAddr,
+	query *dns.Msg,
+) (*dns.Msg, error) {
+	// If we're not waiting for duplicates, our job is easy
+	if !task.WaitDuplicates {
+		return txp.Query(ctx, addr, query)
+	}
+
+	// Otherwise, we need to reading duplicate responses
+	// until the overall timeout says we should bail, which
+	// happens through context expiration.
+	var (
+		resp0 *dns.Msg
+		err0  error
+		once  sync.Once
+	)
+	respch := txp.QueryWithDuplicates(ctx, addr, query)
+	for entry := range respch {
+		resp, err := entry.Msg, entry.Err
+		once.Do(func() {
+			resp0, err0 = resp, err
+		})
+	}
+	if resp0 == nil && err0 == nil {
+		return nil, errors.New("received nil response and nil error")
+	}
+	return resp0, err0
+}
+
 // formatShort returns a short string representation of the DNS response.
 func (task *Task) formatShort(response *dns.Msg) string {
 	var builder strings.Builder