From 4795a2c8fbf49cede106641e97892c008dccd2e9 Mon Sep 17 00:00:00 2001
From: Stanislav Chzhen <s.chzhen@adguard.com>
Date: Tue, 21 Jan 2025 17:48:00 +0300
Subject: [PATCH] all: query statistics

---
 go.mod                       |  3 ++-
 go.sum                       |  4 ++--
 internal/dnsforward/stats.go | 43 +++++++++++++++++++++++++++++++-----
 3 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/go.mod b/go.mod
index 1ed09910..d55a5452 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,8 @@ module github.com/AdguardTeam/AdGuardHome
 go 1.23.5
 
 require (
-	github.com/AdguardTeam/dnsproxy v0.73.5
+	// TODO!!
+	github.com/AdguardTeam/dnsproxy v0.74.2-0.20250116174805-966cabfa8953
 	github.com/AdguardTeam/golibs v0.31.0
 	github.com/AdguardTeam/urlfilter v0.20.0
 	github.com/NYTimes/gziphandler v1.1.1
diff --git a/go.sum b/go.sum
index ed6ca788..18157b2a 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,5 @@
-github.com/AdguardTeam/dnsproxy v0.73.5 h1:3EiVaTfmuW6PcJGbqloR6ZdHACsrYkm9z0eH8ZQTZnQ=
-github.com/AdguardTeam/dnsproxy v0.73.5/go.mod h1:Oqw+k7LyjDObfYzXYCkpgtirbzbUrmotr92jrb3N09I=
+github.com/AdguardTeam/dnsproxy v0.74.2-0.20250116174805-966cabfa8953 h1:oWKRUtLrKqUO0g3Vh/uEvuxH4wpqh5mV2r7fwqmABF8=
+github.com/AdguardTeam/dnsproxy v0.74.2-0.20250116174805-966cabfa8953/go.mod h1:Oqw+k7LyjDObfYzXYCkpgtirbzbUrmotr92jrb3N09I=
 github.com/AdguardTeam/golibs v0.31.0 h1:Z0oPfLTLw6iZmpE58dePy2Bel0MaX+lnDwtFEE5EmIo=
 github.com/AdguardTeam/golibs v0.31.0/go.mod h1:wIkZ9o2UnppeW6/YD7yJB71dYbMhiuC1Fh/I2ElW7GQ=
 github.com/AdguardTeam/urlfilter v0.20.0 h1:X32qiuVCVd8WDYCEsbdZKfXMzwdVqrdulamtUi4rmzs=
diff --git a/internal/dnsforward/stats.go b/internal/dnsforward/stats.go
index ffcbc6ef..622421c4 100644
--- a/internal/dnsforward/stats.go
+++ b/internal/dnsforward/stats.go
@@ -122,9 +122,14 @@ func (s *Server) logQuery(dctx *dnsContext, ip net.IP, processingTime time.Durat
 
 	if pctx.Upstream != nil {
 		p.Upstream = pctx.Upstream.Address()
-	} else if cachedUps := pctx.CachedUpstreamAddr; cachedUps != "" {
-		p.Upstream = pctx.CachedUpstreamAddr
-		p.Cached = true
+	}
+
+	if qs := pctx.QueryStatistics(); qs != nil {
+		ms := qs.Main()
+		if len(ms) == 1 && ms[0].IsCached {
+			p.Upstream = ms[0].Address
+			p.Cached = true
+		}
 	}
 
 	s.queryLog.Add(p)
@@ -138,11 +143,10 @@ func (s *Server) updateStats(dctx *dnsContext, clientIP string, processingTime t
 		Domain:         aghnet.NormalizeDomain(pctx.Req.Question[0].Name),
 		Result:         stats.RNotFiltered,
 		ProcessingTime: processingTime,
-		UpstreamTime:   pctx.QueryDuration,
 	}
 
 	if pctx.Upstream != nil {
-		e.Upstream = pctx.Upstream.Address()
+		e.Upstream, e.UpstreamTime = upstreamDur(pctx)
 	}
 
 	if clientID := dctx.clientID; clientID != "" {
@@ -167,3 +171,32 @@ func (s *Server) updateStats(dctx *dnsContext, clientIP string, processingTime t
 
 	s.stats.Update(e)
 }
+
+// upstreamDur returns the upstream DNS server address and the DNS lookup
+// duration.  If the upstream address is empty, it means the request was served
+// from the cache.
+func upstreamDur(pctx *proxy.DNSContext) (upstream string, dur time.Duration) {
+	if pctx.Upstream == nil {
+		return "", 0
+	}
+
+	qs := pctx.QueryStatistics()
+	if qs == nil {
+		return "", 0
+	}
+
+	addr := pctx.Upstream.Address()
+	for _, u := range qs.Main() {
+		if u.Address == addr {
+			return u.Address, u.QueryDuration
+		}
+	}
+
+	for _, u := range qs.Fallback() {
+		if u.Address == addr {
+			return u.Address, u.QueryDuration
+		}
+	}
+
+	return "", 0
+}