fancy-dns/internal/blocklist/set_test.go

package blocklist

import (
	"fmt"
	"sync"
	"testing"
)

// mustBuild compiles rules into an immutable Set.
func mustBuild(rules ...Rule) *Set {
	b := NewBuilder()
	for _, r := range rules {
		b.Add(r, "test", nil)
	}
	return b.Build()
}

func block(name string) Rule { return Rule{Name: name, Action: ActionBlock} }

func TestSetMatch(t *testing.T) {
	s := mustBuild(
		block("example.com"),                                   // exact: the name and nothing else
		block("*.wild.example.com"),                            // strict subdomains only
		block("**.root.example.net"),                           // apex and subdomains
		block("**.sub.example.org"),                            // whole subtree
		block("ads.example.com"),                               // more specific than the exact "example.com"
		Rule{Name: "tracker.example.com", Action: ActionAllow}, // exact allow
	)

	tests := []struct {
		name     string
		query    string
		want     string // stored rule name, "" = no match
		kind     PatternKind
		action   Action
		explains string
	}{
		{name: "exact apex", query: "example.com", want: "example.com", kind: PatternExact},
		{name: "exact apex trailing dot", query: "example.com.", want: "example.com", kind: PatternExact},
		{name: "exact apex uppercase", query: "ExAmPlE.CoM", want: "example.com", kind: PatternExact},
		{name: "exact rule does not cover subdomains", query: "sub.example.com", want: "", kind: PatternExact},
		{name: "sub wildcard matches subdomain", query: "a.wild.example.com", want: "wild.example.com", kind: PatternSub},
		{name: "sub wildcard does not match apex", query: "wild.example.com", want: "", kind: PatternExact},
		{name: "root wildcard matches apex", query: "root.example.net", want: "root.example.net", kind: PatternRoot},
		{name: "root wildcard matches subdomain", query: "deep.a.root.example.net", want: "root.example.net", kind: PatternRoot},
		{name: "subtree", query: "a.b.sub.example.org", want: "sub.example.org", kind: PatternRoot},
		{name: "subtree apex", query: "sub.example.org", want: "sub.example.org", kind: PatternRoot},
		{name: "most specific exact wins over root", query: "ads.example.com", want: "ads.example.com", kind: PatternExact},
		{name: "allow rule reported with its action", query: "tracker.example.com", want: "tracker.example.com", kind: PatternExact, action: ActionAllow},
		{name: "unrelated name", query: "unrelated.test", want: ""},
		{name: "empty name", query: "", want: ""},
		{name: "the root label is normalised away", query: ".", want: ""},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			got := s.Match(tc.query)
			if tc.want == "" {
				if got != nil {
					t.Fatalf("Match(%q) = %q, want no match", tc.query, got.Name)
				}
				return
			}
			if got == nil {
				t.Fatalf("Match(%q) = nil, want %q", tc.query, tc.want)
			}
			if got.Name != tc.want {
				t.Errorf("Match(%q) rule = %q, want %q", tc.query, got.Name, tc.want)
			}
			wantAction := tc.action
			if wantAction == ActionBlock && got.Action != ActionAllow {
				wantAction = ActionBlock
			}
			if got.Action != wantAction {
				t.Errorf("Match(%q) action = %v, want %v", tc.query, got.Action, wantAction)
			}
			if k := s.MatchPattern(tc.query); k != tc.kind {
				t.Errorf("MatchPattern(%q) = %v, want %v", tc.query, k, tc.kind)
			}
		})
	}
}

func TestSetMatchLeadingDotAndSubdomainToggle(t *testing.T) {
	// splitPattern maps a leading dot to the "**." (root) form.
	s := mustBuild(block(".dot.example.com"))
	if e := s.Match("dot.example.com"); e == nil {
		t.Errorf("leading-dot rule did not match its apex")
	}
	if e := s.Match("x.dot.example.com"); e == nil {
		t.Errorf("leading-dot rule did not match a subdomain")
	}

	// With Subdomains off the parsers emit exact rules, so nothing below the
	// apex is covered.
	exact := mustBuild(hostsRule("ads.example.com", ParseOptions{Subdomains: false}))
	if e := exact.Match("sub.ads.example.com"); e != nil {
		t.Errorf("exact hosts rule matched a subdomain: %q", e.Name)
	}
	if e := exact.Match("ads.example.com"); e == nil {
		t.Errorf("exact hosts rule did not match its own name")
	}
	if exact.HasWildcards() {
		t.Errorf("HasWildcards() = true for an exact-only set")
	}
}

