Pull request 1803: 5685-fix-safe-search

Updates #5685.

Squashed commit of the following:

commit 5312147abfa0914c896acbf1e88f8c8f1af90f2b
Author: Ainar Garipov <A.Garipov@AdGuard.COM>
Date:   Thu Apr 6 14:09:44 2023 +0300

    safesearch: imp tests, logs

commit 298b5d24ce292c5f83ebe33d1e92329e4b3c1acc
Author: Ainar Garipov <A.Garipov@AdGuard.COM>
Date:   Wed Apr 5 20:36:16 2023 +0300

    safesearch: fix filters, logging

commit 63d6ca5d694d45705473f2f0410e9e0b49cf7346
Author: Ainar Garipov <A.Garipov@AdGuard.COM>
Date:   Wed Apr 5 20:24:47 2023 +0300

    all: dry; fix logs

commit fdbf2f364fd0484b47b3161bf6f4581856fdf47b
Author: Ainar Garipov <A.Garipov@AdGuard.COM>
Date:   Wed Apr 5 20:01:08 2023 +0300

    all: fix safe search update

internal/filtering/safesearch/rules/bing.txt

@@ -1 +1 @@
-|www.bing.com^$dnsrewrite=NOERROR;CNAME;strict.bing.com
+|www.bing.com^$dnsrewrite=NOERROR;CNAME;strict.bing.com

internal/filtering/safesearch/rules/duckduckgo.txt

@@ -1,3 +1,3 @@
 |duckduckgo.com^$dnsrewrite=NOERROR;CNAME;safe.duckduckgo.com
 |start.duckduckgo.com^$dnsrewrite=NOERROR;CNAME;safe.duckduckgo.com
-|www.duckduckgo.com^$dnsrewrite=NOERROR;CNAME;safe.duckduckgo.com
+|www.duckduckgo.com^$dnsrewrite=NOERROR;CNAME;safe.duckduckgo.com

internal/filtering/safesearch/rules/google.txt

@@ -188,4 +188,4 @@
 |www.google.tt^$dnsrewrite=NOERROR;CNAME;forcesafesearch.google.com
 |www.google.vg^$dnsrewrite=NOERROR;CNAME;forcesafesearch.google.com
 |www.google.vu^$dnsrewrite=NOERROR;CNAME;forcesafesearch.google.com
-|www.google.ws^$dnsrewrite=NOERROR;CNAME;forcesafesearch.google.com
+|www.google.ws^$dnsrewrite=NOERROR;CNAME;forcesafesearch.google.com

internal/filtering/safesearch/rules/pixabay.txt

@@ -1 +1 @@
-|pixabay.com^$dnsrewrite=NOERROR;CNAME;safesearch.pixabay.com
+|pixabay.com^$dnsrewrite=NOERROR;CNAME;safesearch.pixabay.com

internal/filtering/safesearch/rules/yandex.txt

@@ -49,4 +49,4 @@
 |yandex.ru^$dnsrewrite=NOERROR;A;213.180.193.56
 |yandex.tj^$dnsrewrite=NOERROR;A;213.180.193.56
 |yandex.tm^$dnsrewrite=NOERROR;A;213.180.193.56
-|yandex.uz^$dnsrewrite=NOERROR;A;213.180.193.56
+|yandex.uz^$dnsrewrite=NOERROR;A;213.180.193.56

internal/filtering/safesearch/rules/youtube.txt

@@ -2,4 +2,4 @@
 |m.youtube.com^$dnsrewrite=NOERROR;CNAME;restrictmoderate.youtube.com
 |youtubei.googleapis.com^$dnsrewrite=NOERROR;CNAME;restrictmoderate.youtube.com
 |youtube.googleapis.com^$dnsrewrite=NOERROR;CNAME;restrictmoderate.youtube.com
-|www.youtube-nocookie.com^$dnsrewrite=NOERROR;CNAME;restrictmoderate.youtube.com
+|www.youtube-nocookie.com^$dnsrewrite=NOERROR;CNAME;restrictmoderate.youtube.com

