Pull request: custom upstreams cache

Merge in DNS/adguard-home from custom-ups-cache to master

Squashed commit of the following:

commit 98428a87520f70cb522701d8eccfe4c529be1e40
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Nov 17 10:53:32 2023 +0200

    all: upd dep

commit 775a639af4a2a45220b17e8b0037edc126ff62e4
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Nov 17 09:52:31 2023 +0200

    dnsforward: imp test

commit e9e2a58b48e8588dfcb28df319d4651e1fe77af5
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Nov 17 09:44:46 2023 +0200

    docs: changelog

commit a6d67218f037c8fec29e5fa2967476d63c3cfc32
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Nov 17 09:37:17 2023 +0200

    all: upd dep

commit b101ff6e0cf393dacdee6fb68d33ba8f11c36280
Merge: d61f4eb88 8bb1aad73
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Thu Nov 16 15:54:05 2023 +0200

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

commit d61f4eb8871f8ae8504259998bf9015b29001cfb
Merge: 567a8a4af fdf60eeed
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Mon Nov 13 10:32:22 2023 +0200

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

commit 567a8a4af34ad001d0e6d7d2efdc123205569e8c
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Mon Nov 13 10:30:24 2023 +0200

    home: imp code

commit a3c16facbebc166e5c0c731c1e892b61c0950d9e
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Nov 10 14:34:04 2023 +0200

    all: imp code

commit 84160eafee1d0f2d0cd3f025f2d5070e4f597ad6
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Nov 10 14:31:26 2023 +0200

    all: conf custom ups cache

commit b7f6581901ebad96c87e765a305a1fa5b336efbb
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Nov 10 14:29:47 2023 +0200

    all: conf custom ups cache

commit d07df945d4e7614a679ef5dc77756096abf1e66c
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Nov 10 09:26:29 2023 +0200

    all: docs

commit 998124bac08889c7d354dd1a099929726725bccc
Merge: f665e2f85 53170d871
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Nov 10 09:24:28 2023 +0200

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

commit f665e2f85bce12d95f80aba6614b6bfd4874b122
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Nov 10 09:22:46 2023 +0200

    all: conf custom ups cache

commit a4b26973bef4f3b339198ffbe52a50baca303daf
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Thu Nov 9 12:46:39 2023 +0200

    all: conf custom ups cache

internal/aghtest/interface.go

@@ -125,7 +125,7 @@ type ClientsContainer struct {
 	OnUpstreamConfigByID func(
 		id string,
 		boot upstream.Resolver,
-	) (conf *proxy.UpstreamConfig, err error)
+	) (conf *proxy.CustomUpstreamConfig, err error)
 }
 
 // UpstreamConfigByID implements the [dnsforward.ClientsContainer] interface