func TestSetDropsUnusablePatterns(t *testing.T) {
	// "." is "the whole internet" and is never useful: it is dropped. A bare
	// TLD, on the other hand, is kept and matches the TLD itself.
	s := mustBuild(block("."), block("*.com"), block("com"), block(""), block("   "))
	if s.Len() != 2 {
		t.Fatalf("Len() = %d, want 2 (exact \"com\" and sub \"com\")", s.Len())
	}
	if e := s.Match("com"); e == nil || e.Name != "com" {
		t.Errorf("Match(\"com\") = %v, want the exact rule \"com\"", e)
	}
	if e := s.Match("ads.com"); e == nil || e.Name != "com" {
		t.Errorf("Match(\"ads.com\") = %v, want sub rule \"com\"", e)
	}
}

func TestSetMatchesWholeSubtree(t *testing.T) {
	// A single root rule on a TLD covers its whole subtree.
	s := mustBuild(block("**.com"))
	for _, q := range []string{"com", "example.com", "a.b.c.example.com", "EXAMPLE.COM."} {
		if e := s.Match(q); e == nil || e.Name != "com" {
			t.Errorf("Match(%q) = %v, want the \"**.com\" rule", q, e)
		}
	}
	for _, q := range []string{"de", "example.de", "com.evil.test"} {
		if e := s.Match(q); e != nil {
			t.Errorf("Match(%q) = %q, want no match", q, e.Name)
		}
	}
	// The apex is covered by "**." but not by "*.".
	sub := mustBuild(block("*.com"))
	if e := sub.Match("com"); e != nil {
		t.Errorf("\"*.com\" matched the apex: %q", e.Name)
	}
	if e := sub.Match("example.com"); e == nil {
		t.Errorf("\"*.com\" did not match a subdomain")
	}
	// A single label or an empty name never reaches any map.
	if e := s.Match(""); e != nil {
		t.Errorf("Match(\"\") = %q, want nil", e.Name)
	}
}

func TestNormalizeName(t *testing.T) {
	tests := []struct{ in, want string }{
		{"example.com", "example.com"},
		{"example.com.", "example.com"},
		{"EXAMPLE.com..", "example.com."}, // one trailing dot, then lowercased
		{"", ""},
		{".", ""},
		{"..", "."},
	}
	for _, tc := range tests {
		if got := NormalizeName(tc.in); got != tc.want {
			t.Errorf("NormalizeName(%q) = %q, want %q", tc.in, got, tc.want)
		}
	}
	// Match must treat "example.com." and "example.com" as the same name,
	// and must be case-insensitive.
	s := mustBuild(block("example.com"))
	if s.Match("example.com.") == nil || s.Match("EXAMPLE.COM") == nil {
		t.Errorf("Match is not case/trailing-dot insensitive")
	}
}

func TestAllowOverBlockPriority(t *testing.T) {
	b := NewBuilder()
	b.Add(Rule{Name: "cdn.example.com", Action: ActionBlock}, "blocklist", nil)
	b.Add(Rule{Name: "cdn.example.com", Action: ActionAllow}, "allowlist", []string{"user"})
	s := b.Build()

	e := s.Match("cdn.example.com")
	if e == nil {
		t.Fatal("no rule matched")
	}
	if !e.IsAllow() {
		t.Errorf("IsAllow() = false for %v", e.Action)
	}
	if e.Action != ActionAllow {
		t.Errorf("action = %v, want allow", e.Action)
	}
	if e.Source != "allowlist" {
		t.Errorf("Source = %q, want %q (the winning rule's list)", e.Source, "allowlist")
	}
	if e.Priority != 10 {
		t.Errorf("Priority = %d, want 10 (allow default)", e.Priority)
	}
	if got := e.Categories; len(got) != 1 || got[0] != "user" {
		t.Errorf("Categories = %v, want [user]", got)
	}
	c := s.Counts()
	if c.Allow != 1 || c.Block != 0 {
		t.Errorf("Counts() = %+v, want Allow=1 Block=0 (block rule was replaced)", c)
	}
	if s.Duplicates() != 1 {
		t.Errorf("Duplicates() = %d, want 1", s.Duplicates())
	}
	if s.Lookup("cdn.example.com").Action != ActionAllow {
		t.Errorf("Lookup() disagrees with Match()")
	}
}

