From df7492e9de063e3c1b582e9f014554bc6cba7e90 Mon Sep 17 00:00:00 2001
From: Stanislav Chzhen <s.chzhen@adguard.com>
Date: Tue, 24 Sep 2024 19:06:37 +0300
Subject: [PATCH] all: imp code

---
 internal/client/storage.go      | 21 +++++++++++++++++----
 internal/client/storage_test.go |  3 ++-
 internal/home/clients.go        |  7 +------
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/internal/client/storage.go b/internal/client/storage.go
index 473acf76..4d65de48 100644
--- a/internal/client/storage.go
+++ b/internal/client/storage.go
@@ -83,8 +83,8 @@ type HostsContainer interface {
 
 // StorageConfig is the client storage configuration structure.
 type StorageConfig struct {
-	// DHCP is used to update [SourceDHCP] runtime client information.  It must
-	// not be nil.
+	// DHCP is used to match IPs against MACs of persistent clients and update
+	// [SourceDHCP] runtime client information.  It must not be nil.
 	DHCP DHCP
 
 	// EtcHosts is used to update [SourceHostsFile] runtime client information.
@@ -100,6 +100,10 @@ type StorageConfig struct {
 	// ARPClientsUpdatePeriod defines how often [SourceARP] runtime client
 	// information is updated.
 	ARPClientsUpdatePeriod time.Duration
+
+	// RuntimeSourceDHCP specifies whether to update [SourceDHCP] information
+	// of runtime clients.
+	RuntimeSourceDHCP bool
 }
 
 // Storage contains information about persistent and runtime clients.
@@ -131,6 +135,10 @@ type Storage struct {
 	// arpClientsUpdatePeriod defines how often [SourceARP] runtime client
 	// information is updated.  It must be greater than zero.
 	arpClientsUpdatePeriod time.Duration
+
+	// runtimeSourceDHCP specifies whether to update [SourceDHCP] information
+	// of runtime clients.
+	runtimeSourceDHCP bool
 }
 
 // NewStorage returns initialized client storage.  conf must not be nil.
@@ -143,8 +151,9 @@ func NewStorage(conf *StorageConfig) (s *Storage, err error) {
 		dhcp:                   conf.DHCP,
 		etcHosts:               conf.EtcHosts,
 		arpDB:                  conf.ARPDB,
-		arpClientsUpdatePeriod: conf.ARPClientsUpdatePeriod,
 		done:                   make(chan struct{}),
+		arpClientsUpdatePeriod: conf.ARPClientsUpdatePeriod,
+		runtimeSourceDHCP:      conf.RuntimeSourceDHCP,
 	}
 
 	for i, p := range conf.InitialClients {
@@ -306,7 +315,7 @@ func (s *Storage) UpdateAddress(ip netip.Addr, host string, info *whois.Info) {
 
 // UpdateDHCP updates [SourceDHCP] runtime client information.
 func (s *Storage) UpdateDHCP() {
-	if s.dhcp == nil {
+	if s.dhcp == nil || !s.runtimeSourceDHCP {
 		return
 	}
 
@@ -545,6 +554,10 @@ func (s *Storage) ClientRuntime(ip netip.Addr) (rc *Runtime) {
 		return rc.clone()
 	}
 
+	if !s.runtimeSourceDHCP {
+		return nil
+	}
+
 	host := s.dhcp.HostByIP(ip)
 	if host == "" {
 		return nil
diff --git a/internal/client/storage_test.go b/internal/client/storage_test.go
index c5f09c22..c8773ce2 100644
--- a/internal/client/storage_test.go
+++ b/internal/client/storage_test.go
@@ -364,7 +364,8 @@ func TestClientsDHCP(t *testing.T) {
 	}
 
 	storage, err := client.NewStorage(&client.StorageConfig{
-		DHCP: d,
+		DHCP:              d,
+		RuntimeSourceDHCP: true,
 	})
 	require.NoError(t, err)
 
diff --git a/internal/home/clients.go b/internal/home/clients.go
index d7d4e628..63f30dd4 100644
--- a/internal/home/clients.go
+++ b/internal/home/clients.go
@@ -27,9 +27,6 @@ type clientsContainer struct {
 	// storage stores information about persistent clients.
 	storage *client.Storage
 
-	// dhcp is the DHCP service implementation.
-	dhcp client.DHCP
-
 	// clientChecker checks if a client is blocked by the current access
 	// settings.
 	clientChecker BlockedClientChecker
@@ -75,9 +72,6 @@ func (clients *clientsContainer) Init(
 		return errors.Error("clients container already initialized")
 	}
 
-	// TODO(e.burkov):  Use [dhcpsvc] implementation when it's ready.
-	clients.dhcp = dhcpServer
-
 	confClients := make([]*client.Persistent, 0, len(objects))
 	for i, o := range objects {
 		var p *client.Persistent
@@ -106,6 +100,7 @@ func (clients *clientsContainer) Init(
 		EtcHosts:               EtcHosts,
 		ARPDB:                  arpDB,
 		ARPClientsUpdatePeriod: arpClientsUpdatePeriod,
+		RuntimeSourceDHCP:      config.Clients.Sources.DHCP,
 	})
 	if err != nil {
 		return fmt.Errorf("init client storage: %w", err)