From 9f93a21bf6b3aa4bc4015aa840095f9c4694ffe0 Mon Sep 17 00:00:00 2001
From: Ainar Garipov <A.Garipov@AdGuard.COM>
Date: Thu, 13 Jul 2023 18:18:49 +0300
Subject: [PATCH] all: imp client resolving

---
 CHANGELOG.md                                  |   3 +
 internal/dnsforward/config.go                 |   5 +-
 internal/dnsforward/dnsforward.go             |  16 +-
 internal/dnsforward/{dns.go => process.go}    |  33 +++-
 .../{dns_test.go => process_internal_test.go} |   0
 internal/home/clientaddr.go                   | 146 +++++++++++++++
 internal/home/clients.go                      |   2 +-
 internal/home/dns.go                          | 168 +++---------------
 internal/home/home.go                         |  12 +-
 9 files changed, 221 insertions(+), 164 deletions(-)
 rename internal/dnsforward/{dns.go => process.go} (97%)
 rename internal/dnsforward/{dns_test.go => process_internal_test.go} (100%)
 create mode 100644 internal/home/clientaddr.go

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 851caebd..4f2360b4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,8 @@ NOTE: Add new changes BELOW THIS COMMENT.
 
 ### Fixed
 
+- Occasional client information lookup failures leading to DNS resolving getting
+  stuck ([#6006]).
 - `bufio.Scanner: token too long` errors when trying to add filtering-rule lists
   with lines over 1024 bytes long ([#6003]).
 
@@ -34,6 +36,7 @@ NOTE: Add new changes BELOW THIS COMMENT.
   the `Dockerfile`.
 
 [#6003]: https://github.com/AdguardTeam/AdGuardHome/issues/6003
+[#6006]: https://github.com/AdguardTeam/AdGuardHome/issues/6006
 
 <!--
 NOTE: Add new changes ABOVE THIS COMMENT.
diff --git a/internal/dnsforward/config.go b/internal/dnsforward/config.go
index aa3f4808..bcdf0b5d 100644
--- a/internal/dnsforward/config.go
+++ b/internal/dnsforward/config.go
@@ -270,7 +270,10 @@ type ServerConfig struct {
 	UDPListenAddrs []*net.UDPAddr        // UDP listen address
 	TCPListenAddrs []*net.TCPAddr        // TCP listen address
 	UpstreamConfig *proxy.UpstreamConfig // Upstream DNS servers config
-	OnDNSRequest   func(d *proxy.DNSContext)
+
+	// ClientIPs, if not nil, is used to send clients' IP addresses to other
+	// parts of AdGuard Home that may use it for resolving rDNS, WHOIS, etc.
+	ClientIPs chan netip.Addr
 
 	FilteringConfig
 	TLSConfig
diff --git a/internal/dnsforward/dnsforward.go b/internal/dnsforward/dnsforward.go
index 70abd660..992c36bf 100644
--- a/internal/dnsforward/dnsforward.go
+++ b/internal/dnsforward/dnsforward.go
@@ -99,6 +99,10 @@ type Server struct {
 	// must be a valid domain name plus dots on each side.
 	localDomainSuffix string
 
+	// ClientIPs, if not nil, is used to send clients' IP addresses to other
+	// parts of AdGuard Home that may use it for resolving rDNS, WHOIS, etc.
+	clientIPs chan<- netip.Addr
+
 	ipset          ipsetCtx
 	privateNets    netutil.SubnetSet
 	localResolvers *proxy.Proxy
@@ -318,7 +322,8 @@ func (s *Server) Exchange(ip netip.Addr) (host string, err error) {
 			Qclass: dns.ClassINET,
 		}},
 	}
-	ctx := &proxy.DNSContext{
+
+	dctx := &proxy.DNSContext{
 		Proto:     "udp",
 		Req:       req,
 		StartTime: time.Now(),
@@ -336,11 +341,11 @@ func (s *Server) Exchange(ip netip.Addr) (host string, err error) {
 		resolver = s.internalProxy
 	}
 
-	if err = resolver.Resolve(ctx); err != nil {
+	if err = resolver.Resolve(dctx); err != nil {
 		return "", err
 	}
 
-	return hostFromPTR(ctx.Res)
+	return hostFromPTR(dctx.Res)
 }
 
 // hostFromPTR returns domain name from the PTR response or error.
@@ -555,6 +560,8 @@ func (s *Server) Prepare(conf *ServerConfig) (err error) {
 
 	s.recDetector.clear()
 
+	s.clientIPs = s.conf.ClientIPs
+
 	return nil
 }
 
@@ -696,6 +703,9 @@ func (s *Server) Reconfigure(conf *ServerConfig) error {
 	// TODO(a.garipov): This whole piece of API is weird and needs to be remade.
 	if conf == nil {
 		conf = &s.conf
+	} else if s.clientIPs != nil {
+		close(s.clientIPs)
+		s.clientIPs = nil
 	}
 
 	err = s.Prepare(conf)
diff --git a/internal/dnsforward/dns.go b/internal/dnsforward/process.go
similarity index 97%
rename from internal/dnsforward/dns.go
rename to internal/dnsforward/process.go
index 6007801f..4891b144 100644
--- a/internal/dnsforward/dns.go
+++ b/internal/dnsforward/process.go
@@ -30,6 +30,7 @@ type dnsContext struct {
 	setts *filtering.Settings
 
 	result *filtering.Result
+
 	// origResp is the response received from upstream.  It is set when the
 	// response is modified by filters.
 	origResp *dns.Msg
@@ -48,13 +49,13 @@ type dnsContext struct {
 	// clientID is the ClientID from DoH, DoQ, or DoT, if provided.
 	clientID string
 
+	// startTime is the time at which the processing of the request has started.
+	startTime time.Time
+
 	// origQuestion is the question received from the client.  It is set
 	// when the request is modified by rewrites.
 	origQuestion dns.Question
 
-	// startTime is the time at which the processing of the request has started.
-	startTime time.Time
-
 	// protectionEnabled shows if the filtering is enabled, and if the
 	// server's DNS filter is ready.
 	protectionEnabled bool
@@ -177,9 +178,7 @@ func (s *Server) processInitial(dctx *dnsContext) (rc resultCode) {
 		return resultCodeFinish
 	}
 
-	if s.conf.OnDNSRequest != nil {
-		s.conf.OnDNSRequest(pctx)
-	}
+	s.processClientIP(pctx.Addr)
 
 	// Disable Mozilla DoH.
 	//
@@ -218,6 +217,28 @@ func (s *Server) processInitial(dctx *dnsContext) (rc resultCode) {
 	return resultCodeSuccess
 }
 
+// processClientIP sends the client IP address to s.clientIPs, if needed.
+func (s *Server) processClientIP(addr net.Addr) {
+	clientIP := netutil.NetAddrToAddrPort(addr).Addr()
+	if clientIP == (netip.Addr{}) {
+		log.Info("dnsforward: warning: bad client addr %q", addr)
+
+		return
+	}
+
+	// Do not assign s.clientIPs to a local variable to then use, since this
+	// lock also serializes the closure of s.clientIPs.
+	s.serverLock.RLock()
+	defer s.serverLock.RUnlock()
+
+	select {
+	case s.clientIPs <- clientIP:
+		// Go on.
+	default:
+		log.Debug("dnsforward: client ip channel is nil or full; len: %d", len(s.clientIPs))
+	}
+}
+
 func (s *Server) setTableHostToIP(t hostToIPTable) {
 	s.tableHostToIPLock.Lock()
 	defer s.tableHostToIPLock.Unlock()
diff --git a/internal/dnsforward/dns_test.go b/internal/dnsforward/process_internal_test.go
similarity index 100%
rename from internal/dnsforward/dns_test.go
rename to internal/dnsforward/process_internal_test.go
diff --git a/internal/home/clientaddr.go b/internal/home/clientaddr.go
new file mode 100644
index 00000000..d2e4f021
--- /dev/null
+++ b/internal/home/clientaddr.go
@@ -0,0 +1,146 @@
+package home
+
+import (
+	"context"
+	"net/netip"
+	"time"
+
+	"github.com/AdguardTeam/AdGuardHome/internal/rdns"
+	"github.com/AdguardTeam/AdGuardHome/internal/whois"
+	"github.com/AdguardTeam/golibs/log"
+)
+
+// TODO(a.garipov): It is currently hard to add tests for this structure due to
+// strong coupling between it and Context.dnsServer with Context.clients.
+// Resolve this coupling and add proper testing.
+
+// clientAddrProcessor processes incoming client addresses with rDNS and WHOIS,
+// if configured.
+type clientAddrProcessor struct {
+	rdns  rdns.Interface
+	whois whois.Interface
+}
+
+const (
+	// defaultQueueSize is the size of queue of IPs for rDNS and WHOIS
+	// processing.
+	defaultQueueSize = 255
+
+	// defaultCacheSize is the maximum size of the cache for rDNS and WHOIS
+	// processing.  It must be greater than zero.
+	defaultCacheSize = 10_000
+
+	// defaultIPTTL is the Time to Live duration for IP addresses cached by
+	// rDNS and WHOIS.
+	defaultIPTTL = 1 * time.Hour
+)
+
+// newClientAddrProcessor returns a new client address processor.  c must not be
+// nil.
+func newClientAddrProcessor(c *clientSourcesConfig) (p *clientAddrProcessor) {
+	p = &clientAddrProcessor{}
+
+	if c.RDNS {
+		p.rdns = rdns.New(&rdns.Config{
+			Exchanger: Context.dnsServer,
+			CacheSize: defaultCacheSize,
+			CacheTTL:  defaultIPTTL,
+		})
+	} else {
+		p.rdns = rdns.Empty{}
+	}
+
+	if c.WHOIS {
+		// TODO(s.chzhen):  Consider making configurable.
+		const (
+			// defaultTimeout is the timeout for WHOIS requests.
+			defaultTimeout = 5 * time.Second
+
+			// defaultMaxConnReadSize is an upper limit in bytes for reading from a
+			// net.Conn.
+			defaultMaxConnReadSize = 64 * 1024
+
+			// defaultMaxRedirects is the maximum redirects count.
+			defaultMaxRedirects = 5
+
+			// defaultMaxInfoLen is the maximum length of whois.Info fields.
+			defaultMaxInfoLen = 250
+		)
+
+		p.whois = whois.New(&whois.Config{
+			DialContext:     customDialContext,
+			ServerAddr:      whois.DefaultServer,
+			Port:            whois.DefaultPort,
+			Timeout:         defaultTimeout,
+			CacheSize:       defaultCacheSize,
+			MaxConnReadSize: defaultMaxConnReadSize,
+			MaxRedirects:    defaultMaxRedirects,
+			MaxInfoLen:      defaultMaxInfoLen,
+			CacheTTL:        defaultIPTTL,
+		})
+	} else {
+		p.whois = whois.Empty{}
+	}
+
+	return p
+}
+
+// process processes the incoming client IP-address information.  It is intended
+// to be used as a goroutine.
+func (p *clientAddrProcessor) process(clientIPs <-chan netip.Addr) {
+	defer log.OnPanic("clientAddrProcessor.process")
+
+	log.Info("home: processing client addresses")
+
+	for ip := range clientIPs {
+		p.processRDNS(ip)
+		p.processWHOIS(ip)
+	}
+
+	log.Info("home: finished processing client addresses")
+}
+
+// processRDNS resolves the clients' IP addresses using reverse DNS.
+func (p *clientAddrProcessor) processRDNS(ip netip.Addr) {
+	start := time.Now()
+	log.Debug("home: processing client %s with rdns", ip)
+	defer func() {
+		log.Debug("home: finished processing client %s with rdns in %s", ip, time.Since(start))
+	}()
+
+	ok := Context.dnsServer.ShouldResolveClient(ip)
+	if !ok {
+		return
+	}
+
+	host, changed := p.rdns.Process(ip)
+	if host == "" || !changed {
+		return
+	}
+
+	ok = Context.clients.AddHost(ip, host, ClientSourceRDNS)
+	if ok {
+		return
+	}
+
+	log.Debug("dns: setting rdns info for client %q: already set with higher priority source", ip)
+}
+
+// processWHOIS looks up the information aobut clients' IP addresses in the
+// WHOIS databases.
+func (p *clientAddrProcessor) processWHOIS(ip netip.Addr) {
+	start := time.Now()
+	log.Debug("home: processing client %s with whois", ip)
+	defer func() {
+		log.Debug("home: finished processing client %s with whois in %s", ip, time.Since(start))
+	}()
+
+	// TODO(s.chzhen):  Move the timeout logic from WHOIS configuration to the
+	// context.
+	info, changed := p.whois.Process(context.Background(), ip)
+	if info == nil || !changed {
+		return
+	}
+
+	Context.clients.setWHOISInfo(ip, info)
+}
diff --git a/internal/home/clients.go b/internal/home/clients.go
index 1cd3add0..0d1010db 100644
--- a/internal/home/clients.go
+++ b/internal/home/clients.go
@@ -141,7 +141,7 @@ func (clients *clientsContainer) handleHostsUpdates() {
 	}
 }
 
-// webHandlersRegistered prevents a [clientsContainer] from regisering its web
+// webHandlersRegistered prevents a [clientsContainer] from registering its web
 // handlers more than once.
 //
 // TODO(a.garipov): Refactor HTTP handler registration logic.
diff --git a/internal/home/dns.go b/internal/home/dns.go
index fbbda423..fd74be8a 100644
--- a/internal/home/dns.go
+++ b/internal/home/dns.go
@@ -17,10 +17,7 @@ import (
 	"github.com/AdguardTeam/AdGuardHome/internal/dnsforward"
 	"github.com/AdguardTeam/AdGuardHome/internal/filtering"
 	"github.com/AdguardTeam/AdGuardHome/internal/querylog"
-	"github.com/AdguardTeam/AdGuardHome/internal/rdns"
 	"github.com/AdguardTeam/AdGuardHome/internal/stats"
-	"github.com/AdguardTeam/AdGuardHome/internal/whois"
-	"github.com/AdguardTeam/dnsproxy/proxy"
 	"github.com/AdguardTeam/golibs/errors"
 	"github.com/AdguardTeam/golibs/log"
 	"github.com/AdguardTeam/golibs/netutil"
@@ -154,134 +151,32 @@ func initDNSServer(
 
 	Context.clients.dnsServer = Context.dnsServer
 
-	dnsConf, err := generateServerConfig(tlsConf, httpReg)
+	dnsConf, err := newServerConfig(tlsConf, httpReg)
 	if err != nil {
 		closeDNSServer()
 
-		return fmt.Errorf("generateServerConfig: %w", err)
+		return fmt.Errorf("newServerConfig: %w", err)
 	}
 
-	err = Context.dnsServer.Prepare(&dnsConf)
+	err = Context.dnsServer.Prepare(dnsConf)
 	if err != nil {
 		closeDNSServer()
 
 		return fmt.Errorf("dnsServer.Prepare: %w", err)
 	}
 
-	initRDNS()
-	initWHOIS()
+	clientIPs := dnsConf.ClientIPs
+	addrProc := newClientAddrProcessor(config.Clients.Sources)
+	go addrProc.process(clientIPs)
+
+	const topClientsNumber = 100
+	for _, ip := range Context.stats.TopClientsIP(topClientsNumber) {
+		clientIPs <- ip
+	}
 
 	return nil
 }
 
-const (
-	// defaultQueueSize is the size of queue of IPs for rDNS and WHOIS
-	// processing.
-	defaultQueueSize = 255
-
-	// defaultCacheSize is the maximum size of the cache for rDNS and WHOIS
-	// processing.  It must be greater than zero.
-	defaultCacheSize = 10_000
-
-	// defaultIPTTL is the Time to Live duration for IP addresses cached by
-	// rDNS and WHOIS.
-	defaultIPTTL = 1 * time.Hour
-)
-
-// initRDNS initializes the rDNS.
-func initRDNS() {
-	Context.rdnsCh = make(chan netip.Addr, defaultQueueSize)
-
-	// TODO(s.chzhen):  Add ability to disable it on dns server configuration
-	// update in [dnsforward] package.
-	r := rdns.New(&rdns.Config{
-		Exchanger: Context.dnsServer,
-		CacheSize: defaultCacheSize,
-		CacheTTL:  defaultIPTTL,
-	})
-
-	go processRDNS(r)
-}
-
-// processRDNS processes reverse DNS lookup queries.  It is intended to be used
-// as a goroutine.
-func processRDNS(r rdns.Interface) {
-	defer log.OnPanic("rdns")
-
-	for ip := range Context.rdnsCh {
-		ok := Context.dnsServer.ShouldResolveClient(ip)
-		if !ok {
-			continue
-		}
-
-		host, changed := r.Process(ip)
-		if host == "" || !changed {
-			continue
-		}
-
-		ok = Context.clients.AddHost(ip, host, ClientSourceRDNS)
-		if ok {
-			continue
-		}
-
-		log.Debug(
-			"dns: can't set rdns info for client %q: already set with higher priority source",
-			ip,
-		)
-	}
-}
-
-// initWHOIS initializes the WHOIS.
-//
-// TODO(s.chzhen):  Consider making configurable.
-func initWHOIS() {
-	const (
-		// defaultTimeout is the timeout for WHOIS requests.
-		defaultTimeout = 5 * time.Second
-
-		// defaultMaxConnReadSize is an upper limit in bytes for reading from
-		// net.Conn.
-		defaultMaxConnReadSize = 64 * 1024
-
-		// defaultMaxRedirects is the maximum redirects count.
-		defaultMaxRedirects = 5
-
-		// defaultMaxInfoLen is the maximum length of whois.Info fields.
-		defaultMaxInfoLen = 250
-	)
-
-	Context.whoisCh = make(chan netip.Addr, defaultQueueSize)
-
-	var w whois.Interface
-
-	if config.Clients.Sources.WHOIS {
-		w = whois.New(&whois.Config{
-			DialContext:     customDialContext,
-			ServerAddr:      whois.DefaultServer,
-			Port:            whois.DefaultPort,
-			Timeout:         defaultTimeout,
-			CacheSize:       defaultCacheSize,
-			MaxConnReadSize: defaultMaxConnReadSize,
-			MaxRedirects:    defaultMaxRedirects,
-			MaxInfoLen:      defaultMaxInfoLen,
-			CacheTTL:        defaultIPTTL,
-		})
-	} else {
-		w = whois.Empty{}
-	}
-
-	go func() {
-		defer log.OnPanic("whois")
-
-		for ip := range Context.whoisCh {
-			info, changed := w.Process(context.Background(), ip)
-			if info != nil && changed {
-				Context.clients.setWHOISInfo(ip, info)
-			}
-		}
-	}()
-}
-
 // parseSubnetSet parses a slice of subnets.  If the slice is empty, it returns
 // a subnet set that matches all locally served networks, see
 // [netutil.IsLocallyServed].
@@ -312,17 +207,6 @@ func isRunning() bool {
 	return Context.dnsServer != nil && Context.dnsServer.IsRunning()
 }
 
-func onDNSRequest(pctx *proxy.DNSContext) {
-	ip := netutil.NetAddrToAddrPort(pctx.Addr).Addr()
-	if ip == (netip.Addr{}) {
-		// This would be quite weird if we get here.
-		return
-	}
-
-	Context.rdnsCh <- ip
-	Context.whoisCh <- ip
-}
-
 func ipsToTCPAddrs(ips []netip.Addr, port int) (tcpAddrs []*net.TCPAddr) {
 	if ips == nil {
 		return nil
@@ -349,19 +233,20 @@ func ipsToUDPAddrs(ips []netip.Addr, port int) (udpAddrs []*net.UDPAddr) {
 	return udpAddrs
 }
 
-func generateServerConfig(
+func newServerConfig(
 	tlsConf *tlsConfigSettings,
 	httpReg aghhttp.RegisterFunc,
-) (newConf dnsforward.ServerConfig, err error) {
+) (newConf *dnsforward.ServerConfig, err error) {
 	dnsConf := config.DNS
 	hosts := aghalg.CoalesceSlice(dnsConf.BindHosts, []netip.Addr{netutil.IPv4Localhost()})
-	newConf = dnsforward.ServerConfig{
+	clientIPs := make(chan netip.Addr, defaultQueueSize)
+	newConf = &dnsforward.ServerConfig{
 		UDPListenAddrs:  ipsToUDPAddrs(hosts, dnsConf.Port),
 		TCPListenAddrs:  ipsToTCPAddrs(hosts, dnsConf.Port),
 		FilteringConfig: dnsConf.FilteringConfig,
 		ConfigModified:  onConfigModified,
 		HTTPRegister:    httpReg,
-		OnDNSRequest:    onDNSRequest,
+		ClientIPs:       clientIPs,
 		UseDNS64:        config.DNS.UseDNS64,
 		DNS64Prefixes:   config.DNS.DNS64Prefixes,
 	}
@@ -385,9 +270,9 @@ func generateServerConfig(
 		if tlsConf.PortDNSCrypt != 0 {
 			newConf.DNSCryptConfig, err = newDNSCrypt(hosts, *tlsConf)
 			if err != nil {
-				// Don't wrap the error, because it's already
-				// wrapped by newDNSCrypt.
-				return dnsforward.ServerConfig{}, err
+				// Don't wrap the error, because it's already wrapped by
+				// newDNSCrypt.
+				return nil, err
 			}
 		}
 	}
@@ -556,31 +441,26 @@ func startDNSServer() error {
 	Context.stats.Start()
 	Context.queryLog.Start()
 
-	const topClientsNumber = 100 // the number of clients to get
-	for _, ip := range Context.stats.TopClientsIP(topClientsNumber) {
-		Context.rdnsCh <- ip
-		Context.whoisCh <- ip
-	}
-
 	return nil
 }
 
 func reconfigureDNSServer() (err error) {
-	var newConf dnsforward.ServerConfig
-
 	tlsConf := &tlsConfigSettings{}
 	Context.tls.WriteDiskConfig(tlsConf)
 
-	newConf, err = generateServerConfig(tlsConf, httpRegister)
+	newConf, err := newServerConfig(tlsConf, httpRegister)
 	if err != nil {
 		return fmt.Errorf("generating forwarding dns server config: %w", err)
 	}
 
-	err = Context.dnsServer.Reconfigure(&newConf)
+	err = Context.dnsServer.Reconfigure(newConf)
 	if err != nil {
 		return fmt.Errorf("starting forwarding dns server: %w", err)
 	}
 
+	addrProc := newClientAddrProcessor(config.Clients.Sources)
+	go addrProc.process(newConf.ClientIPs)
+
 	return nil
 }
 
diff --git a/internal/home/home.go b/internal/home/home.go
index bdc0f86c..09a14c81 100644
--- a/internal/home/home.go
+++ b/internal/home/home.go
@@ -82,12 +82,6 @@ type homeContext struct {
 	client           *http.Client
 	appSignalChannel chan os.Signal // Channel for receiving OS signals by the console app
 
-	// rdnsCh is the channel for receiving IPs for rDNS processing.
-	rdnsCh chan netip.Addr
-
-	// whoisCh is the channel for receiving IPs for WHOIS processing.
-	whoisCh chan netip.Addr
-
 	// tlsCipherIDs are the ID of the cipher suites that AdGuard Home must use.
 	tlsCipherIDs []uint16
 
@@ -634,10 +628,10 @@ func run(opts options, clientBuildFS fs.FS) {
 		Context.tls.start()
 
 		go func() {
-			sErr := startDNSServer()
-			if sErr != nil {
+			startErr := startDNSServer()
+			if startErr != nil {
 				closeDNSServer()
-				fatalOnError(sErr)
+				fatalOnError(startErr)
 			}
 		}()