func TestPriorityTieKeepsFirstRule(t *testing.T) {
	b := NewBuilder()
	b.Add(Rule{Name: "tie.example.com", Action: ActionBlock, Priority: 5}, "first", nil)
	b.Add(Rule{Name: "tie.example.com", Action: ActionBlock, Priority: 5}, "second", nil)
	s := b.Build()
	if s.Lookup("tie.example.com").Source != "first" {
		t.Errorf("on a priority tie the later rule won: source = %q", s.Lookup("tie.example.com").Source)
	}
	if s.Duplicates() != 1 {
		t.Errorf("Duplicates() = %d, want 1", s.Duplicates())
	}
}

func TestPassthroughCountsAsAllow(t *testing.T) {
	b := NewBuilder()
	b.Add(Rule{Name: "pt.example.com", Action: ActionPassthrough}, "rpz", nil)
	s := b.Build()
	e := s.Match("pt.example.com")
	if e == nil || !e.IsAllow() {
		t.Fatalf("passthrough rule: IsAllow() = false, e = %v", e)
	}
	if c := s.Counts(); c.Passthru != 1 || c.Allow != 0 {
		t.Errorf("Counts() = %+v, want Passthru=1", c)
	}
}

func TestHits(t *testing.T) {
	s := mustBuild(block("hit.example.com"), block("*.wild.hit.example.com"))
	const n = 3
	for i := 0; i < n; i++ {
		if s.Match("hit.example.com") == nil {
			t.Fatal("setup: no match")
		}
	}
	e := s.Match("hit.example.com")
	if got := e.Hits(); got != n+1 {
		t.Errorf("Hits() = %d, want %d (every Match counts, including this one)", got, n+1)
	}
	// Non-matching queries must not bump any counter.
	s.Match("miss.example.com")
	if got := s.Match("hit.example.com").Hits(); got != n+2 {
		t.Errorf("Hits() = %d after a miss, want %d", got, n+2)
	}
	// A wildcard rule has its own counter, and the shared *Entry is visible
	// through every pointer to it.
	if got := s.Match("a.wild.hit.example.com").Hits(); got != 1 {
		t.Errorf("wildcard Hits() = %d, want 1", got)
	}
	if got := s.Match("hit.example.com").Hits(); got != n+3 {
		t.Errorf("exact Hits() = %d, want %d", got, n+3)
	}
	// A zero-valued Entry is safe.
	var zero Entry
	if zero.Hits() != 0 {
		t.Errorf("zero Entry Hits() = %d, want 0", zero.Hits())
	}
	if got := (*Entry)(nil).Hits(); got != 0 {
		t.Errorf("nil Entry Hits() = %d, want 0", got)
	}
}

func TestDuplicateDomainsAcrossLists(t *testing.T) {
	// The same domain from two lists compiles to ONE rule: the dedup happens
	// in Builder.Add, keyed on the normalised pattern, and the losing rule is
	// dropped (Duplicates counts it) rather than kept as a second entry.
	b := NewBuilder()
	b.Add(Rule{Name: "Dup.Example.com.", Action: ActionBlock}, "list-a", nil)
	b.Add(Rule{Name: "dup.example.com", Action: ActionBlock}, "list-b", nil)
	s := b.Build()

	if s.Len() != 1 {
		t.Fatalf("Len() = %d, want 1 (duplicates are deduped)", s.Len())
	}
	if s.Duplicates() != 1 {
		t.Errorf("Duplicates() = %d, want 1", s.Duplicates())
	}
	e := s.Match("dup.example.com")
	if e == nil {
		t.Fatal("no match")
	}
	if e.Source != "list-a" {
		t.Errorf("Source = %q, want list-a (first list wins on a tie)", e.Source)
	}
	// One Entry => one hit counter, regardless of how many lists listed it.
	if got := e.Hits(); got != 1 {
		t.Errorf("Hits() = %d, want 1 (deduped rules share one counter)", got)
	}
	// A higher-priority rule from a later list replaces the earlier one.
	b2 := NewBuilder()
	b2.Add(Rule{Name: "dup.example.com", Action: ActionBlock}, "list-a", nil)
	b2.Add(Rule{Name: "dup.example.com", Action: ActionAllow}, "list-b", nil)
	s2 := b2.Build()
	if a := s2.Match("dup.example.com"); a == nil || a.Source != "list-b" {
		t.Errorf("higher-priority later rule did not win: %+v", a)
	}
	if s2.Duplicates() != 1 {
		t.Errorf("Duplicates() = %d, want 1", s2.Duplicates())
	}
}

