From d26c480d03944653bcf78674724da4ad74f2a502 Mon Sep 17 00:00:00 2001
From: Ainar Garipov <a.garipov@adguard.com>
Date: Mon, 19 Jun 2023 15:45:01 +0300
Subject: [PATCH] Pull request 1877: 5913-fix-safesearch-ipv6

Updates #5913.

Squashed commit of the following:

commit a0ab1320ea22dc1b4e2804ef2d14e0091daa6a1e
Author: Ainar Garipov <A.Garipov@AdGuard.COM>
Date:   Mon Jun 19 15:23:44 2023 +0300

    all: fmt; typo

commit 3a2e561c535bbbd2b2eeeaa1a6f423bc123b1a6b
Author: Ainar Garipov <A.Garipov@AdGuard.COM>
Date:   Mon Jun 19 15:16:28 2023 +0300

    all: fix safesearch for ipv6
---
 CHANGELOG.md                                     |  8 ++++----
 internal/dnsforward/msg.go                       | 15 ++++++---------
 internal/filtering/rewrite/storage.go            |  2 +-
 internal/filtering/safesearch/safesearch.go      | 11 +++++++++--
 internal/filtering/safesearch/safesearch_test.go | 11 +++++++++++
 5 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d23c574c..12ac29a6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -82,11 +82,11 @@ In this release, the schema version has changed from 20 to 21.
 
 ### Fixed
 
- - DNSCrypt upstream not resetting the client and resolver information on
-   dialing errors ([#5872]).
+- Safe Search not working with `AAAA` queries for Yandex domains ([#5913]).
 
 [#951]:  https://github.com/AdguardTeam/AdGuardHome/issues/951
 [#1577]: https://github.com/AdguardTeam/AdGuardHome/issues/1577
+[#5913]: https://github.com/AdguardTeam/AdGuardHome/issues/5913
 
 <!--
 NOTE: Add new changes ABOVE THIS COMMENT.
@@ -98,8 +98,8 @@ NOTE: Add new changes ABOVE THIS COMMENT.
 
 ### Fixed
 
- - DNSCrypt upstream not resetting the client and resolver information on
-   dialing errors ([#5872]).
+- DNSCrypt upstream not resetting the client and resolver information on
+  dialing errors ([#5872]).
 
 
 
diff --git a/internal/dnsforward/msg.go b/internal/dnsforward/msg.go
index e6bd178e..36c82f0b 100644
--- a/internal/dnsforward/msg.go
+++ b/internal/dnsforward/msg.go
@@ -57,16 +57,13 @@ func (s *Server) genDNSFilterMessage(
 		return s.genBlockedHost(req, s.conf.SafeBrowsingBlockHost, dctx)
 	case filtering.FilteredParental:
 		return s.genBlockedHost(req, s.conf.ParentalBlockHost, dctx)
+	case filtering.FilteredSafeSearch:
+		// If Safe Search generated the necessary IP addresses, use them.
+		// Otherwise, if there were no errors, there are no addresses for the
+		// requested IP version, so produce a NODATA response.
+		return s.genResponseWithIPs(req, ipsFromRules(res.Rules))
 	default:
-		// If the query was filtered by Safe Search, filtering also must return
-		// the IP addresses that must be used in response.  Return them
-		// regardless of the filtering method.
-		ips := ipsFromRules(res.Rules)
-		if res.Reason == filtering.FilteredSafeSearch && len(ips) > 0 {
-			return s.genResponseWithIPs(req, ips)
-		}
-
-		return s.genForBlockingMode(req, ips)
+		return s.genForBlockingMode(req, ipsFromRules(res.Rules))
 	}
 }
 
diff --git a/internal/filtering/rewrite/storage.go b/internal/filtering/rewrite/storage.go
index 0afb8b6e..31b2e094 100644
--- a/internal/filtering/rewrite/storage.go
+++ b/internal/filtering/rewrite/storage.go
@@ -84,7 +84,7 @@ func (s *DefaultStorage) MatchRequest(dReq *urlfilter.DNSRequest) (rws []*rules.
 		return nil
 	}
 
-	// TODO(a.garipov): Check cnames for cycles on initialisation.
+	// TODO(a.garipov): Check cnames for cycles on initialization.
 	cnames := stringutil.NewSet()
 	host := dReq.Hostname
 	for len(rrules) > 0 && rrules[0].DNSRewrite != nil && rrules[0].DNSRewrite.NewCNAME != "" {
diff --git a/internal/filtering/safesearch/safesearch.go b/internal/filtering/safesearch/safesearch.go
index 880f406c..a1ac6404 100644
--- a/internal/filtering/safesearch/safesearch.go
+++ b/internal/filtering/safesearch/safesearch.go
@@ -203,7 +203,7 @@ func (ss *Default) CheckHost(
 		return res, nil
 	}
 
-	return filtering.Result{}, fmt.Errorf("no ipv4 addresses for %q", host)
+	return filtering.Result{}, fmt.Errorf("no ip addresses for %q", host)
 }
 
 // searchHost looks up DNS rewrites in the internal DNS filtering engine.
@@ -255,7 +255,14 @@ func (ss *Default) newResult(
 
 	host := rewrite.NewCNAME
 	if host == "" {
-		return nil, nil
+		// If there is a rewrite, but it's neither a CNAME one nor one matching
+		// the IP version, then it's a service that only has one type of IP
+		// record but not the other.  Return the empty result to be converted
+		// into a NODATA response.
+		//
+		// TODO(a.garipov): Use the main rewrite result mechanism used in
+		// [dnsforward.Server.filterDNSRequest].
+		return res, nil
 	}
 
 	ss.log(log.DEBUG, "resolving %q", host)
diff --git a/internal/filtering/safesearch/safesearch_test.go b/internal/filtering/safesearch/safesearch_test.go
index 8e1aea2e..b0775c60 100644
--- a/internal/filtering/safesearch/safesearch_test.go
+++ b/internal/filtering/safesearch/safesearch_test.go
@@ -71,6 +71,17 @@ func TestDefault_CheckHost_yandex(t *testing.T) {
 	}
 }
 
+func TestDefault_CheckHost_yandexAAAA(t *testing.T) {
+	conf := testConf
+	ss, err := safesearch.NewDefault(conf, "", testCacheSize, testCacheTTL)
+	require.NoError(t, err)
+
+	res, err := ss.CheckHost("www.yandex.ru", dns.TypeAAAA)
+	require.NoError(t, err)
+
+	assert.True(t, res.IsFiltered)
+}
+
 func TestDefault_CheckHost_google(t *testing.T) {
 	resolver := &aghtest.TestResolver{}
 	ip, _ := resolver.HostToIPs("forcesafesearch.google.com")