Pull request: proxy: custom upstream cache

Merge in GO/dnsproxy from custom-ups-cache to master

Squashed commit of the following:

commit 26ee425
Author: Dimitry Kolyshev <[email protected]>
Date:   Tue Nov 14 10:22:17 2023 +0200

    proxy: custom upstream cache

commit ae2f902
Author: Dimitry Kolyshev <[email protected]>
Date:   Mon Nov 13 14:31:28 2023 +0200

    proxy: imp docs

commit 6991d32
Author: Dimitry Kolyshev <[email protected]>
Date:   Mon Nov 13 10:29:57 2023 +0200

    proxy: imp code

commit 34662cc
Merge: cb80501 65b5293
Author: Dimitry Kolyshev <[email protected]>
Date:   Mon Nov 13 10:10:41 2023 +0200

    Merge remote-tracking branch 'origin/master' into custom-ups-cache

commit cb80501
Merge: a933441 ea0f2b1
Author: Dimitry Kolyshev <[email protected]>
Date:   Fri Nov 10 15:32:25 2023 +0200

    Merge remote-tracking branch 'origin/master' into custom-ups-cache

commit a933441
Author: Dimitry Kolyshev <[email protected]>
Date:   Fri Nov 10 15:17:19 2023 +0200

    proxy: imp docs

commit 1cff28e
Author: Dimitry Kolyshev <[email protected]>
Date:   Fri Nov 10 14:31:16 2023 +0200

    proxy: imp code

commit 62fed0c
Author: Dimitry Kolyshev <[email protected]>
Date:   Fri Nov 10 14:07:20 2023 +0200

    proxy: imp code

commit f1aaa42
Merge: d6e3a57 b490fb4
Author: Dimitry Kolyshev <[email protected]>
Date:   Thu Nov 9 12:04:13 2023 +0200

    Merge remote-tracking branch 'origin/master' into custom-ups-cache

commit d6e3a57
Author: Dimitry Kolyshev <[email protected]>
Date:   Wed Nov 8 12:11:43 2023 +0200

    proxy: imp code

commit 0b11894
Author: Dimitry Kolyshev <[email protected]>
Date:   Wed Nov 8 10:37:28 2023 +0200

    proxy: imp code

commit 1775ae9
Author: Dimitry Kolyshev <[email protected]>
Date:   Tue Nov 7 14:44:03 2023 +0200

    proxy: custom upstream cache

Author: Mizzick

proxy/dnscontext.go

