From 2c46bc92fe6d493d48f704ad4779bcd8de23ebc6 Mon Sep 17 00:00:00 2001
From: Stanislav Chzhen <s.chzhen@adguard.com>
Date: Fri, 25 Apr 2025 21:33:36 +0300
Subject: [PATCH] Pull request 2399: AGDNS-2686-fix-custom-upstream-cache

Merge in DNS/adguard-home from AGDNS-2686-fix-custom-upstream-cache to master

Squashed commit of the following:

commit 11ad20a225e0e21a59552dc885fbcb2d3acc1cef
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Thu Apr 24 18:51:35 2025 +0300

    client: imp docs

commit e6d73f2d7a9f2ea181b321dd0029cf6c42ddeba5
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Thu Apr 24 17:36:05 2025 +0300

    all: imp chlog

commit b8fdd884b801db28f03efb00bd871df2332cf40a
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Thu Apr 24 17:05:11 2025 +0300

    client: fix dhcp clients cache

commit 1760699fcb8e61580a48e61037b805b8aa8ca8b4
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Thu Apr 24 15:37:50 2025 +0300

    all: upd chlog

commit c6f049c200736032e2d78a2023db7d8cc6c32917
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Thu Apr 24 15:33:53 2025 +0300

    client: imp tests

commit 7432de722292ef74bbdf5fbd875ea67d55b29040
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Thu Apr 24 15:32:26 2025 +0300

    client: fix custom upstream cache
---
 CHANGELOG.md                       |  2 +
 internal/client/persistent.go      |  9 ++--
 internal/client/storage.go         |  2 +-
 internal/client/storage_test.go    | 76 +++++++++++++++++++++++-------
 internal/client/upstreammanager.go |  1 +
 5 files changed, 69 insertions(+), 21 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 388ff3dc..99b793f7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -20,6 +20,8 @@ NOTE: Add new changes BELOW THIS COMMENT.
 
 ### Fixed
 
+- DNS cache not working for custom upstream configurations.
+
 - Validation process for the DNS-over-TLS, DNS-over-QUIC, and HTTPS ports on the *Encryption Settings* page.
 
 <!--
