DNS cookies support

README.md

@@ -67,6 +67,8 @@ Usage of ./dnsproxy:
         If specified, dnsproxy will act as a DNS64 server.
   --dns64-prefix=subnet
         Prefix used to handle DNS64. If not specified, dnsproxy uses the 'Well-Known Prefix' 64:ff9b::.  Can be specified multiple times.
+  --disable-dns-cookies
+        Disable DNS cookies support (which is enabled by default).
   --dnscrypt-config=path/-g path
         Path to a file with DNSCrypt configuration. You can generate one using https://github.com/ameshkov/dnscrypt.
   --dnscrypt-port=port/-y port
@@ -304,6 +306,12 @@ Loads upstreams list from a file.
 ./dnsproxy -l 127.0.0.1 -p 5353 -u ./upstreams.txt
 ```
 
+DNS cookies are enabled by default. Disable them if needed:
+
+```shell
+./dnsproxy --disable-dns-cookies
+```
+
 ### DNS64 server
 
 `dnsproxy` is capable of working as a DNS64 server.

internal/cmd/args.go

@@ -24,6 +24,8 @@ const (
 	httpsUserinfoIdx
 	dnsCryptConfigPathIdx
 	ednsAddrIdx
+	disableDNSCookiesIdx
+	dnsCookieSecretIdx
 	upstreamModeIdx
 	listenAddrsIdx
 	listenPortsIdx
@@ -132,6 +134,18 @@ var commandLineOptions = []*commandLineOption{
 		short:       "",
 		valueType:   "address",
 	},
+	disableDNSCookiesIdx: {
+		description: "Disable EDNS Cookies support.",
+		long:        "disable-dns-cookies",
+		short:       "",
+		valueType:   "",
+	},
+	dnsCookieSecretIdx: {
+		description: "Hex-encoded 16-byte secret used to generate EDNS Cookies.",
+		long:        "dns-cookie-secret",
+		short:       "",
+		valueType:   "hex",
+	},
 	upstreamModeIdx: {
 		description: "Defines the upstreams logic mode, possible values: load_balance, parallel, " +
 			"fastest_addr (default: load_balance).",
@@ -421,6 +435,8 @@ func parseCmdLineOptions(conf *configuration) (err error) {
 		httpsUserinfoIdx:            &conf.HTTPSUserinfo,
 		dnsCryptConfigPathIdx:       &conf.DNSCryptConfigPath,
 		ednsAddrIdx:                 &conf.EDNSAddr,
+		disableDNSCookiesIdx:        &conf.DisableDNSCookies,
+		dnsCookieSecretIdx:          &conf.DNSCookieSecret,
 		upstreamModeIdx:             &conf.UpstreamMode,
 		listenAddrsIdx:              &conf.ListenAddrs,
 		listenPortsIdx:              &conf.ListenPorts,

internal/cmd/config.go

@@ -39,6 +39,13 @@ type configuration struct {
 	// EDNSAddr is the custom EDNS Client Address to send.
 	EDNSAddr string `yaml:"edns-addr"`
 
+	// DisableDNSCookies strips EDNS Cookies from both requests and responses.
+	DisableDNSCookies bool `yaml:"disable-dns-cookies"`
+
+	// DNSCookieSecret is a hex-encoded 16-byte secret used to generate server
+	// cookies.
+	DNSCookieSecret string `yaml:"dns-cookie-secret"`
+
 	// UpstreamMode determines the logic through which upstreams will be used.
 	// If not specified the [proxy.UpstreamModeLoadBalance] is used.
 	UpstreamMode string `yaml:"upstream-mode"`

internal/cmd/proxy.go

@@ -75,6 +75,8 @@ func createProxyConfig(
 			netip.MustParsePrefix("0.0.0.0/0"),
 			netip.MustParsePrefix("::0/0"),
 		},
+		DisableDNSCookies:      conf.DisableDNSCookies,
+		DNSCookieSecret:        conf.DNSCookieSecret,
 		EnableEDNSClientSubnet: conf.EnableEDNSSubnet,
 		UDPBufferSize:          conf.UDPBufferSize,
 		HTTPSServerName:        conf.HTTPSServerName,

proxy/config.go

@@ -175,6 +175,14 @@ type Config struct {
 	// EDNSAddr is the ECS IP used in request.
 	EDNSAddr net.IP
 
+	// DisableDNSCookies strips EDNS Cookies from both requests and responses.
+	// If false, the proxy responds with EDNS Cookies per RFC 7873 and RFC 9018.
+	DisableDNSCookies bool `yaml:"disable-dns-cookies" json:"disable_dns_cookies"`
+
+	// DNSCookieSecret is a hex-encoded 16-byte secret used to generate server
+	// cookies.  If empty, a random secret is generated on start.
+	DNSCookieSecret string `yaml:"dns-cookie-secret" json:"dns_cookie_secret"`
+
 	// TODO(s.chzhen):  Extract ratelimit settings to a separate structure.
 
 	// RatelimitSubnetLenIPv4 is a subnet length for IPv4 addresses used for

proxy/cookie.go

@@ -0,0 +1,243 @@
+package proxy
+
+import (
+	"crypto/hmac"
+	"crypto/rand"
+	"crypto/sha256"
+	"encoding/hex"
+	"fmt"
+	"net/netip"
+
+	"github.com/miekg/dns"
+)
+
+const (
+	// clientCookieLen is the length of the client cookie in bytes.
+	clientCookieLen = 8
+
+	// serverCookieLen is the length of the server cookie in bytes.
+	serverCookieLen = 16
+
+	// cookieSecretLen is the size of the secret used to generate server
+	// cookies.
+	cookieSecretLen = 16
+
+	// doBitMask is the mask of the DO bit in the TTL field of an OPT record.
+	doBitMask = uint32(1 << 15)
+)
+
+// parseCookie returns the client and server cookies from m if present.  The
+// client cookie is required to be at least eight bytes long.
+func parseCookie(m *dns.Msg) (client, server []byte) {
+	if m == nil {
+		return nil, nil
+	}
+
+	opt := m.IsEdns0()
+	if opt == nil {
+		return nil, nil
+	}
+
+	for _, o := range opt.Option {
+		c, ok := o.(*dns.EDNS0_COOKIE)
+		if !ok {
+			continue
+		}
+
+		raw, err := hex.DecodeString(c.Cookie)
+		if err != nil || len(raw) < clientCookieLen {
+			continue
+		}
+
+		client = raw[:clientCookieLen]
+		if len(raw) > clientCookieLen {
+			server = raw[clientCookieLen:]
+		}
+
+		return client, server
+	}
+
+	return nil, nil
+}
+
+// stripCookie removes any EDNS cookie options from m.
+func stripCookie(m *dns.Msg) {
+	if m == nil {
+		return
+	}
+
+	opt := m.IsEdns0()
+	if opt == nil {
+		return
+	}
+
+	opts := opt.Option[:0]
+	for _, o := range opt.Option {
+		if o.Option() == dns.EDNS0COOKIE {
+			continue
+		}
+
+		opts = append(opts, o)
+	}
+
+	opt.Option = opts
+}
+
+// handleRequestCookies parses client cookies and strips cookie options before
+// forwarding requests upstream.
+func (p *Proxy) handleRequestCookies(dctx *DNSContext) {
+	if dctx == nil || dctx.Req == nil {
+		return
+	}
+
+	if !p.DisableDNSCookies {
+		if client, _ := parseCookie(dctx.Req); len(client) > 0 {
+			dctx.reqClientCookie = client
+		}
+	}
+
+	stripCookie(dctx.Req)
+}
+
+// handleResponseCookies sets the response cookie if needed or strips cookie
+// options when cookies are disabled.
+func (p *Proxy) handleResponseCookies(dctx *DNSContext) {
+	if dctx == nil || dctx.Res == nil {
+		return
+	}
+
+	if p.DisableDNSCookies {
+		stripCookie(dctx.Res)
+
+		return
+	}
+
+	if len(dctx.reqClientCookie) == 0 {
+		stripCookie(dctx.Res)
+
+		return
+	}
+
+	udpSize := dctx.udpSize
+	if udpSize == 0 {
+		udpSize = defaultUDPBufSize
+	}
+
+	stripCookie(dctx.Res)
+
+	server := p.serverCookie(dctx.Addr.Addr(), dctx.reqClientCookie)
+	if len(server) == 0 {
+		return
+	}
+
+	setCookie(dctx.Res, dctx.reqClientCookie, server, udpSize, dctx.doBit)
+}
+
+// serverCookie returns the server cookie for the provided address and client
+// cookie using an HMAC-SHA256.
+func (p *Proxy) serverCookie(ip netip.Addr, client []byte) (server []byte) {
+	if len(client) < clientCookieLen || !ip.IsValid() {
+		return nil
+	}
+
+	p.cookieMu.Lock()
+	secret := p.cookieSecret
+	p.cookieMu.Unlock()
+
+	if len(secret) != cookieSecretLen {
+		return nil
+	}
+
+	h := hmac.New(sha256.New, secret)
+	_, _ = h.Write(client)
+	_, _ = h.Write(ip.AsSlice())
+
+	// Truncate to sixteen bytes.
+	return h.Sum(nil)[:serverCookieLen]
+}
+
+// setCookie ensures msg has the OPT record and sets the cookie option to
+// client+server.  udpSize is the UDP buffer size advertised to the client.
+func setCookie(msg *dns.Msg, client, server []byte, udpSize uint16, do bool) {
+	if msg == nil || len(client) == 0 || len(server) == 0 {
+		return
+	}
+
+	opt := msg.IsEdns0()
+	if opt == nil {
+		msg.SetEdns0(udpSize, do)
+		opt = msg.IsEdns0()
+	} else {
+		opt.SetUDPSize(udpSize)
+		if do {
+			opt.SetDo()
+		} else {
+			opt.Hdr.Ttl &^= doBitMask
+		}
+	}
+
+	stripCookie(msg)
+
+	cookie := make([]byte, 0, len(client)+len(server))
+	cookie = append(cookie, client...)
+	cookie = append(cookie, server...)
+
+	opt = msg.IsEdns0()
+	opt.Option = append(opt.Option, &dns.EDNS0_COOKIE{
+		Code:   dns.EDNS0COOKIE,
+		Cookie: hex.EncodeToString(cookie),
+	})
+}
+
+// decodeCookieSecret decodes hexSecret and validates its length.
+func decodeCookieSecret(hexSecret string) (secret []byte, err error) {
+	secret, err = hex.DecodeString(hexSecret)
+	if err != nil {
+		return nil, fmt.Errorf("decoding dns cookie secret: %w", err)
+	}
+
+	if len(secret) != cookieSecretLen {
+		return nil, fmt.Errorf(
+			"decoding dns cookie secret: invalid length %d, want %d",
+			len(secret),
+			cookieSecretLen,
+		)
+	}
+
+	return secret, nil
+}
+
+// generateCookieSecret returns a new random secret suitable for cookie
+// generation.
+func generateCookieSecret() (secret []byte, err error) {
+	secret = make([]byte, cookieSecretLen)
+	_, err = rand.Read(secret)
+	if err != nil {
+		return nil, fmt.Errorf("reading random secret: %w", err)
+	}
+
+	return secret, nil
+}
+
+// initCookieSecret prepares the cookie secret depending on the configuration.
+func (p *Proxy) initCookieSecret() (err error) {
+	if p.DisableDNSCookies {
+		return nil
+	}
+
+	var secret []byte
+	if p.DNSCookieSecret != "" {
+		secret, err = decodeCookieSecret(p.DNSCookieSecret)
+	} else {
+		secret, err = generateCookieSecret()
+	}
+	if err != nil {
+		return fmt.Errorf("creating dns cookie secret: %w", err)
+	}
+
+	p.cookieMu.Lock()
+	p.cookieSecret = secret
+	p.cookieMu.Unlock()
+
+	return nil
+}