@@ -47,9 +47,10 @@ type DNSContext struct {
 	// ReqECS is the EDNS Client Subnet used in the request.
 	ReqECS *net.IPNet
 
-	// CustomUpstreamConfig is only used for current request.  The Resolve
-	// method of Proxy uses it instead of the default servers if it's not nil.
-	CustomUpstreamConfig *UpstreamConfig
+	// CustomUpstreamConfig is the upstreams configuration used only for current
+	// request.  The Resolve method of Proxy uses it instead of the default
+	// servers if it's not nil.
+	CustomUpstreamConfig *CustomUpstreamConfig
 
 	// Req is the request message.
 	Req *dns.Msg
@@ -157,3 +158,53 @@ const (
 	// DoQv1 represents DoQ v1.0: https://www.rfc-editor.org/rfc/rfc9250.html.
 	DoQv1 DoQVersion = 0x01
 )
+
+// CustomUpstreamConfig contains upstreams configuration with an optional cache.
+type CustomUpstreamConfig struct {
+	// upstream is the upstream configuration.
+	upstream *UpstreamConfig
+
+	// cache is an optional cache for upstreams in the current configuration.
+	// It is disabled if nil.
+	//
+	// TODO(d.kolyshev): Move this cache to [UpstreamConfig].
+	cache *cache
+}
+
+// NewCustomUpstreamConfig returns new custom upstream configuration.
+func NewCustomUpstreamConfig(
+	u *UpstreamConfig,
+	cacheEnabled bool,
+	cacheSize int,
+	enableEDNSClientSubnet bool,
+) (c *CustomUpstreamConfig) {
+	var customCache *cache
+	if cacheEnabled {
+		// TODO(d.kolyshev): Support optimistic with newOptimisticResolver.
+		customCache = newCache(cacheSize, enableEDNSClientSubnet, false)
+	}
+
+	return &CustomUpstreamConfig{
+		upstream: u,
+		cache:    customCache,
+	}
+}
+
+// Close closes the custom upstream config.
+func (c *CustomUpstreamConfig) Close() (err error) {
+	if c.upstream == nil {
+		return nil
+	}
+
+	return c.upstream.Close()
+}
+
+// ClearCache removes all items from the cache.
+func (c *CustomUpstreamConfig) ClearCache() {
+	if c.cache == nil {
+		return
+	}
+
+	c.cache.clearItems()
+	c.cache.clearItemsWithSubnet()
+}

proxy/proxy.go

@@ -139,6 +139,8 @@ type Proxy struct {
 	// --
 
 	// cache is used to cache requests.  It is disabled if nil.
+	//
+	// TODO(d.kolyshev): Move this cache to [Proxy.UpstreamConfig] field.
 	cache *cache
 
 	// shortFlighter is used to resolve the expired cached requests without
@@ -494,7 +496,7 @@ func (p *Proxy) selectUpstreams(d *DNSContext) (upstreams []upstream.Upstream) {
 
 		if custom := d.CustomUpstreamConfig; custom != nil {
 			// Try to use custom.
-			upstreams = getUpstreams(custom, host)
+			upstreams = getUpstreams(custom.upstream, host)
 			if len(upstreams) > 0 {
 				return upstreams
 			}
@@ -622,7 +624,6 @@ func (p *Proxy) Resolve(dctx *DNSContext) (err error) {
 
 	dctx.calcFlagsAndSize()
 
-	// Use cache only if it's enabled and the query doesn't use custom upstream.
 	// Also don't lookup the cache for responses with DNSSEC checking disabled
 	// since only validated responses are cached and those may be not the
 	// desired result for user specifying CD flag.
@@ -676,9 +677,12 @@ func (p *Proxy) cacheWorks(dctx *DNSContext) (ok bool) {
 	switch {
 	case p.cache == nil:
 		reason = "disabled"
-	case dctx.CustomUpstreamConfig != nil:
+	case dctx.CustomUpstreamConfig != nil && dctx.CustomUpstreamConfig.cache == nil:
+		// In case of custom upstream cache is not configured, the global proxy
+		// cache cannot be used because different upstreams can return different
+		// results.
 		// See https://github.com/AdguardTeam/dnsproxy/issues/169.
-		reason = "custom upstreams used"
+		reason = "custom upstreams cache is not configured"
 	case dctx.Req.CheckingDisabled:
 		reason = "dnssec check disabled"
 	default:

proxy/proxy_test.go

@@ -846,8 +846,13 @@ func TestProxy_ReplyFromUpstream_badResponse(t *testing.T) {
 	}
 
 	d := &DNSContext{
-		CustomUpstreamConfig: &UpstreamConfig{Upstreams: []upstream.Upstream{u}},
-		Req:                  createHostTestMessage("host"),
+		CustomUpstreamConfig: NewCustomUpstreamConfig(
+			&UpstreamConfig{Upstreams: []upstream.Upstream{u}},
+			false,
+			0,
+			false,
+		),
+		Req: createHostTestMessage("host"),
 		Addr: &net.TCPAddr{
 			IP: net.IP{1, 2, 3, 0},
 		},
@@ -869,7 +874,7 @@ func TestExchangeCustomUpstreamConfig(t *testing.T) {
 	testutil.CleanupAndRequireSuccess(t, prx.Stop)
 
 	ansIP := net.IP{4, 3, 2, 1}
-	u := testUpstream{
+	u := &testUpstream{
 		ans: []dns.RR{&dns.A{
 			Hdr: dns.RR_Header{
 				Rrtype: dns.TypeA,
@@ -881,14 +886,86 @@ func TestExchangeCustomUpstreamConfig(t *testing.T) {
 	}
 
 	d := DNSContext{
-		CustomUpstreamConfig: &UpstreamConfig{Upstreams: []upstream.Upstream{&u}},
+		CustomUpstreamConfig: NewCustomUpstreamConfig(
+			&UpstreamConfig{Upstreams: []upstream.Upstream{u}},
+			false,
+			0,
+			false,
+		),
+		Req:  createHostTestMessage("host"),
+		Addr: &net.TCPAddr{IP: net.IP{1, 2, 3, 0}},
+	}
+
+	err = prx.Resolve(&d)
+	require.NoError(t, err)
+
+	assert.Equal(t, ansIP, getIPFromResponse(d.Res))
+}
+
+func TestExchangeCustomUpstreamConfigCache(t *testing.T) {
+	prx := createTestProxy(t, nil)
+	prx.CacheEnabled = true
+	prx.initCache()
+
+	err := prx.Start()
+	require.NoError(t, err)
+	testutil.CleanupAndRequireSuccess(t, prx.Stop)
+
+	var count int
+
+	ansIP := net.IP{4, 3, 2, 1}
+	exchangeFunc := func(m *dns.Msg) (resp *dns.Msg, err error) {
+		resp = &dns.Msg{}
+		resp.SetReply(m)
+		resp.Answer = append(resp.Answer, &dns.A{
+			Hdr: dns.RR_Header{
+				Name:   m.Question[0].Name,
+				Class:  dns.ClassINET,
+				Rrtype: dns.TypeA,
+				Ttl:    defaultTestTTL,
+			},
+			A: ansIP,
+		})
+
+		count++
+
+		return resp, nil
+	}
+	u := &funcUpstream{
+		exchangeFunc: exchangeFunc,
+	}
+
+	customUpstreamConfig := NewCustomUpstreamConfig(
+		&UpstreamConfig{Upstreams: []upstream.Upstream{u}},
+		true,
+		defaultCacheSize,
+		prx.EnableEDNSClientSubnet,
+	)
+
+	d := DNSContext{
+		CustomUpstreamConfig: customUpstreamConfig,
 		Req:                  createHostTestMessage("host"),
 		Addr:                 &net.TCPAddr{IP: net.IP{1, 2, 3, 0}},
 	}
 
 	err = prx.Resolve(&d)
 	require.NoError(t, err)
 
+	assert.Equal(t, 1, count)
+	assert.Equal(t, ansIP, getIPFromResponse(d.Res))
+
+	err = prx.Resolve(&d)
+	require.NoError(t, err)
+
+	assert.Equal(t, 1, count)
+	assert.Equal(t, ansIP, getIPFromResponse(d.Res))
+
+	customUpstreamConfig.ClearCache()
+
+	err = prx.Resolve(&d)
+	require.NoError(t, err)
+
+	assert.Equal(t, 2, count)
 	assert.Equal(t, ansIP, getIPFromResponse(d.Res))
 }
 

proxy/proxycache.go

@@ -7,22 +7,34 @@ import (
 	"github.com/AdguardTeam/golibs/netutil"
 )
 
-// replyFromCache tries to get the response from general or subnet cache.
-// Returns true on success.
+// cacheForContext returns cache object for the given context.
+func (p *Proxy) cacheForContext(d *DNSContext) (c *cache) {
+	if d.CustomUpstreamConfig != nil && d.CustomUpstreamConfig.cache != nil {
+		return d.CustomUpstreamConfig.cache
+	}
+
+	return p.cache
+}
+
+// replyFromCache tries to get the response from general or subnet cache.  In
+// case the cache is present in d, it's used first.  Returns true on success.
 func (p *Proxy) replyFromCache(d *DNSContext) (hit bool) {
+	dctxCache := p.cacheForContext(d)
+
 	var ci *cacheItem
 	var hitMsg string
 	var expired bool
 	var key []byte
 
+	// TODO(d.kolyshev): Use EnableEDNSClientSubnet from dctxCache.
 	if !p.Config.EnableEDNSClientSubnet {
-		ci, expired, key = p.cache.get(d.Req)
+		ci, expired, key = dctxCache.get(d.Req)
 		hitMsg = "serving cached response"
 	} else if d.ReqECS != nil {
-		ci, expired, key = p.cache.getWithSubnet(d.Req, d.ReqECS)
+		ci, expired, key = dctxCache.getWithSubnet(d.Req, d.ReqECS)
 		hitMsg = "serving response from subnet cache"
 	} else {
-		ci, expired, key = p.cache.get(d.Req)
+		ci, expired, key = dctxCache.get(d.Req)
 		hitMsg = "serving response from general cache"
 	}
 
@@ -35,7 +47,7 @@ func (p *Proxy) replyFromCache(d *DNSContext) (hit bool) {
 
 	log.Debug("dnsproxy: cache: %s", hitMsg)
 
-	if p.cache.optimistic && expired {
+	if dctxCache.optimistic && expired {
 		// Build a reduced clone of the current context to avoid data race.
 		minCtxClone := &DNSContext{
 			// It is only read inside the optimistic resolver.
@@ -53,10 +65,13 @@ func (p *Proxy) replyFromCache(d *DNSContext) (hit bool) {
 	return hit
 }
 
-// cacheResp stores the response from d in general or subnet cache.
+// cacheResp stores the response from d in general or subnet cache.  In case the
+// cache is present in d, it's used first.
 func (p *Proxy) cacheResp(d *DNSContext) {
+	dctxCache := p.cacheForContext(d)
+
 	if !p.EnableEDNSClientSubnet {
-		p.cache.set(d.Res, d.Upstream)
+		dctxCache.set(d.Res, d.Upstream)
 
 		return
 	}
@@ -92,13 +107,13 @@ func (p *Proxy) cacheResp(d *DNSContext) {
 
 		log.Debug("dnsproxy: cache: ecs option in response: %s", ecs)
 
-		p.cache.setWithSubnet(d.Res, d.Upstream, ecs)
+		dctxCache.setWithSubnet(d.Res, d.Upstream, ecs)
 	case d.ReqECS != nil:
 		// Cache the response for all subnets since the server doesn't support
 		// EDNS Client Subnet option.
-		p.cache.setWithSubnet(d.Res, d.Upstream, &net.IPNet{IP: nil, Mask: nil})
+		dctxCache.setWithSubnet(d.Res, d.Upstream, &net.IPNet{IP: nil, Mask: nil})
 	default:
-		p.cache.set(d.Res, d.Upstream)
+		dctxCache.set(d.Res, d.Upstream)
 	}
 }