diff --git a/internal/client/persistent.go b/internal/client/persistent.go
index 56bd8f47..145dcc4b 100644
--- a/internal/client/persistent.go
+++ b/internal/client/persistent.go
@@ -70,7 +70,9 @@ type Persistent struct {
 	// Tags is a list of client tags that categorize the client.
 	Tags []string
 
-	// Upstreams is a list of custom upstream DNS servers for the client.
+	// Upstreams is a list of custom upstream DNS servers for the client.  If
+	// it's empty, the custom upstream cache is disabled, regardless of the
+	// value of UpstreamsCacheEnabled.
 	Upstreams []string
 
 	// IPs is a list of IP addresses that identify the client.  The client must
@@ -94,10 +96,11 @@ type Persistent struct {
 	// UID is the unique identifier of the persistent client.
 	UID UID
 
-	// UpstreamsCacheSize is the cache size for custom upstreams.
+	// UpstreamsCacheSize defines the size of the custom upstream cache.
 	UpstreamsCacheSize uint32
 
-	// UpstreamsCacheEnabled specifies whether custom upstreams are used.
+	// UpstreamsCacheEnabled specifies whether the custom upstream cache is
+	// used.  If true, the list of Upstreams should not be empty.
 	UpstreamsCacheEnabled bool
 
 	// UseOwnSettings specifies whether custom filtering settings are used.
diff --git a/internal/client/storage.go b/internal/client/storage.go
index c5382648..ef7d5209 100644
--- a/internal/client/storage.go
+++ b/internal/client/storage.go
@@ -781,7 +781,7 @@ func (s *Storage) CustomUpstreamConfig(
 
 	c, ok := s.index.findByClientID(ClientID(id))
 	if !ok {
-		c, ok = s.index.findByIP(addr)
+		c, ok = s.findByIP(addr)
 	}
 
 	if !ok {
diff --git a/internal/client/storage_test.go b/internal/client/storage_test.go
index 7d6051ca..b1f68e41 100644
--- a/internal/client/storage_test.go
+++ b/internal/client/storage_test.go
@@ -1169,29 +1169,20 @@ func TestStorage_RangeByName(t *testing.T) {
 
 func TestStorage_CustomUpstreamConfig(t *testing.T) {
 	const (
-		existingName     = "existing_name"
-		existingClientID = "existing_client_id"
-
+		existingClientID    = "existing_client_id"
 		nonExistingClientID = "non_existing_client_id"
 	)
 
 	var (
-		existingClientUID = client.MustNewUID()
-		existingIP        = netip.MustParseAddr("192.0.2.1")
-
+		existingIP    = netip.MustParseAddr("192.0.2.1")
 		nonExistingIP = netip.MustParseAddr("192.0.2.255")
 
+		dhcpCliIP  = netip.MustParseAddr("192.0.2.2")
+		dhcpCliMAC = errors.Must(net.ParseMAC("02:00:00:00:00:00"))
+
 		testUpstreamTimeout = time.Second
 	)
 
-	existingClient := &client.Persistent{
-		Name:      existingName,
-		IPs:       []netip.Addr{existingIP},
-		ClientIDs: []client.ClientID{existingClientID},
-		UID:       existingClientUID,
-		Upstreams: []string{"192.0.2.0"},
-	}
-
 	date := time.Now()
 	clock := &faketime.Clock{
 		OnNow: func() (now time.Time) {
@@ -1201,7 +1192,30 @@ func TestStorage_CustomUpstreamConfig(t *testing.T) {
 		},
 	}
 
-	s := newTestStorage(t, clock)
+	ipToMAC := map[netip.Addr]net.HardwareAddr{
+		dhcpCliIP: dhcpCliMAC,
+	}
+
+	dhcp := &testDHCP{
+		OnLeases: func() (ls []*dhcpsvc.Lease) {
+			panic("not implemented")
+		},
+		OnHostBy: func(ip netip.Addr) (host string) {
+			panic("not implemented")
+		},
+		OnMACBy: func(ip netip.Addr) (mac net.HardwareAddr) {
+			return ipToMAC[ip]
+		},
+	}
+
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
+	s, err := client.NewStorage(ctx, &client.StorageConfig{
+		Logger: slogutil.NewDiscardLogger(),
+		Clock:  clock,
+		DHCP:   dhcp,
+	})
+	require.NoError(t, err)
+
 	s.UpdateCommonUpstreamConfig(&client.CommonUpstreamConfig{
 		UpstreamTimeout: testUpstreamTimeout,
 	})
@@ -1210,8 +1224,21 @@ func TestStorage_CustomUpstreamConfig(t *testing.T) {
 		return s.Shutdown(testutil.ContextWithTimeout(t, testTimeout))
 	})
 
-	ctx := testutil.ContextWithTimeout(t, testTimeout)
-	err := s.Add(ctx, existingClient)
+	err = s.Add(ctx, &client.Persistent{
+		Name:      "client_first",
+		IPs:       []netip.Addr{existingIP},
+		ClientIDs: []client.ClientID{existingClientID},
+		UID:       client.MustNewUID(),
+		Upstreams: []string{"192.0.2.0"},
+	})
+	require.NoError(t, err)
+
+	err = s.Add(ctx, &client.Persistent{
+		Name:      "client_second",
+		MACs:      []net.HardwareAddr{dhcpCliMAC},
+		UID:       client.MustNewUID(),
+		Upstreams: []string{"192.0.2.0"},
+	})
 	require.NoError(t, err)
 
 	testCases := []struct {
@@ -1229,6 +1256,11 @@ func TestStorage_CustomUpstreamConfig(t *testing.T) {
 		cliID:       "",
 		cliAddr:     existingIP,
 		wantNilConf: assert.NotNil,
+	}, {
+		name:        "client_dhcp",
+		cliID:       "",
+		cliAddr:     dhcpCliIP,
+		wantNilConf: assert.NotNil,
 	}, {
 		name:        "non_existing_client_id",
 		cliID:       nonExistingClientID,
@@ -1261,6 +1293,16 @@ func TestStorage_CustomUpstreamConfig(t *testing.T) {
 
 		assert.NotEqual(t, conf, updConf)
 	})
+
+	t.Run("same_custom_config", func(t *testing.T) {
+		firstConf := s.CustomUpstreamConfig(existingClientID, existingIP)
+		require.NotNil(t, firstConf)
+
+		secondConf := s.CustomUpstreamConfig(existingClientID, existingIP)
+		require.NotNil(t, secondConf)
+
+		assert.Same(t, firstConf, secondConf)
+	})
 }
 
 func BenchmarkFindParams_Set(b *testing.B) {
diff --git a/internal/client/upstreammanager.go b/internal/client/upstreammanager.go
index b804dbdc..38ad4231 100644
--- a/internal/client/upstreammanager.go
+++ b/internal/client/upstreammanager.go
@@ -138,6 +138,7 @@ func (m *upstreamManager) customUpstreamConfig(uid UID) (proxyConf *proxy.Custom
 
 	proxyConf = newCustomUpstreamConfig(cliConf, m.commonConf)
 	cliConf.proxyConf = proxyConf
+	cliConf.commonConfUpdate = m.confUpdate
 	cliConf.isChanged = false
 
 	return proxyConf