func TestSetIntrospection(t *testing.T) {
	b := NewBuilder()
	b.Add(block("exact.example.com"), "a", nil)
	b.Add(block("*.sub.example.com"), "b", nil)
	b.Add(block("**.root.example.com"), "c", nil)
	b.NoteSource(SourceMeta{Name: "a", Format: "hosts", Rules: 1, Bytes: 42, Enabled: true})
	s := b.Build()

	if got := s.Sources(); len(got) != 3 || got[0] != "a" || got[2] != "c" {
		t.Errorf("Sources() = %v, want insertion order [a b c]", got)
	}
	m, ok := s.SourceMeta("b")
	if !ok {
		t.Errorf("SourceMeta(\"b\") missing")
	}
	_ = m
	if a, ok := s.SourceMeta("a"); !ok || a.Rules != 1 || a.Bytes != 42 || a.Format != "hosts" {
		t.Errorf("SourceMeta(\"a\") = %+v, ok=%v", a, ok)
	}
	if _, ok := s.SourceMeta("nope"); ok {
		t.Errorf("SourceMeta(\"nope\") reported ok")
	}
	if c := s.Counts(); c.Exact != 1 || c.Sub != 1 || c.Root != 1 || c.Block != 3 {
		t.Errorf("Counts() = %+v, want 1/1/1 exact/sub/root and 3 block", c)
	}
	st := s.Stats()
	if st.Total != 3 || st.Exact != 1 || st.Subdomains != 1 || st.RootRules != 1 {
		t.Errorf("Stats() = %+v", st)
	}
	if st.BuiltAt.IsZero() {
		t.Errorf("Stats().BuiltAt is zero")
	}

	// Each renders patterns back with their wildcard prefix.
	seen := map[string]PatternKind{}
	s.Each(func(pattern string, e *Entry, kind PatternKind) bool {
		seen[pattern] = kind
		if got := kind.Pattern(e.Name); got != pattern {
			t.Errorf("Pattern(%q) = %q, want %q", e.Name, got, pattern)
		}
		return true
	})
	if seen["exact.example.com"] != PatternExact || seen["sub.example.com"] != PatternSub ||
		seen["root.example.com"] != PatternRoot {
		t.Errorf("Each() visited %v", seen)
	}
	if PatternSub.String() != "subdomain" || PatternRoot.String() != "domain" || PatternExact.String() != "exact" {
		t.Errorf("PatternKind.String() values changed")
	}

	// Lookup is exact-only: it does not walk wildcards.
	if s.Lookup("deep.sub.example.com") != nil {
		t.Errorf("Lookup() walked wildcards")
	}
	if s.Lookup("EXACT.example.com.") == nil {
		t.Errorf("Lookup() is not normalised")
	}
}

func TestDigestStableAndContentSensitive(t *testing.T) {
	a := mustBuild(block("x.example.com"), block("*.y.example.com"), block("**.z.example.com"))
	b := mustBuild(block("**.z.example.com"), block("x.example.com"), block("*.y.example.com"))
	if a.Digest() != b.Digest() {
		t.Errorf("Digest() depends on insertion order: %q vs %q", a.Digest(), b.Digest())
	}
	c := mustBuild(block("x.example.com"), block("*.y.example.com"), block("**.different.example.com"))
	if a.Digest() == c.Digest() {
		t.Errorf("Digest() ignored a pattern change")
	}
	// The action is part of the digest.
	d := mustBuild(Rule{Name: "x.example.com", Action: ActionAllow})
	if a.Digest() == d.Digest() {
		t.Errorf("Digest() ignored the action change")
	}
}

func TestNilSetIsSafe(t *testing.T) {
	var s *Set
	if s.Match("example.com") != nil || s.Len() != 0 || s.HasWildcards() ||
		s.DuplicatesN() != 0 || s.Sources() != nil || s.AllSourceMeta() != nil ||
		s.Digest() != "" || !s.BuiltAt().IsZero() {
		t.Errorf("nil *Set accessors are not nil-safe")
	}
	if _, ok := s.SourceMeta("x"); ok {
		t.Errorf("nil Set SourceMeta reported ok")
	}
	if s.Lookup("x") != nil || s.MatchPattern("x") != PatternExact {
		t.Errorf("nil Set lookup misbehaved")
	}
	if st := s.Stats(); st.Total != 0 || st.PerSource == nil {
		t.Errorf("nil Stats() = %+v", st)
	}
	s.Each(func(string, *Entry, PatternKind) bool { t.Errorf("nil Each called fn"); return true })
}

