filtering: imp hashprefix

This commit is contained in:
Dimitry Kolyshev 2025-06-30 11:39:58 +04:00
parent b88d8799c0
commit 84544cce8a
3 changed files with 25 additions and 16 deletions

View File

@ -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)
}

View File

@ -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
}

View File

@ -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{}