From 0cd6781a9a94bbb0cf1d0bd310fc50d665da4674 Mon Sep 17 00:00:00 2001
From: Simon Zolin <s.zolin@adguard.com>
Date: Tue, 19 Nov 2019 14:09:10 +0300
Subject: [PATCH] * QueryLog.Add() now receives net.IP, not net.Addr

---
 dnsforward/dnsforward.go | 13 ++++++++++++-
 querylog/qlog.go         | 16 ++--------------
 querylog/querylog.go     |  2 +-
 3 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/dnsforward/dnsforward.go b/dnsforward/dnsforward.go
index 3ecd80d5..cb86c235 100644
--- a/dnsforward/dnsforward.go
+++ b/dnsforward/dnsforward.go
@@ -462,7 +462,7 @@ func (s *Server) handleDNSRequest(p *proxy.Proxy, d *proxy.DNSContext) error {
 		if d.Upstream != nil {
 			upstreamAddr = d.Upstream.Address()
 		}
-		s.queryLog.Add(msg, d.Res, res, elapsed, d.Addr, upstreamAddr)
+		s.queryLog.Add(msg, d.Res, res, elapsed, getIP(d.Addr), upstreamAddr)
 	}
 
 	s.updateStats(d, elapsed, *res)
@@ -471,6 +471,17 @@ func (s *Server) handleDNSRequest(p *proxy.Proxy, d *proxy.DNSContext) error {
 	return nil
 }
 
+// Get IP address from net.Addr
+func getIP(addr net.Addr) net.IP {
+	switch addr := addr.(type) {
+	case *net.UDPAddr:
+		return addr.IP
+	case *net.TCPAddr:
+		return addr.IP
+	}
+	return nil
+}
+
 func (s *Server) updateStats(d *proxy.DNSContext, elapsed time.Duration, res dnsfilter.Result) {
 	if s.stats == nil {
 		return
diff --git a/querylog/qlog.go b/querylog/qlog.go
index e53d685c..151b16ff 100644
--- a/querylog/qlog.go
+++ b/querylog/qlog.go
@@ -107,18 +107,7 @@ type logEntry struct {
 	Upstream string `json:",omitempty"` // if empty, means it was cached
 }
 
-// getIPString is a helper function that extracts IP address from net.Addr
-func getIPString(addr net.Addr) string {
-	switch addr := addr.(type) {
-	case *net.UDPAddr:
-		return addr.IP.String()
-	case *net.TCPAddr:
-		return addr.IP.String()
-	}
-	return ""
-}
-
-func (l *queryLog) Add(question *dns.Msg, answer *dns.Msg, result *dnsfilter.Result, elapsed time.Duration, addr net.Addr, upstream string) {
+func (l *queryLog) Add(question *dns.Msg, answer *dns.Msg, result *dnsfilter.Result, elapsed time.Duration, ip net.IP, upstream string) {
 	if !l.conf.Enabled {
 		return
 	}
@@ -130,7 +119,6 @@ func (l *queryLog) Add(question *dns.Msg, answer *dns.Msg, result *dnsfilter.Res
 
 	var a []byte
 	var err error
-	ip := getIPString(addr)
 
 	if answer != nil {
 		a, err = answer.Pack()
@@ -146,7 +134,7 @@ func (l *queryLog) Add(question *dns.Msg, answer *dns.Msg, result *dnsfilter.Res
 
 	now := time.Now()
 	entry := logEntry{
-		IP:   ip,
+		IP:   ip.String(),
 		Time: now,
 
 		Answer:   a,
diff --git a/querylog/querylog.go b/querylog/querylog.go
index 26bd55a0..5158a211 100644
--- a/querylog/querylog.go
+++ b/querylog/querylog.go
@@ -21,7 +21,7 @@ type QueryLog interface {
 	Close()
 
 	// Add a log entry
-	Add(question *dns.Msg, answer *dns.Msg, result *dnsfilter.Result, elapsed time.Duration, addr net.Addr, upstream string)
+	Add(question *dns.Msg, answer *dns.Msg, result *dnsfilter.Result, elapsed time.Duration, ip net.IP, upstream string)
 
 	// WriteDiskConfig - write configuration
 	WriteDiskConfig(dc *DiskConfig)