diff --git a/internal/filtering/hashprefix/cache.go b/internal/filtering/hashprefix/cache.go index d77a25e2..99d6cc73 100644 --- a/internal/filtering/hashprefix/cache.go +++ b/internal/filtering/hashprefix/cache.go @@ -1,6 +1,7 @@ package hashprefix import ( + "context" "encoding/binary" "time" ) @@ -89,7 +90,7 @@ func (c *Checker) findInCache( } // storeInCache caches hashes. -func (c *Checker) storeInCache(hashesToRequest, respHashes []hostnameHash) { +func (c *Checker) storeInCache(ctx context.Context, hashesToRequest, respHashes []hostnameHash) { hashToStore := make(map[prefix][]hostnameHash) for _, hash := range respHashes { @@ -100,7 +101,7 @@ func (c *Checker) storeInCache(hashesToRequest, respHashes []hostnameHash) { } for pref, hash := range hashToStore { - c.setCache(pref, hash) + c.setCache(ctx, pref, hash) } for _, hash := range hashesToRequest { @@ -109,18 +110,18 @@ func (c *Checker) storeInCache(hashesToRequest, respHashes []hostnameHash) { var pref prefix copy(pref[:], hash[:]) - c.setCache(pref, nil) + c.setCache(ctx, pref, nil) } } } // setCache stores hash in cache. -func (c *Checker) setCache(pref prefix, hashes []hostnameHash) { +func (c *Checker) setCache(ctx context.Context, pref prefix, hashes []hostnameHash) { item := &cacheItem{ expiry: time.Now().Add(c.cacheTime), hashes: hashes, } c.cache.Set(pref[:], fromCacheItem(item)) - c.logger.Debug("stored in cache", "pref", pref) + c.logger.DebugContext(ctx, "stored in cache", "pref", pref) } diff --git a/internal/filtering/hashprefix/hashprefix.go b/internal/filtering/hashprefix/hashprefix.go index 75935a03..0a85c417 100644 --- a/internal/filtering/hashprefix/hashprefix.go +++ b/internal/filtering/hashprefix/hashprefix.go @@ -2,6 +2,7 @@ package hashprefix import ( + "context" "crypto/sha256" "encoding/hex" "fmt" @@ -103,20 +104,22 @@ func New(conf *Config) (c *Checker) { // Check returns true if request for the host should be blocked. func (c *Checker) Check(host string) (ok bool, err error) { + ctx := context.TODO() + hashes := hostnameToHashes(host) l := c.logger.With("host", host) found, blocked, hashesToRequest := c.findInCache(hashes) if found { - l.Debug("found in cache", "blocked", blocked) + l.DebugContext(ctx, "found in cache", "blocked", blocked) return blocked, nil } question := c.getQuestion(hashesToRequest) - c.logger.Debug("checking", "question", question) + l.DebugContext(ctx, "checking", "question", question) req := (&dns.Msg{}).SetQuestion(question, dns.TypeTXT) resp, err := c.upstream.Exchange(req) @@ -124,9 +127,9 @@ func (c *Checker) Check(host string) (ok bool, err error) { return false, fmt.Errorf("getting hashes: %w", err) } - matched, receivedHashes := c.processAnswer(l, hashesToRequest, resp) + matched, receivedHashes := c.processAnswer(ctx, l, hashesToRequest, resp) - c.storeInCache(hashesToRequest, receivedHashes) + c.storeInCache(ctx, hashesToRequest, receivedHashes) return matched, nil } @@ -187,6 +190,7 @@ func (c *Checker) getQuestion(hashes []hostnameHash) (q string) { // processAnswer returns true if DNS response matches the hash, and received // hashed hostnames from the upstream. l must not be nil. func (c *Checker) processAnswer( + ctx context.Context, l *slog.Logger, hashesToRequest []hostnameHash, resp *dns.Msg, @@ -201,14 +205,14 @@ func (c *Checker) processAnswer( txtCount++ - receivedHashes = c.appendHashesFromTXT(l, receivedHashes, txt) + receivedHashes = c.appendHashesFromTXT(ctx, l, receivedHashes, txt) } - l.Debug("processing answer with TXT", "txt_count", txtCount) + l.DebugContext(ctx, "processing answer with TXT", "txt_count", txtCount) matched = findMatch(hashesToRequest, receivedHashes) if matched { - l.Debug("matched") + l.DebugContext(ctx, "matched") return true, receivedHashes } @@ -218,22 +222,23 @@ func (c *Checker) processAnswer( // appendHashesFromTXT appends received hashed hostnames. l must not be nil. func (c *Checker) appendHashesFromTXT( + ctx context.Context, l *slog.Logger, hashes []hostnameHash, txt *dns.TXT, ) (receivedHashes []hostnameHash) { - l.Debug("received hashes", "txt", txt.Txt) + l.DebugContext(ctx, "received hashes", "txt", txt.Txt) for _, t := range txt.Txt { if len(t) != hexSize { - l.Debug("wrong hex size", "len", len(t), "txt", t) + l.DebugContext(ctx, "wrong hex size", "len", len(t), "txt", t) continue } buf, err := hex.DecodeString(t) if err != nil { - l.Debug("decoding hex string", "txt", t, slogutil.KeyError, err) + l.DebugContext(ctx, "decoding hex string", "txt", t, slogutil.KeyError, err) continue } diff --git a/internal/filtering/hashprefix/hashprefix_internal_test.go b/internal/filtering/hashprefix/hashprefix_internal_test.go index 6f50d92e..f1ab3b86 100644 --- a/internal/filtering/hashprefix/hashprefix_internal_test.go +++ b/internal/filtering/hashprefix/hashprefix_internal_test.go @@ -11,6 +11,7 @@ import ( "github.com/AdguardTeam/AdGuardHome/internal/aghtest" "github.com/AdguardTeam/golibs/cache" "github.com/AdguardTeam/golibs/logutil/slogutil" + "github.com/AdguardTeam/golibs/testutil" "github.com/miekg/dns" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -96,6 +97,8 @@ func TestHostnameToHashes(t *testing.T) { } func TestChecker_storeInCache(t *testing.T) { + const testTimeout = 1 * time.Second + c := New(&Config{ Logger: slogutil.NewDiscardLogger(), CacheTime: cacheTime, @@ -114,7 +117,7 @@ func TestChecker_storeInCache(t *testing.T) { hashesArray = append(hashesArray, hash4) hash2 := sha256.Sum256([]byte("host.com")) hashesArray = append(hashesArray, hash2) - c.storeInCache(hashes, hashesArray) + c.storeInCache(testutil.ContextWithTimeout(t, testTimeout), hashes, hashesArray) // match "3.sub.host.com" or "host.com" from cache hashes = []hostnameHash{}