Pull request #434: 8384-fix-do-cache

Merge in GO/dnsproxy from 8384-fix-do-cache to master

Updates AdguardTeam/AdGuardHome#8384.

Squashed commit of the following:

commit e36ac40
Merge: 29d546b bd64c3a
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Thu May 21 18:40:11 2026 +0300

    Merge branch 'master' into 8384-fix-do-cache

commit 29d546b
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Thu May 21 17:09:12 2026 +0300

    proxy: add do bit to cache key

commit e6245e9
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Wed May 20 22:17:17 2026 +0300

    proxy: imp doc

commit 58b5643
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Wed May 20 21:04:01 2026 +0300

    proxy: fix caching without do

Author: EugeneOne1

proxy/cache.go

@@ -326,14 +326,15 @@ func createCache(cacheSize int) (glc glcache.Cache) {
 	return glcache.New(conf)
 }
 
-// set stores response and upstream in the cache.  l must not be nil.
-func (c *cache) set(m *dns.Msg, u upstream.Upstream, l *slog.Logger) {
+// set stores response and upstream for req in the cache.  u and l must not be
+// nil.
+func (c *cache) set(req, m *dns.Msg, u upstream.Upstream, l *slog.Logger) {
 	item := c.respToItem(m, u, l)
 	if item == nil {
 		return
 	}
 
-	key := msgToKey(m)
+	key := msgToKey(req)
 	packed := item.pack()
 
 	c.itemsLock.Lock()
@@ -343,16 +344,16 @@ func (c *cache) set(m *dns.Msg, u upstream.Upstream, l *slog.Logger) {
 }
 
 // setWithSubnet stores response and upstream with subnet in the cache.  The
-// given subnet mask and IP address are used to calculate the cache key.  l must
-// not be nil.
-func (c *cache) setWithSubnet(m *dns.Msg, u upstream.Upstream, subnet *net.IPNet, l *slog.Logger) {
+// given subnet mask and IP address are used to calculate the cache key.  u, n,
+// and l must not be nil.
+func (c *cache) setWithSubnet(req, m *dns.Msg, u upstream.Upstream, n *net.IPNet, l *slog.Logger) {
 	item := c.respToItem(m, u, l)
 	if item == nil {
 		return
 	}
 
-	pref, _ := subnet.Mask.Size()
-	key := msgToKeyWithSubnet(m, subnet.IP.Mask(subnet.Mask), pref)
+	pref, _ := n.Mask.Size()
+	key := msgToKeyWithSubnet(req, n.IP.Mask(n.Mask), pref)
 	packed := item.pack()
 
 	c.itemsWithSubnetLock.Lock()
@@ -536,12 +537,16 @@ func respectTTLOverrides(ttl, cacheMinTTL, cacheMaxTTL uint32) uint32 {
 func msgToKey(m *dns.Msg) (b []byte) {
 	q := m.Question[0]
 	name := q.Name
-	b = make([]byte, packedMsgLenSz+packedMsgLenSz+len(name))
+	b = make([]byte, 1+packedMsgLenSz+packedMsgLenSz+len(name))
+
+	// Put the DO flag.
+	opt := m.IsEdns0()
+	b[0] = mathutil.BoolToNumber[byte](opt != nil && opt.Do())
 
 	// Put QTYPE, QCLASS, and QNAME.
-	binary.BigEndian.PutUint16(b, q.Qtype)
-	binary.BigEndian.PutUint16(b[packedMsgLenSz:], q.Qclass)
-	copy(b[2*packedMsgLenSz:], strings.ToLower(name))
+	binary.BigEndian.PutUint16(b[1:], q.Qtype)
+	binary.BigEndian.PutUint16(b[1+packedMsgLenSz:], q.Qclass)
+	copy(b[1+2*packedMsgLenSz:], strings.ToLower(name))
 
 	return b
 }
@@ -573,9 +578,7 @@ func msgToKeyWithSubnet(m *dns.Msg, ecsIP net.IP, mask int) (key []byte) {
 	key[0] = mathutil.BoolToNumber[byte](opt != nil && opt.Do())
 
 	// Put Qtype.
-	//
-	// TODO(d.kolyshev): We should put Qtype in key[1:].
-	binary.BigEndian.PutUint16(key[:], q.Qtype)
+	binary.BigEndian.PutUint16(key[1:], q.Qtype)
 
 	// Put Qclass.
 	binary.BigEndian.PutUint16(key[1+packedMsgLenSz:], q.Qclass)

proxy/cache_internal_test.go

@@ -57,23 +57,24 @@ func TestServeCached(t *testing.T) {
 		TCPListenAddr:  []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
 		UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
 		TrustedProxies: defaultTrustedProxies,
-		DNSSECEnabled:  true,
+		DNSSECEnabled:  false,
 		CacheEnabled:   true,
 	})
 
 	// Start listening.
 	servicetest.RequireRun(t, dnsProxy, testTimeout)
 
+	// Create a DNS request.
+	request := (&dns.Msg{}).SetQuestion("google.com.", dns.TypeA)
+	request.SetEdns0(defaultUDPBufSize, false)
+
 	// Fill the cache.
 	reply := (&dns.Msg{
-		MsgHdr: dns.MsgHdr{
-			Response: true,
-		},
 		Answer: []dns.RR{newRR(t, "google.com.", dns.TypeA, 3600, net.IP{8, 8, 8, 8})},
-	}).SetQuestion("google.com.", dns.TypeA)
+	}).SetReply(request)
 	reply.SetEdns0(defaultUDPBufSize, false)
 
-	dnsProxy.cache.set(reply, upstreamWithAddr, testLogger)
+	dnsProxy.cache.set(request, reply, upstreamWithAddr, testLogger)
 
 	// Create a DNS-over-UDP client connection.
 	addr := dnsProxy.Addr(ProtoUDP)
@@ -82,10 +83,6 @@ func TestServeCached(t *testing.T) {
 		Timeout: testTimeout,
 	}
 
-	// Create a DNS request.
-	request := (&dns.Msg{}).SetQuestion("google.com.", dns.TypeA)
-	request.SetEdns0(defaultUDPBufSize, false)
-
 	r, _, err := client.Exchange(request, addr.String())
 	require.NoErrorf(t, err, "error in the first request: %s", err)
 
@@ -187,25 +184,24 @@ func TestCache_expired(t *testing.T) {
 func TestCacheDO(t *testing.T) {
 	testCache := newTestCache(t, nil)
 
+	// Make a request.
+	request := (&dns.Msg{}).SetQuestion("google.com.", dns.TypeA)
+
 	// Fill the cache.
 	reply := (&dns.Msg{
 		MsgHdr: dns.MsgHdr{
 			Response: true,
 		},
 		Answer: []dns.RR{newRR(t, "google.com.", dns.TypeA, 3600, net.IP{8, 8, 8, 8})},
 	}).SetQuestion("google.com.", dns.TypeA)
-	reply.SetEdns0(4096, true)
+	reply.SetEdns0(4096, false)
 
 	// Store in cache.
-	testCache.set(reply, upstreamWithAddr, testLogger)
-
-	// Make a request.
-	request := (&dns.Msg{}).SetQuestion("google.com.", dns.TypeA)
+	testCache.set(request, reply, upstreamWithAddr, testLogger)
 
 	t.Run("without_do", func(t *testing.T) {
-		ci, expired, key := testCache.get(request)
+		ci, expired, _ := testCache.get(request)
 		assert.False(t, expired)
-		assert.Equal(t, msgToKey(request), key)
 		assert.NotNil(t, ci)
 	})
 
@@ -217,30 +213,26 @@ func TestCacheDO(t *testing.T) {
 
 		request.SetEdns0(4096, true)
 
-		ci, expired, key := testCache.get(request)
+		ci, expired, _ := testCache.get(request)
+		require.Nil(t, ci)
 		assert.False(t, expired)
-		assert.Equal(t, msgToKey(request), key)
-
-		require.NotNil(t, ci)
-
-		assert.Equal(t, testUpsAddr, ci.u)
 	})
 }
 
 func TestCacheCNAME(t *testing.T) {
 	testCache := newTestCache(t, nil)
 
+	// Create a DNS request.
+	request := (&dns.Msg{}).SetQuestion("google.com.", dns.TypeA)
+
 	// Fill the cache
 	reply := (&dns.Msg{
 		MsgHdr: dns.MsgHdr{
 			Response: true,
 		},
 		Answer: []dns.RR{newRR(t, "google.com.", dns.TypeCNAME, 3600, "test.google.com.")},
 	}).SetQuestion("google.com.", dns.TypeA)
-	testCache.set(reply, upstreamWithAddr, testLogger)
-
-	// Create a DNS request.
-	request := (&dns.Msg{}).SetQuestion("google.com.", dns.TypeA)
+	testCache.set(request, reply, upstreamWithAddr, testLogger)
 
 	t.Run("no_cnames", func(t *testing.T) {
 		r, expired, _ := testCache.get(request)
@@ -250,7 +242,7 @@ func TestCacheCNAME(t *testing.T) {
 
 	// Now fill the cache with a cacheable CNAME response.
 	reply.Answer = append(reply.Answer, newRR(t, "google.com.", dns.TypeA, 3600, net.IP{8, 8, 8, 8}))
-	testCache.set(reply, upstreamWithAddr, testLogger)
+	testCache.set(request, reply, upstreamWithAddr, testLogger)
 
 	// We are testing that a proper CNAME response gets cached
 	t.Run("cnames_exist", func(t *testing.T) {
@@ -273,7 +265,7 @@ func TestCache_uncacheable(t *testing.T) {
 	reply := (&dns.Msg{}).SetRcode(request, dns.RcodeBadAlg)
 
 	// We are testing that SERVFAIL responses aren't cached
-	testCache.set(reply, upstreamWithAddr, testLogger)
+	testCache.set(request, reply, upstreamWithAddr, testLogger)
 
 	r, expired, _ := testCache.get(request)
 	assert.Nil(t, r)
@@ -330,30 +322,35 @@ func TestCacheExpiration(t *testing.T) {
 		newRR(t, "google.com.", dns.TypeA, 1, net.IP{8, 8, 8, 8}),
 		newRR(t, "yandex.com.", dns.TypeA, 1, net.IP{213, 180, 204, 62}),
 	}
-	replies := make([]*dns.Msg, len(rrs))
-	for i, rr := range rrs {
-		rep := (&dns.Msg{
+	requests := make([]*dns.Msg, 0, len(rrs))
+	responses := make([]*dns.Msg, 0, len(rrs))
+	for _, rr := range rrs {
+		req := (&dns.Msg{
 			MsgHdr: dns.MsgHdr{
-				Response: true,
+				Id: dns.Id(),
 			},
-			Answer: []dns.RR{dns.Copy(rr)},
 		}).SetQuestion(rr.Header().Name, dns.TypeA)
-		dnsProxy.cache.set(rep, upstreamWithAddr, testLogger)
-		replies[i] = rep
+		requests = append(requests, req)
+
+		resp := (&dns.Msg{}).SetReply(req)
+		resp.Answer = append(resp.Answer, rr)
+		responses = append(responses, resp)
+
+		dnsProxy.cache.set(req, resp, upstreamWithAddr, testLogger)
 	}
 
-	for _, r := range replies {
+	for i, r := range requests {
 		ci, expired, key := dnsProxy.cache.get(r)
 		require.NotNil(t, ci)
 
 		assert.False(t, expired)
 		assert.Equal(t, msgToKey(ci.m), key)
 
-		requireEqualMsgs(t, ci.m, r)
+		requireEqualMsgs(t, ci.m, responses[i])
 	}
 
 	assert.Eventually(t, func() bool {
-		for _, r := range replies {
+		for _, r := range requests {
 			if ci, _, _ := dnsProxy.cache.get(r); ci != nil {
 				return false
 			}
@@ -550,13 +547,15 @@ func (tests testCases) run(t *testing.T) {
 	testCache := newTestCache(t, nil)
 
 	for _, res := range tests.cache {
+		request := (&dns.Msg{}).SetQuestion(res.q, res.t)
+
 		reply := (&dns.Msg{
 			MsgHdr: dns.MsgHdr{
 				Response: true,
 			},
 			Answer: res.a,
-		}).SetQuestion(res.q, res.t)
-		testCache.set(reply, upstreamWithAddr, testLogger)
+		}).SetReply(request)
+		testCache.set(request, reply, upstreamWithAddr, testLogger)
 	}
 
 	for _, tc := range tests.cases {
@@ -579,7 +578,7 @@ func (tests testCases) run(t *testing.T) {
 			Answer: tc.a,
 		}).SetQuestion(tc.q, tc.t)
 
-		testCache.set(reply, upstreamWithAddr, testLogger)
+		testCache.set(request, reply, upstreamWithAddr, testLogger)
 
 		requireEqualMsgs(t, ci.m, reply)
 	}
@@ -619,27 +618,33 @@ func setAndGetCache(t *testing.T, c *cache, g *sync.WaitGroup, host, ip string)
 
 	ipAddr := net.ParseIP(ip)
 
+	req := (&dns.Msg{
+		MsgHdr: dns.MsgHdr{
+			Id: dns.Id(),
+		},
+	}).SetQuestion(host, dns.TypeA)
+
 	dnsMsg := (&dns.Msg{
 		MsgHdr: dns.MsgHdr{
 			Response: true,
 		},
 		Answer: []dns.RR{newRR(t, host, dns.TypeA, 1, ipAddr)},
 	}).SetQuestion(host, dns.TypeA)
 
-	c.set(dnsMsg, upstreamWithAddr, testLogger)
+	c.set(req, dnsMsg, upstreamWithAddr, testLogger)
 
 	for range 2 {
 		ci, expired, key := c.get(dnsMsg)
 		require.NotNilf(t, ci, "no cache found for %s", host)
 
 		assert.False(t, expired)
-		assert.Equal(t, msgToKey(dnsMsg), key)
+		assert.Equal(t, msgToKey(req), key)
 
 		requireEqualMsgs(t, ci.m, dnsMsg)
 	}
 
 	assert.Eventuallyf(t, func() bool {
-		ci, _, _ := c.get(dnsMsg)
+		ci, _, _ := c.get(req)
 
 		return ci == nil
 	}, cacheTimeout, cacheTick, "cache for %s should already be removed", host)
@@ -665,7 +670,7 @@ func TestCache_getWithSubnet(t *testing.T) {
 	resp := (&dns.Msg{
 		Answer: []dns.RR{newRR(t, testFQDN, dns.TypeA, 1, net.IP{1, 1, 1, 1})},
 	}).SetReply(req)
-	c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: ip1234, Mask: mask16}, testLogger)
+	c.setWithSubnet(req, resp, upstreamWithAddr, &net.IPNet{IP: ip1234, Mask: mask16}, testLogger)
 
 	t.Run("different_ip", func(t *testing.T) {
 		ci, expired, key := c.getWithSubnet(req, &net.IPNet{IP: ip2234, Mask: mask24})
@@ -678,13 +683,13 @@ func TestCache_getWithSubnet(t *testing.T) {
 	resp = (&dns.Msg{
 		Answer: []dns.RR{newRR(t, testFQDN, dns.TypeA, 1, net.IP{2, 2, 2, 2})},
 	}).SetReply(req)
-	c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: ip2234, Mask: mask16}, testLogger)
+	c.setWithSubnet(req, resp, upstreamWithAddr, &net.IPNet{IP: ip2234, Mask: mask16}, testLogger)
 
 	// Add a response entry without subnet.
 	resp = (&dns.Msg{
 		Answer: []dns.RR{newRR(t, testFQDN, dns.TypeA, 1, net.IP{3, 3, 3, 3})},
 	}).SetReply(req)
-	c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: nil, Mask: nil}, testLogger)
+	c.setWithSubnet(req, resp, upstreamWithAddr, &net.IPNet{IP: nil, Mask: nil}, testLogger)
 
 	t.Run("with_subnet_1", func(t *testing.T) {
 		ci, expired, key := c.getWithSubnet(req, &net.IPNet{IP: ip1234, Mask: mask24})
@@ -751,6 +756,7 @@ func TestCache_getWithSubnet_mask(t *testing.T) {
 
 	// Cache IP network that contains the testIP.
 	c.setWithSubnet(
+		req,
 		resp,
 		upstreamWithAddr,
 		&net.IPNet{IP: cachedIP, Mask: cidrMask},

proxy/dnscontext.go

@@ -165,6 +165,9 @@ func (dctx *DNSContext) scrub() {
 	}
 
 	// We should guarantee that all the values we need are calculated.
+	//
+	// TODO(e.burkov):  It's theoretically guaranteed already, investigate the
+	// need to this call.
 	dctx.calcFlagsAndSize()
 
 	// RFC-6891 (https://tools.ietf.org/html/rfc6891) states that response

proxy/proxy.go

@@ -701,6 +701,15 @@ func (p *Proxy) Resolve(ctx context.Context, dctx *DNSContext) (err error) {
 
 	cacheWorks := p.cacheWorks(dctx)
 	if cacheWorks {
+		// Request for DNSSEC from the upstream to cache the
+		// DNSSEC resource records as well.  In case of disabled DNSSEC,
+		// requesting and therefore caching of DNSSEC resource records depends
+		// on the DO bit of the initiating query.
+		//
+		// See https://datatracker.ietf.org/doc/html/rfc4035#section-4.5 and
+		// https://datatracker.ietf.org/doc/html/rfc3225#section-3.
+		p.addDO(dctx.Req)
+
 		// Only add pending requests if the cache is enabled, since this is a
 		// mitigation against cache poisoning.
 		//
@@ -714,14 +723,11 @@ func (p *Proxy) Resolve(ctx context.Context, dctx *DNSContext) (err error) {
 
 		if p.replyFromCache(dctx) {
 			// Complete the response from cache.
+			filterMsg(dctx.Res, dctx.Res, dctx.adBit, dctx.doBit, 0)
 			dctx.scrub()
 
 			return nil
 		}
-
-		// On cache miss request for DNSSEC from the upstream to cache it
-		// afterwards.
-		p.addDO(dctx.Req)
 	}
 
 	var ok bool
@@ -800,12 +806,6 @@ func (p *Proxy) cacheWorks(dctx *DNSContext) (ok bool) {
 		// Don't cache the requests intended for local upstream servers, those
 		// should be fast enough as is.
 		reason = "requested address is private"
-	case !p.DNSSECEnabled && !dctx.doBit:
-		// Don't cache the responses without DNSSEC RRs if DNSSEC is disabled
-		// and DO bit is not set, since those responses may differ from the ones
-		// with DNSSEC RRs and thus may be not the desired result for user.  In
-		// case DNSSEC is enabled in the proxy, the DO bit will be enforced.
-		reason = "dnssec disabled"
 	case dctx.Req.CheckingDisabled:
 		// Also don't lookup the cache for responses with DNSSEC checking
 		// disabled since only validated responses are cached and those may be