internal/filtering/safesearch/safesearch.go

@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"net"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/AdguardTeam/AdGuardHome/internal/filtering"
@@ -53,44 +54,85 @@ func isServiceProtected(s filtering.SafeSearchConfig, service Service) (ok bool)
 	}
 }
 
-// DefaultSafeSearch is the default safesearch struct.
-type DefaultSafeSearch struct {
-	engine          *urlfilter.DNSEngine
-	safeSearchCache cache.Cache
-	resolver        filtering.Resolver
-	cacheTime       time.Duration
+// Default is the default safe search filter that uses filtering rules with the
+// dnsrewrite modifier.
+type Default struct {
+	// mu protects engine.
+	mu *sync.RWMutex
+
+	// engine is the filtering engine that contains the DNS rewrite rules.
+	// engine may be nil, which means that this safe search filter is disabled.
+	engine *urlfilter.DNSEngine
+
+	cache     cache.Cache
+	resolver  filtering.Resolver
+	logPrefix string
+	cacheTTL  time.Duration
 }
 
-// NewDefaultSafeSearch returns new safesearch struct.  CacheTime is an element
-// TTL (in minutes).
-func NewDefaultSafeSearch(
+// NewDefault returns an initialized default safe search filter.  name is used
+// for logging.
+func NewDefault(
 	conf filtering.SafeSearchConfig,
+	name string,
 	cacheSize uint,
-	cacheTime time.Duration,
-) (ss *DefaultSafeSearch, err error) {
-	engine, err := newEngine(filtering.SafeSearchListID, conf)
-	if err != nil {
-		return nil, err
-	}
-
+	cacheTTL time.Duration,
+) (ss *Default, err error) {
 	var resolver filtering.Resolver = net.DefaultResolver
 	if conf.CustomResolver != nil {
 		resolver = conf.CustomResolver
 	}
 
-	return &DefaultSafeSearch{
-		engine: engine,
-		safeSearchCache: cache.New(cache.Config{
+	ss = &Default{
+		mu: &sync.RWMutex{},
+
+		cache: cache.New(cache.Config{
 			EnableLRU: true,
 			MaxSize:   cacheSize,
 		}),
-		cacheTime: cacheTime,
-		resolver:  resolver,
-	}, nil
+		resolver: resolver,
+		// Use %s, because the client safe-search names already contain double
+		// quotes.
+		logPrefix: fmt.Sprintf("safesearch %s: ", name),
+		cacheTTL:  cacheTTL,
+	}
+
+	err = ss.resetEngine(filtering.SafeSearchListID, conf)
+	if err != nil {
+		// Don't wrap the error, because it's informative enough as is.
+		return nil, err
+	}
+
+	return ss, nil
 }
 
-// newEngine creates new engine for provided safe search configuration.
-func newEngine(listID int, conf filtering.SafeSearchConfig) (engine *urlfilter.DNSEngine, err error) {
+// log is a helper for logging  that includes the name of the safe search
+// filter.  level must be one of [log.DEBUG], [log.INFO], and [log.ERROR].
+func (ss *Default) log(level log.Level, msg string, args ...any) {
+	switch level {
+	case log.DEBUG:
+		log.Debug(ss.logPrefix+msg, args...)
+	case log.INFO:
+		log.Info(ss.logPrefix+msg, args...)
+	case log.ERROR:
+		log.Error(ss.logPrefix+msg, args...)
+	default:
+		panic(fmt.Errorf("safesearch: unsupported logging level %d", level))
+	}
+}
+
+// resetEngine creates new engine for provided safe search configuration and
+// sets it in ss.
+func (ss *Default) resetEngine(
+	listID int,
+	conf filtering.SafeSearchConfig,
+) (err error) {
+	if !conf.Enabled {
+		ss.log(log.INFO, "disabled")
+
+		return nil
+	}
+
 	var sb strings.Builder
 	for service, serviceRules := range safeSearchRules {
 		if isServiceProtected(conf, service) {
@@ -106,20 +148,73 @@ func newEngine(listID int, conf filtering.SafeSearchConfig) (engine *urlfilter.D
 
 	rs, err := filterlist.NewRuleStorage([]filterlist.RuleList{strList})
 	if err != nil {
-		return nil, fmt.Errorf("creating rule storage: %w", err)
+		return fmt.Errorf("creating rule storage: %w", err)
 	}
 
-	engine = urlfilter.NewDNSEngine(rs)
-	log.Info("safesearch: filter %d: reset %d rules", listID, engine.RulesCount)
+	ss.engine = urlfilter.NewDNSEngine(rs)
 
-	return engine, nil
+	ss.log(log.INFO, "reset %d rules", ss.engine.RulesCount)
+
+	return nil
 }
 
 // type check
-var _ filtering.SafeSearch = (*DefaultSafeSearch)(nil)
+var _ filtering.SafeSearch = (*Default)(nil)
+
+// CheckHost implements the [filtering.SafeSearch] interface for
+// *DefaultSafeSearch.
+func (ss *Default) CheckHost(
+	host string,
+	qtype rules.RRType,
+) (res filtering.Result, err error) {
+	start := time.Now()
+	defer func() {
+		ss.log(log.DEBUG, "lookup for %q finished in %s", host, time.Since(start))
+	}()
+
+	if qtype != dns.TypeA && qtype != dns.TypeAAAA {
+		return filtering.Result{}, fmt.Errorf("unsupported question type %s", dns.Type(qtype))
+	}
+
+	// Check cache. Return cached result if it was found
+	cachedValue, isFound := ss.getCachedResult(host, qtype)
+	if isFound {
+		ss.log(log.DEBUG, "found in cache: %q", host)
+
+		return cachedValue, nil
+	}
+
+	rewrite := ss.searchHost(host, qtype)
+	if rewrite == nil {
+		return filtering.Result{}, nil
+	}
+
+	fltRes, err := ss.newResult(rewrite, qtype)
+	if err != nil {
+		ss.log(log.DEBUG, "looking up addresses for %q: %s", host, err)
+
+		return filtering.Result{}, err
+	}
+
+	if fltRes != nil {
+		res = *fltRes
+		ss.setCacheResult(host, qtype, res)
+
+		return res, nil
+	}
+
+	return filtering.Result{}, fmt.Errorf("no ipv4 addresses for %q", host)
+}
+
+// searchHost looks up DNS rewrites in the internal DNS filtering engine.
+func (ss *Default) searchHost(host string, qtype rules.RRType) (res *rules.DNSRewrite) {
+	ss.mu.RLock()
+	defer ss.mu.RUnlock()
+
+	if ss.engine == nil {
+		return nil
+	}
 
-// SearchHost implements the [filtering.SafeSearch] interface for *DefaultSafeSearch.
-func (ss *DefaultSafeSearch) SearchHost(host string, qtype uint16) (res *rules.DNSRewrite) {
 	r, _ := ss.engine.MatchRequest(&urlfilter.DNSRequest{
 		Hostname: strings.ToLower(host),
 		DNSType:  qtype,
@@ -133,51 +228,11 @@ func (ss *DefaultSafeSearch) SearchHost(host string, qtype uint16) (res *rules.D
 	return nil
 }
 
-// CheckHost implements the [filtering.SafeSearch] interface for
-// *DefaultSafeSearch.
-func (ss *DefaultSafeSearch) CheckHost(
-	host string,
-	qtype uint16,
-) (res filtering.Result, err error) {
-	if log.GetLevel() >= log.DEBUG {
-		timer := log.StartTimer()
-		defer timer.LogElapsed("safesearch: lookup for %s", host)
-	}
-
-	// Check cache. Return cached result if it was found
-	cachedValue, isFound := ss.getCachedResult(host)
-	if isFound {
-		log.Debug("safesearch: found in cache: %s", host)
-
-		return cachedValue, nil
-	}
-
-	rewrite := ss.SearchHost(host, qtype)
-	if rewrite == nil {
-		return filtering.Result{}, nil
-	}
-
-	dRes, err := ss.newResult(rewrite, qtype)
-	if err != nil {
-		log.Debug("safesearch: failed to lookup addresses for %s: %s", host, err)
-
-		return filtering.Result{}, err
-	}
-
-	if dRes != nil {
-		res = *dRes
-		ss.setCacheResult(host, res)
-
-		return res, nil
-	}
-
-	return filtering.Result{}, fmt.Errorf("no ipv4 addresses in safe search response for %s", host)
-}
-
-// newResult creates Result object from rewrite rule.
-func (ss *DefaultSafeSearch) newResult(
+// newResult creates Result object from rewrite rule.  qtype must be either
+// [dns.TypeA] or [dns.TypeAAAA].
+func (ss *Default) newResult(
 	rewrite *rules.DNSRewrite,
-	qtype uint16,
+	qtype rules.RRType,
 ) (res *filtering.Result, err error) {
 	res = &filtering.Result{
 		Rules: []*filtering.ResultRule{{
@@ -187,7 +242,7 @@ func (ss *DefaultSafeSearch) newResult(
 		IsFiltered: true,
 	}
 
-	if rewrite.RRType == qtype && (qtype == dns.TypeA || qtype == dns.TypeAAAA) {
+	if rewrite.RRType == qtype {
 		ip, ok := rewrite.Value.(net.IP)
 		if !ok || ip == nil {
 			return nil, nil
@@ -198,17 +253,25 @@ func (ss *DefaultSafeSearch) newResult(
 		return res, nil
 	}
 
-	if rewrite.NewCNAME == "" {
+	host := rewrite.NewCNAME
+	if host == "" {
 		return nil, nil
 	}
 
-	ips, err := ss.resolver.LookupIP(context.Background(), "ip", rewrite.NewCNAME)
+	ss.log(log.DEBUG, "resolving %q", host)
+
+	ips, err := ss.resolver.LookupIP(context.Background(), qtypeToProto(qtype), host)
 	if err != nil {
 		return nil, err
 	}
 
+	ss.log(log.DEBUG, "resolved %s", ips)
+
 	for _, ip := range ips {
-		if ip = ip.To4(); ip == nil {
+		// TODO(a.garipov): Remove this filtering once the resolver we use
+		// actually learns about network.
+		ip = fitToProto(ip, qtype)
+		if ip == nil {
 			continue
 		}
 
@@ -220,38 +283,71 @@ func (ss *DefaultSafeSearch) newResult(
 	return nil, nil
 }
 
-// setCacheResult stores data in cache for host.
-func (ss *DefaultSafeSearch) setCacheResult(host string, res filtering.Result) {
-	expire := uint32(time.Now().Add(ss.cacheTime).Unix())
+// qtypeToProto returns "ip4" for [dns.TypeA] and "ip6" for [dns.TypeAAAA].
+// It panics for other types.
+func qtypeToProto(qtype rules.RRType) (proto string) {
+	switch qtype {
+	case dns.TypeA:
+		return "ip4"
+	case dns.TypeAAAA:
+		return "ip6"
+	default:
+		panic(fmt.Errorf("safesearch: unsupported question type %s", dns.Type(qtype)))
+	}
+}
+
+// fitToProto returns a non-nil IP address if ip is the correct protocol version
+// for qtype.  qtype is expected to be either [dns.TypeA] or [dns.TypeAAAA].
+func fitToProto(ip net.IP, qtype rules.RRType) (res net.IP) {
+	ip4 := ip.To4()
+	if qtype == dns.TypeA {
+		return ip4
+	}
+
+	if ip4 == nil {
+		return ip
+	}
+
+	return nil
+}
+
+// setCacheResult stores data in cache for host.  qtype is expected to be either
+// [dns.TypeA] or [dns.TypeAAAA].
+func (ss *Default) setCacheResult(host string, qtype rules.RRType, res filtering.Result) {
+	expire := uint32(time.Now().Add(ss.cacheTTL).Unix())
 	exp := make([]byte, 4)
 	binary.BigEndian.PutUint32(exp, expire)
 	buf := bytes.NewBuffer(exp)
 
 	err := gob.NewEncoder(buf).Encode(res)
 	if err != nil {
-		log.Error("safesearch: cache encoding: %s", err)
+		ss.log(log.ERROR, "cache encoding: %s", err)
 
 		return
 	}
 
 	val := buf.Bytes()
-	_ = ss.safeSearchCache.Set([]byte(host), val)
+	_ = ss.cache.Set([]byte(dns.Type(qtype).String()+" "+host), val)
 
-	log.Debug("safesearch: stored in cache: %s (%d bytes)", host, len(val))
+	ss.log(log.DEBUG, "stored in cache: %q, %d bytes", host, len(val))
 }
 
-// getCachedResult returns stored data from cache for host.
-func (ss *DefaultSafeSearch) getCachedResult(host string) (res filtering.Result, ok bool) {
+// getCachedResult returns stored data from cache for host.  qtype is expected
+// to be either [dns.TypeA] or [dns.TypeAAAA].
+func (ss *Default) getCachedResult(
+	host string,
+	qtype rules.RRType,
+) (res filtering.Result, ok bool) {
 	res = filtering.Result{}
 
-	data := ss.safeSearchCache.Get([]byte(host))
+	data := ss.cache.Get([]byte(dns.Type(qtype).String() + " " + host))
 	if data == nil {
 		return res, false
 	}
 
 	exp := binary.BigEndian.Uint32(data[:4])
 	if exp <= uint32(time.Now().Unix()) {
-		ss.safeSearchCache.Del([]byte(host))
+		ss.cache.Del([]byte(host))
 
 		return res, false
 	}
@@ -260,10 +356,27 @@ func (ss *DefaultSafeSearch) getCachedResult(host string) (res filtering.Result,
 
 	err := gob.NewDecoder(buf).Decode(&res)
 	if err != nil {
-		log.Debug("safesearch: cache decoding: %s", err)
+		ss.log(log.ERROR, "cache decoding: %s", err)
 
 		return filtering.Result{}, false
 	}
 
 	return res, true
 }
+
+// Update implements the [filtering.SafeSearch] interface for *Default.  Update
+// ignores the CustomResolver and Enabled fields.
+func (ss *Default) Update(conf filtering.SafeSearchConfig) (err error) {
+	ss.mu.Lock()
+	defer ss.mu.Unlock()
+
+	err = ss.resetEngine(filtering.SafeSearchListID, conf)
+	if err != nil {
+		// Don't wrap the error, because it's informative enough as is.
+		return err
+	}
+
+	ss.cache.Clear()
+
+	return nil
+}

internal/filtering/safesearch/safesearch_internal_test.go

@@ -0,0 +1,137 @@
+package safesearch
+
+import (
+	"context"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/AdguardTeam/AdGuardHome/internal/aghtest"
+	"github.com/AdguardTeam/AdGuardHome/internal/filtering"
+	"github.com/AdguardTeam/urlfilter/rules"
+	"github.com/miekg/dns"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TODO(a.garipov): Move as much of this as possible into proper external tests.
+
+const (
+	// TODO(a.garipov): Add IPv6 tests.
+	testQType     = dns.TypeA
+	testCacheSize = 5000
+	testCacheTTL  = 30 * time.Minute
+)
+
+var defaultSafeSearchConf = filtering.SafeSearchConfig{
+	Enabled:    true,
+	Bing:       true,
+	DuckDuckGo: true,
+	Google:     true,
+	Pixabay:    true,
+	Yandex:     true,
+	YouTube:    true,
+}
+
+var yandexIP = net.IPv4(213, 180, 193, 56)
+
+func newForTest(t testing.TB, ssConf filtering.SafeSearchConfig) (ss *Default) {
+	ss, err := NewDefault(ssConf, "", testCacheSize, testCacheTTL)
+	require.NoError(t, err)
+
+	return ss
+}
+
+func TestSafeSearch(t *testing.T) {
+	ss := newForTest(t, defaultSafeSearchConf)
+	val := ss.searchHost("www.google.com", testQType)
+
+	assert.Equal(t, &rules.DNSRewrite{NewCNAME: "forcesafesearch.google.com"}, val)
+}
+
+func TestSafeSearchCacheYandex(t *testing.T) {
+	const domain = "yandex.ru"
+
+	ss := newForTest(t, filtering.SafeSearchConfig{Enabled: false})
+
+	// Check host with disabled safesearch.
+	res, err := ss.CheckHost(domain, testQType)
+	require.NoError(t, err)
+
+	assert.False(t, res.IsFiltered)
+	assert.Empty(t, res.Rules)
+
+	ss = newForTest(t, defaultSafeSearchConf)
+	res, err = ss.CheckHost(domain, testQType)
+	require.NoError(t, err)
+
+	// For yandex we already know valid IP.
+	require.Len(t, res.Rules, 1)
+
+	assert.Equal(t, res.Rules[0].IP, yandexIP)
+
+	// Check cache.
+	cachedValue, isFound := ss.getCachedResult(domain, testQType)
+	require.True(t, isFound)
+	require.Len(t, cachedValue.Rules, 1)
+
+	assert.Equal(t, cachedValue.Rules[0].IP, yandexIP)
+}
+
+func TestSafeSearchCacheGoogle(t *testing.T) {
+	const domain = "www.google.ru"
+
+	ss := newForTest(t, filtering.SafeSearchConfig{Enabled: false})
+
+	res, err := ss.CheckHost(domain, testQType)
+	require.NoError(t, err)
+
+	assert.False(t, res.IsFiltered)
+	assert.Empty(t, res.Rules)
+
+	resolver := &aghtest.TestResolver{}
+	ss = newForTest(t, defaultSafeSearchConf)
+	ss.resolver = resolver
+
+	// Lookup for safesearch domain.
+	rewrite := ss.searchHost(domain, testQType)
+
+	ips, err := resolver.LookupIP(context.Background(), "ip", rewrite.NewCNAME)
+	require.NoError(t, err)
+
+	var foundIP net.IP
+	for _, ip := range ips {
+		if ip.To4() != nil {
+			foundIP = ip
+
+			break
+		}
+	}
+
+	res, err = ss.CheckHost(domain, testQType)
+	require.NoError(t, err)
+	require.Len(t, res.Rules, 1)
+
+	assert.True(t, res.Rules[0].IP.Equal(foundIP))
+
+	// Check cache.
+	cachedValue, isFound := ss.getCachedResult(domain, testQType)
+	require.True(t, isFound)
+	require.Len(t, cachedValue.Rules, 1)
+
+	assert.True(t, cachedValue.Rules[0].IP.Equal(foundIP))
+}
+
+const googleHost = "www.google.com"
+
+var dnsRewriteSink *rules.DNSRewrite
+
+func BenchmarkSafeSearch(b *testing.B) {
+	ss := newForTest(b, defaultSafeSearchConf)
+
+	for n := 0; n < b.N; n++ {
+		dnsRewriteSink = ss.searchHost(googleHost, testQType)
+	}
+
+	assert.Equal(b, "forcesafesearch.google.com", dnsRewriteSink.NewCNAME)
+}

internal/filtering/safesearch/safesearch_test.go

@@ -1,26 +1,37 @@
-package safesearch
+package safesearch_test
 
 import (
-	"context"
 	"net"
 	"testing"
 	"time"
 
 	"github.com/AdguardTeam/AdGuardHome/internal/aghtest"
 	"github.com/AdguardTeam/AdGuardHome/internal/filtering"
-	"github.com/AdguardTeam/urlfilter/rules"
+	"github.com/AdguardTeam/AdGuardHome/internal/filtering/safesearch"
+	"github.com/AdguardTeam/golibs/testutil"
 	"github.com/miekg/dns"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+func TestMain(m *testing.M) {
+	testutil.DiscardLogOutput(m)
+}
+
+// Common test constants.
 const (
-	safeSearchCacheSize = 5000
-	cacheTime           = 30 * time.Minute
+	// TODO(a.garipov): Add IPv6 tests.
+	testQType     = dns.TypeA
+	testCacheSize = 5000
+	testCacheTTL  = 30 * time.Minute
 )
 
-var defaultSafeSearchConf = filtering.SafeSearchConfig{
-	Enabled:    true,
+// testConf is the default safe search configuration for tests.
+var testConf = filtering.SafeSearchConfig{
+	CustomResolver: nil,
+
+	Enabled: true,
+
 	Bing:       true,
 	DuckDuckGo: true,
 	Google:     true,
@@ -29,25 +40,15 @@ var defaultSafeSearchConf = filtering.SafeSearchConfig{
 	YouTube:    true,
 }
 
+// yandexIP is the expected IP address of Yandex safe search results.  Keep in
+// sync with the rules data.
 var yandexIP = net.IPv4(213, 180, 193, 56)
 
-func newForTest(t testing.TB, ssConf filtering.SafeSearchConfig) (ss *DefaultSafeSearch) {
-	ss, err := NewDefaultSafeSearch(ssConf, safeSearchCacheSize, cacheTime)
+func TestDefault_CheckHost_yandex(t *testing.T) {
+	conf := testConf
+	ss, err := safesearch.NewDefault(conf, "", testCacheSize, testCacheTTL)
 	require.NoError(t, err)
 
-	return ss
-}
-
-func TestSafeSearch(t *testing.T) {
-	ss := newForTest(t, defaultSafeSearchConf)
-	val := ss.SearchHost("www.google.com", dns.TypeA)
-
-	assert.Equal(t, &rules.DNSRewrite{NewCNAME: "forcesafesearch.google.com"}, val)
-}
-
-func TestCheckHostSafeSearchYandex(t *testing.T) {
-	ss := newForTest(t, defaultSafeSearchConf)
-
 	// Check host for each domain.
 	for _, host := range []string{
 		"yandex.ru",
@@ -57,7 +58,8 @@ func TestCheckHostSafeSearchYandex(t *testing.T) {
 		"yandex.kz",
 		"www.yandex.com",
 	} {
-		res, err := ss.CheckHost(host, dns.TypeA)
+		var res filtering.Result
+		res, err = ss.CheckHost(host, testQType)
 		require.NoError(t, err)
 
 		assert.True(t, res.IsFiltered)
@@ -69,12 +71,14 @@ func TestCheckHostSafeSearchYandex(t *testing.T) {
 	}
 }
 
-func TestCheckHostSafeSearchGoogle(t *testing.T) {
+func TestDefault_CheckHost_google(t *testing.T) {
 	resolver := &aghtest.TestResolver{}
 	ip, _ := resolver.HostToIPs("forcesafesearch.google.com")
 
-	ss := newForTest(t, defaultSafeSearchConf)
-	ss.resolver = resolver
+	conf := testConf
+	conf.CustomResolver = resolver
+	ss, err := safesearch.NewDefault(conf, "", testCacheSize, testCacheTTL)
+	require.NoError(t, err)
 
 	// Check host for each domain.
 	for _, host := range []string{
@@ -87,7 +91,8 @@ func TestCheckHostSafeSearchGoogle(t *testing.T) {
 		"www.google.je",
 	} {
 		t.Run(host, func(t *testing.T) {
-			res, err := ss.CheckHost(host, dns.TypeA)
+			var res filtering.Result
+			res, err = ss.CheckHost(host, testQType)
 			require.NoError(t, err)
 
 			assert.True(t, res.IsFiltered)
@@ -100,103 +105,35 @@ func TestCheckHostSafeSearchGoogle(t *testing.T) {
 	}
 }
 
-func TestSafeSearchCacheYandex(t *testing.T) {
-	const domain = "yandex.ru"
-
-	ss := newForTest(t, filtering.SafeSearchConfig{Enabled: false})
-
-	// Check host with disabled safesearch.
-	res, err := ss.CheckHost(domain, dns.TypeA)
+func TestDefault_Update(t *testing.T) {
+	conf := testConf
+	ss, err := safesearch.NewDefault(conf, "", testCacheSize, testCacheTTL)
 	require.NoError(t, err)
 
-	assert.False(t, res.IsFiltered)
-	assert.Empty(t, res.Rules)
-
-	ss = newForTest(t, defaultSafeSearchConf)
-	res, err = ss.CheckHost(domain, dns.TypeA)
+	res, err := ss.CheckHost("www.yandex.com", testQType)
 	require.NoError(t, err)
 
-	// For yandex we already know valid IP.
-	require.Len(t, res.Rules, 1)
+	assert.True(t, res.IsFiltered)
 
-	assert.Equal(t, res.Rules[0].IP, yandexIP)
-
-	// Check cache.
-	cachedValue, isFound := ss.getCachedResult(domain)
-	require.True(t, isFound)
-	require.Len(t, cachedValue.Rules, 1)
-
-	assert.Equal(t, cachedValue.Rules[0].IP, yandexIP)
-}
-
-func TestSafeSearchCacheGoogle(t *testing.T) {
-	const domain = "www.google.ru"
-
-	ss := newForTest(t, filtering.SafeSearchConfig{Enabled: false})
-
-	res, err := ss.CheckHost(domain, dns.TypeA)
-	require.NoError(t, err)
-
-	assert.False(t, res.IsFiltered)
-	assert.Empty(t, res.Rules)
-
-	resolver := &aghtest.TestResolver{}
-	ss = newForTest(t, defaultSafeSearchConf)
-	ss.resolver = resolver
-
-	// Lookup for safesearch domain.
-	rewrite := ss.SearchHost(domain, dns.TypeA)
-
-	ips, err := resolver.LookupIP(context.Background(), "ip", rewrite.NewCNAME)
-	require.NoError(t, err)
-
-	var foundIP net.IP
-	for _, ip := range ips {
-		if ip.To4() != nil {
-			foundIP = ip
-
-			break
-		}
-	}
-
-	res, err = ss.CheckHost(domain, dns.TypeA)
-	require.NoError(t, err)
-	require.Len(t, res.Rules, 1)
-
-	assert.True(t, res.Rules[0].IP.Equal(foundIP))
-
-	// Check cache.
-	cachedValue, isFound := ss.getCachedResult(domain)
-	require.True(t, isFound)
-	require.Len(t, cachedValue.Rules, 1)
-
-	assert.True(t, cachedValue.Rules[0].IP.Equal(foundIP))
-}
-
-const googleHost = "www.google.com"
-
-var dnsRewriteSink *rules.DNSRewrite
-
-func BenchmarkSafeSearch(b *testing.B) {
-	ss := newForTest(b, defaultSafeSearchConf)
-
-	for n := 0; n < b.N; n++ {
-		dnsRewriteSink = ss.SearchHost(googleHost, dns.TypeA)
-	}
-
-	assert.Equal(b, "forcesafesearch.google.com", dnsRewriteSink.NewCNAME)
-}
-
-var dnsRewriteParallelSink *rules.DNSRewrite
-
-func BenchmarkSafeSearch_parallel(b *testing.B) {
-	ss := newForTest(b, defaultSafeSearchConf)
-
-	b.RunParallel(func(pb *testing.PB) {
-		for pb.Next() {
-			dnsRewriteParallelSink = ss.SearchHost(googleHost, dns.TypeA)
-		}
+	err = ss.Update(filtering.SafeSearchConfig{
+		Enabled: true,
+		Google:  false,
 	})
+	require.NoError(t, err)
 
-	assert.Equal(b, "forcesafesearch.google.com", dnsRewriteParallelSink.NewCNAME)
+	res, err = ss.CheckHost("www.yandex.com", testQType)
+	require.NoError(t, err)
+
+	assert.False(t, res.IsFiltered)
+
+	err = ss.Update(filtering.SafeSearchConfig{
+		Enabled: false,
+		Google:  true,
+	})
+	require.NoError(t, err)
+
+	res, err = ss.CheckHost("www.yandex.com", testQType)
+	require.NoError(t, err)
+
+	assert.False(t, res.IsFiltered)
 }