func TestActionStringAndParseAction(t *testing.T) {
	all := []struct {
		a   Action
		str string
	}{
		{ActionBlock, "block"},
		{ActionAllow, "allow"},
		{ActionPassthrough, "passthrough"},
		{ActionDrop, "drop"},
		{ActionNXDomain, "nxdomain"},
		{ActionZero, "zero"},
		{ActionCNAME, "cname"},
		{ActionRedirect, "redirect"},
		{Action(99), "unknown"},
	}
	for _, tc := range all {
		if got := tc.a.String(); got != tc.str {
			t.Errorf("Action(%d).String() = %q, want %q", tc.a, got, tc.str)
		}
	}
	roundtrip := map[string]Action{
		"":            ActionBlock,
		"block":       ActionBlock,
		"ALLOW":       ActionAllow,
		" whitelist ": ActionAllow,
		"pass":        ActionPassthrough,
		"drop":        ActionDrop,
		"nxdomain":    ActionNXDomain,
		"null":        ActionZero,
		"cname":       ActionCNAME,
		"redirect":    ActionRedirect,
	}
	for in, want := range roundtrip {
		got, ok := ParseAction(in)
		if !ok || got != want {
			t.Errorf("ParseAction(%q) = %v,%v; want %v,true", in, got, ok, want)
		}
	}
	if got, ok := ParseAction("nonsense"); ok || got != ActionBlock {
		t.Errorf("ParseAction(\"nonsense\") = %v,%v; want block,false", got, ok)
	}
}

func TestHolderSwapIsAtomic(t *testing.T) {
	old := mustBuild(block("a.example.com"))
	newer := mustBuild(block("b.example.com"))
	h := NewHolder(old)
	if h.Get() != old {
		t.Fatalf("NewHolder did not install the set")
	}
	var empty *Holder
	if empty.Get() != nil {
		t.Errorf("zero Holder is not empty")
	}
	// A reader either sees the whole old set or the whole new one.
	var wg sync.WaitGroup
	stop := make(chan struct{})
	for i := 0; i < 4; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				select {
				case <-stop:
					return
				default:
				}
				// The invariant is that a reader gets *one* of the two
				// published sets, not a blend of them. Comparing pointers is
				// the only way to say that: asking a question the two sets
				// answer differently (does this set match a.example.com?)
				// would report a legitimate swap as a torn one.
				s := h.Get()
				if s != old && s != newer {
					t.Errorf("holder returned neither published set (%p)", s)
					return
				}
			}
		}()
	}
	for i := 0; i < 50; i++ {
		h.Set(newer)
		h.Set(old)
	}
	close(stop)
	wg.Wait()
	if h.Get() != old {
		t.Errorf("holder lost the last swap")
	}
}

func BenchmarkSetMatch(b *testing.B) {
	// A mixed set, like a compiled blocklist: mostly exact rules plus some
	// subdomain and root rules so the suffix walk is exercised.
	bld := NewBuilder()
	for i := 0; i < 2000; i++ {
		bld.Add(block(fmt.Sprintf("host%d.ads-%d.example.com", i, i%7)), "bench", nil)
	}
	for i := 0; i < 200; i++ {
		bld.Add(block(fmt.Sprintf("*.wild%d.example.net", i)), "bench", nil)
		bld.Add(block(fmt.Sprintf("**.root%d.example.org", i)), "bench", nil)
	}
	s := bld.Build()

	queries := []string{
		"host0.ads-0.example.com",    // exact hit
		"host1999.ads-5.example.com", // exact hit, different list
		"a.b.wild199.example.net",    // wildcard suffix hit
		"x.root42.example.org",       // root rule
		"nothing.in.here.test",       // pure miss, walks every suffix
	}
	b.ReportAllocs()
	b.ResetTimer()
	var sink int
	for i := 0; i < b.N; i++ {
		if e := s.Match(queries[i%len(queries)]); e != nil {
			sink++
		}
	}
	if sink == 0 {
		b.Fatal("no matches at all")
	}
}