feat: set DNS cookie manually (#61)

natesales · Dec 13, 2023 · 9e9a6be · 9e9a6be
1 parent 0a4790c
commit 9e9a6be
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 23 deletions.
diff --git a/README.md b/README.md
@@ -36,33 +36,28 @@ All long form (--) flags can be toggled with the dig-standard +[no]flag notation
 Application Options:
   -q, --qname=                    Query name
   -s, --server=                   DNS server(s)
-  -t, --type=                     RR type (e.g. A, AAAA, MX, etc.) or type
-                                  integer
+  -t, --type=                     RR type (e.g. A, AAAA, MX, etc.) or type integer
   -x, --reverse                   Reverse lookup
   -d, --dnssec                    Set the DO (DNSSEC OK) bit in the OPT record
   -n, --nsid                      Set EDNS0 NSID opt
       --subnet=                   Set EDNS0 client subnet
   -c, --chaos                     Use CHAOS query class
-  -C=                             Set query class (default: IN 0x01) (default:
-                                  1)
+  -C=                             Set query class (default: IN 0x01) (default: 1)
   -p, --odoh-proxy=               ODoH proxy
       --timeout=                  Query timeout (default: 10s)
       --pad                       Set EDNS0 padding
       --http3                     Use HTTP/3 for DoH
       --id-check                  Check DNS response ID (default: true)
-      --reuse-conn                Reuse connections across queries to the same
-                                  server (default: true)
+      --reuse-conn                Reuse connections across queries to the same server (default: true)
       --txtconcat                 Concatenate TXT responses
       --qid=                      Set query ID (-1 for random) (default: -1)
   -b, --bootstrap-server=         DNS server to use for bootstrapping
       --bootstrap-timeout=        Bootstrapping timeout (default: 5s)
+      --cookie=                   Set EDNS0 cookie
       --recaxfr                   Perform recursive AXFR
-  -f, --format=                   Output format (pretty, column, json, yaml,
-                                  raw) (default: pretty)
-      --pretty-ttls               Format TTLs in human readable format
-                                  (default: true)
-      --short-ttls                Remove zero components of pretty TTLs.
-                                  (24h0m0s->24h) (default: true)
+  -f, --format=                   Output format (pretty, column, json, yaml, raw) (default: pretty)
+      --pretty-ttls               Format TTLs in human readable format (default: true)
+      --short-ttls                Remove zero components of pretty TTLs. (24h0m0s->24h) (default: true)
       --color                     Enable color output
       --question                  Show question section
       --answer                    Show answer section (default: true)
@@ -72,13 +67,11 @@ Application Options:
       --all                       Show all sections and statistics
   -w                              Resolve ASN/ASName for A and AAAA records
   -r, --short                     Show record values only
-  -R, --resolve-ips               Resolve PTR records for IP addresses in A and
-                                  AAAA records
+  -R, --resolve-ips               Resolve PTR records for IP addresses in A and AAAA records
       --aa                        Set AA (Authoritative Answer) flag in query
       --ad                        Set AD (Authentic Data) flag in query
       --cd                        Set CD (Checking Disabled) flag in query
-      --rd                        Set RD (Recursion Desired) flag in query
-                                  (default: true)
+      --rd                        Set RD (Recursion Desired) flag in query (default: true)
       --ra                        Set RA (Recursion Available) flag in query
       --z                         Set Z (Zero) flag in query
       --t                         Set TC (Truncated) flag in query
@@ -96,15 +89,12 @@ Application Options:
       --http-method=              HTTP method (default: GET)
       --pmtud                     PMTU discovery (default: true)
       --quic-alpn-tokens=         QUIC ALPN tokens (default: doq, doq-i11)
-      --quic-length-prefix        Add RFC 9250 compliant length prefix
-                                  (default: true)
+      --quic-length-prefix        Add RFC 9250 compliant length prefix (default: true)
       --dnscrypt-tcp              Use TCP for DNSCrypt (default UDP)
-      --dnscrypt-udp-size=        Maximum size of a DNS response this client
-                                  can sent or receive (default: 0)
+      --dnscrypt-udp-size=        Maximum size of a DNS response this client can sent or receive (default: 0)
       --dnscrypt-key=             DNSCrypt public key
       --dnscrypt-provider=        DNSCrypt provider name
-      --default-rr-types=         Default record types (default: A, AAAA, NS,
-                                  MX, TXT, CNAME)
+      --default-rr-types=         Default record types (default: A, AAAA, NS, MX, TXT, CNAME)
       --udp-buffer=               Set EDNS0 UDP size in query (default: 1232)
   -v, --verbose                   Show verbose log messages
       --trace                     Show trace log messages

diff --git a/cli/cli.go b/cli/cli.go
@@ -31,7 +31,9 @@ type Flags struct {
 	ID               int           `long:"qid" description:"Set query ID (-1 for random)" default:"-1"`
 	BootstrapServer  string        `short:"b" long:"bootstrap-server" description:"DNS server to use for bootstrapping"`
 	BootstrapTimeout time.Duration `long:"bootstrap-timeout" description:"Bootstrapping timeout" default:"5s"`
+	Cookie           string        `long:"cookie" description:"EDNS0 cookie"`
 
+	// Special query modes
 	RecAXFR bool `long:"recaxfr" description:"Perform recursive AXFR"`
 
 	// Output

diff --git a/resolver.go b/resolver.go
@@ -34,7 +34,7 @@ func createQuery(opts cli.Flags, rrTypes []uint16) []dns.Msg {
 		req.Zero = opts.Zero
 		req.Truncated = opts.Truncated
 
-		if opts.DNSSEC || opts.NSID || opts.Pad || opts.ClientSubnet != "" {
+		if opts.DNSSEC || opts.NSID || opts.Pad || opts.ClientSubnet != "" || opts.Cookie != "" {
 			opt := &dns.OPT{
 				Hdr: dns.RR_Header{
 					Name:   ".",
@@ -92,6 +92,15 @@ func createQuery(opts cli.Flags, rrTypes []uint16) []dns.Msg {
 				}
 				opt.Option = append(opt.Option, ednsSubnet)
 			}
+
+			if opts.Cookie != "" {
+				cookie := &dns.EDNS0_COOKIE{
+					Code:   dns.EDNS0COOKIE,
+					Cookie: opts.Cookie,
+				}
+				opt.Option = append(opt.Option, cookie)
+			}
+
 			req.Extra = append(req.Extra, opt)
 		}