all: safesearch rewrites

AdguardTeam · Apr 16, 2024 · fb6e669 · fb6e669
1 parent 88add21
commit fb6e669
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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)
 				}
 			}
 		})