@@ -133,7 +133,7 @@ type ClientsContainer struct {
 func (c *ClientsContainer) UpstreamConfigByID(
 	id string,
 	boot upstream.Resolver,
-) (conf *proxy.UpstreamConfig, err error) {
+) (conf *proxy.CustomUpstreamConfig, err error) {
 	return c.OnUpstreamConfigByID(id, boot)
 }
 

internal/dnsforward/config.go

@@ -34,7 +34,10 @@ type ClientsContainer interface {
 	// returns nil if there is no custom upstream configuration for the client.
 	// The id is expected to be either a string representation of an IP address
 	// or the ClientID.
-	UpstreamConfigByID(id string, boot upstream.Resolver) (conf *proxy.UpstreamConfig, err error)
+	UpstreamConfigByID(
+		id string,
+		boot upstream.Resolver,
+	) (conf *proxy.CustomUpstreamConfig, err error)
 }
 
 // Config represents the DNS filtering configuration of AdGuard Home. The zero

internal/dnsforward/dnsforward_test.go

@@ -644,10 +644,15 @@ func TestBlockedRequest(t *testing.T) {
 }
 
 func TestServerCustomClientUpstream(t *testing.T) {
+	const defaultCacheSize = 1024 * 1024
+
+	var upsCalledCounter uint32
+
 	forwardConf := ServerConfig{
 		UDPListenAddrs: []*net.UDPAddr{{}},
 		TCPListenAddrs: []*net.TCPAddr{{}},
 		Config: Config{
+			CacheSize: defaultCacheSize,
 			EDNSClientSubnet: &EDNSClientSubnet{
 				Enabled: false,
 			},
@@ -658,34 +663,51 @@ func TestServerCustomClientUpstream(t *testing.T) {
 	}, forwardConf, nil)
 
 	ups := aghtest.NewUpstreamMock(func(req *dns.Msg) (resp *dns.Msg, err error) {
+		atomic.AddUint32(&upsCalledCounter, 1)
+
 		return aghalg.Coalesce(
 			aghtest.MatchedResponse(req, dns.TypeA, "host", "192.168.0.1"),
 			new(dns.Msg).SetRcode(req, dns.RcodeNameError),
 		), nil
 	})
+
+	customUpsConf := proxy.NewCustomUpstreamConfig(
+		&proxy.UpstreamConfig{
+			Upstreams: []upstream.Upstream{ups},
+		},
+		true,
+		defaultCacheSize,
+		forwardConf.EDNSClientSubnet.Enabled,
+	)
+
 	s.conf.ClientsContainer = &aghtest.ClientsContainer{
 		OnUpstreamConfigByID: func(
 			_ string,
 			_ upstream.Resolver,
-		) (conf *proxy.UpstreamConfig, err error) {
-			return &proxy.UpstreamConfig{Upstreams: []upstream.Upstream{ups}}, nil
+		) (conf *proxy.CustomUpstreamConfig, err error) {
+			return customUpsConf, nil
 		},
 	}
+
 	startDeferStop(t, s)
 
-	addr := s.dnsProxy.Addr(proxy.ProtoUDP)
+	addr := s.dnsProxy.Addr(proxy.ProtoUDP).String()
 
 	// Send test request.
 	req := createTestMessage("host.")
 
-	reply, err := dns.Exchange(req, addr.String())
+	reply, err := dns.Exchange(req, addr)
 	require.NoError(t, err)
+	require.NotEmpty(t, reply.Answer)
+	require.Len(t, reply.Answer, 1)
 
 	assert.Equal(t, dns.RcodeSuccess, reply.Rcode)
-	require.NotEmpty(t, reply.Answer)
-
-	require.Len(t, reply.Answer, 1)
 	assert.Equal(t, net.IP{192, 168, 0, 1}, reply.Answer[0].(*dns.A).A)
+	assert.Equal(t, uint32(1), atomic.LoadUint32(&upsCalledCounter))
+
+	_, err = dns.Exchange(req, addr)
+	require.NoError(t, err)
+	assert.Equal(t, uint32(1), atomic.LoadUint32(&upsCalledCounter))
 }
 
 // testCNAMEs is a map of names and CNAMEs necessary for the TestUpstream work.

internal/home/client.go

@@ -14,11 +14,11 @@ import (
 
 // Client contains information about persistent clients.
 type Client struct {
-	// upstreamConfig is the custom upstream config for this client.  If
-	// it's nil, it has not been initialized yet.  If it's non-nil and
-	// empty, there are no valid upstreams.  If it's non-nil and non-empty,
-	// these upstream must be used.
-	upstreamConfig *proxy.UpstreamConfig
+	// upstreamConfig is the custom upstream configuration for this client.  If
+	// it's nil, it has not been initialized yet.  If it's non-nil and empty,
+	// there are no valid upstreams.  If it's non-nil and non-empty, these
+	// upstream must be used.
+	upstreamConfig *proxy.CustomUpstreamConfig
 
 	safeSearchConf filtering.SafeSearchConfig
 	SafeSearch     filtering.SafeSearch
@@ -32,6 +32,9 @@ type Client struct {
 	Tags      []string
 	Upstreams []string
 
+	UpstreamsCacheSize    uint32
+	UpstreamsCacheEnabled bool
+
 	UseOwnSettings        bool
 	FilteringEnabled      bool
 	SafeBrowsingEnabled   bool
@@ -57,8 +60,7 @@ func (c *Client) ShallowClone() (sh *Client) {
 // closeUpstreams closes the client-specific upstream config of c if any.
 func (c *Client) closeUpstreams() (err error) {
 	if c.upstreamConfig != nil {
-		err = c.upstreamConfig.Close()
-		if err != nil {
+		if err = c.upstreamConfig.Close(); err != nil {
 			return fmt.Errorf("closing upstreams of client %q: %w", c.Name, err)
 		}
 	}

internal/home/clients.go

@@ -185,6 +185,14 @@ type clientObject struct {
 	Tags      []string `yaml:"tags"`
 	Upstreams []string `yaml:"upstreams"`
 
+	// UpstreamsCacheSize is the DNS cache size (in bytes).
+	//
+	// TODO(d.kolyshev): Use [datasize.Bytesize].
+	UpstreamsCacheSize uint32 `yaml:"upstreams_cache_size"`
+
+	// UpstreamsCacheEnabled indicates if the DNS cache is enabled.
+	UpstreamsCacheEnabled bool `yaml:"upstreams_cache_enabled"`
+
 	UseGlobalSettings        bool `yaml:"use_global_settings"`
 	FilteringEnabled         bool `yaml:"filtering_enabled"`
 	ParentalEnabled          bool `yaml:"parental_enabled"`
@@ -216,6 +224,8 @@ func (clients *clientsContainer) addFromConfig(
 			UseOwnBlockedServices: !o.UseGlobalBlockedServices,
 			IgnoreQueryLog:        o.IgnoreQueryLog,
 			IgnoreStatistics:      o.IgnoreStatistics,
+			UpstreamsCacheEnabled: o.UpstreamsCacheEnabled,
+			UpstreamsCacheSize:    o.UpstreamsCacheSize,
 		}
 
 		if o.SafeSearchConf.Enabled {
@@ -429,11 +439,12 @@ func (clients *clientsContainer) shouldCountClient(ids []string) (y bool) {
 var _ dnsforward.ClientsContainer = (*clientsContainer)(nil)
 
 // UpstreamConfigByID implements the [dnsforward.ClientsContainer] interface for
-// *clientsContainer.
+// *clientsContainer.  upsConf is nil if the client isn't found or if the client
+// has no custom upstreams.
 func (clients *clientsContainer) UpstreamConfigByID(
 	id string,
 	bootstrap upstream.Resolver,
-) (upsConf *proxy.UpstreamConfig, err error) {
+) (conf *proxy.CustomUpstreamConfig, err error) {
 	clients.lock.Lock()
 	defer clients.lock.Unlock()
 
@@ -449,8 +460,8 @@ func (clients *clientsContainer) UpstreamConfigByID(
 		return nil, nil
 	}
 
-	var conf *proxy.UpstreamConfig
-	conf, err = proxy.ParseUpstreamsConfig(
+	var upsConf *proxy.UpstreamConfig
+	upsConf, err = proxy.ParseUpstreamsConfig(
 		upstreams,
 		&upstream.Options{
 			Bootstrap:    bootstrap,
@@ -464,6 +475,12 @@ func (clients *clientsContainer) UpstreamConfigByID(
 		return nil, err
 	}
 
+	conf = proxy.NewCustomUpstreamConfig(
+		upsConf,
+		c.UpstreamsCacheEnabled,
+		int(c.UpstreamsCacheSize),
+		config.DNS.EDNSClientSubnet.Enabled,
+	)
 	c.upstreamConfig = conf
 
 	return conf, nil

internal/home/clients_internal_test.go

@@ -355,13 +355,11 @@ func TestClientsCustomUpstream(t *testing.T) {
 	require.NoError(t, err)
 	assert.True(t, ok)
 
-	config, err := clients.UpstreamConfigByID("1.2.3.4", net.DefaultResolver)
-	assert.Nil(t, config)
+	upsConf, err := clients.UpstreamConfigByID("1.2.3.4", net.DefaultResolver)
+	assert.Nil(t, upsConf)
 	assert.NoError(t, err)
 
-	config, err = clients.UpstreamConfigByID("1.1.1.1", net.DefaultResolver)
-	require.NotNil(t, config)
+	upsConf, err = clients.UpstreamConfigByID("1.1.1.1", net.DefaultResolver)
+	require.NotNil(t, upsConf)
 	assert.NoError(t, err)
-	assert.Len(t, config.Upstreams, 1)
-	assert.Len(t, config.DomainReservedUpstreams, 1)
 }