From fb6e66971b1b984147ec400ceaff856e7b5710c7 Mon Sep 17 00:00:00 2001
From: Dimitry Kolyshev <dkolyshev@adguard.com>
Date: Tue, 16 Apr 2024 10:08:51 +0200
Subject: [PATCH] all: safesearch rewrites

---
 internal/dnsforward/filter.go                    | 6 +++---
 internal/dnsforward/process.go                   | 9 ++++-----
 internal/filtering/filtering.go                  | 2 ++
 internal/filtering/safesearch/safesearch_test.go | 5 +++--
 4 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/internal/dnsforward/filter.go b/internal/dnsforward/filter.go
index cdd4bd33..faff248b 100644
--- a/internal/dnsforward/filter.go
+++ b/internal/dnsforward/filter.go
@@ -40,9 +40,6 @@ func (s *Server) filterDNSRequest(dctx *dnsContext) (res *filtering.Result, err
 	// TODO(a.garipov): Make CheckHost return a pointer.
 	res = &resVal
 	switch {
-	case res.IsFiltered && res.CanonName == "":
-		log.Debug("dnsforward: host %q is filtered, reason: %q", host, res.Reason)
-		pctx.Res = s.genDNSFilterMessage(pctx, res)
 	case res.Reason.In(
 		filtering.Rewritten,
 		filtering.RewrittenRule,
@@ -53,6 +50,9 @@ func (s *Server) filterDNSRequest(dctx *dnsContext) (res *filtering.Result, err
 		// original question is readded in processFilteringAfterResponse.
 		dctx.origQuestion = q
 		req.Question[0].Name = dns.Fqdn(res.CanonName)
+	case res.IsFiltered:
+		log.Debug("dnsforward: host %q is filtered, reason: %q", host, res.Reason)
+		pctx.Res = s.genDNSFilterMessage(pctx, res)
 	case res.Reason.In(filtering.Rewritten, filtering.FilteredSafeSearch):
 		pctx.Res = s.getCNAMEWithIPs(req, res.IPList, res.CanonName)
 	case res.Reason.In(filtering.RewrittenRule, filtering.RewrittenAutoHosts):
diff --git a/internal/dnsforward/process.go b/internal/dnsforward/process.go
index 1dd9e664..7b53fd94 100644
--- a/internal/dnsforward/process.go
+++ b/internal/dnsforward/process.go
@@ -609,11 +609,10 @@ func (s *Server) processFilteringAfterResponse(dctx *dnsContext) (rc resultCode)
 
 		pctx := dctx.proxyCtx
 		pctx.Req.Question[0], pctx.Res.Question[0] = dctx.origQuestion, dctx.origQuestion
-		if len(pctx.Res.Answer) > 0 {
-			rr := s.genAnswerCNAME(pctx.Req, res.CanonName)
-			answer := append([]dns.RR{rr}, pctx.Res.Answer...)
-			pctx.Res.Answer = answer
-		}
+
+		rr := s.genAnswerCNAME(pctx.Req, res.CanonName)
+		answer := append([]dns.RR{rr}, pctx.Res.Answer...)
+		pctx.Res.Answer = answer
 
 		return resultCodeSuccess
 	default:
diff --git a/internal/filtering/filtering.go b/internal/filtering/filtering.go
index 4ea57eea..55404b74 100644
--- a/internal/filtering/filtering.go
+++ b/internal/filtering/filtering.go
@@ -559,6 +559,8 @@ type Result struct {
 	Reason Reason `json:",omitempty"`
 
 	// IsFiltered is true if the request is filtered.
+	//
+	// TODO(d.kolyshev): Get rid of this flag.
 	IsFiltered bool `json:",omitempty"`
 }
 
diff --git a/internal/filtering/safesearch/safesearch_test.go b/internal/filtering/safesearch/safesearch_test.go
index ae62da56..9526c791 100644
--- a/internal/filtering/safesearch/safesearch_test.go
+++ b/internal/filtering/safesearch/safesearch_test.go
@@ -92,8 +92,9 @@ func TestDefault_CheckHost_yandex(t *testing.T) {
 				} else {
 					require.Len(t, res.Rules, 1)
 
-					assert.Equal(t, tc.want, res.Rules[0].IP)
-					assert.Equal(t, rulelist.URLFilterIDSafeSearch, res.Rules[0].FilterListID)
+					rule := res.Rules[0]
+					assert.Equal(t, tc.want, rule.IP)
+					assert.Equal(t, rulelist.URLFilterIDSafeSearch, rule.FilterListID)
 				}
 			}
 		})