counters/internal/domains/domains_test.go

package domains

import (
	"encoding/binary"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/colinmarc/cdb"
)

// All expectations below come from the real production-shaped CDB shipped in
// test_files/, dumped with `cdb -d -m` (tinycdb). 93 549 records, every value
// "<org_id>|O/i,O/o,D/i,D/o,A/i,A/o[,<local>@/<flow>...]".
const testCDB = "../../test_files/rspamd_domains_db.cdb"

func openTestDB(t *testing.T) *DB {
	t.Helper()
	db, err := Open(testCDB)
	if err != nil {
		t.Fatalf("Open(%s): %v", testCDB, err)
	}
	t.Cleanup(func() { db.Close() })
	return db
}

func mustLookup(t *testing.T, db *DB, domain string) *Domain {
	t.Helper()
	d, err := db.Lookup(domain)
	if err != nil {
		t.Fatalf("Lookup(%q): %v", domain, err)
	}
	return d
}

func TestOpen_MissingFile(t *testing.T) {
	if _, err := Open("/nonexistent/rspamd_domains_db"); err == nil {
		t.Fatal("expected error for missing file")
	}
}

func TestDB_PathAndCount(t *testing.T) {
	db := openTestDB(t)
	if db.Path() != testCDB {
		t.Errorf("Path() = %q", db.Path())
	}
	if got := db.Count(); got != 93549 {
		t.Errorf("Count() = %d, want 93549", got)
	}
}

func TestLookup_RealRecord(t *testing.T) {
	db := openTestDB(t)
	d := mustLookup(t, db, "yagolivas.com")
	if d.Name != "yagolivas.com" || d.OrgID != 1049 {
		t.Fatalf("got %+v", d)
	}
	want := Limits{"O/i": 200, "O/o": 20, "D/i": 200, "D/o": 20, "A/i": 200, "A/o": 20}
	if len(d.Limits) != len(want) {
		t.Fatalf("limits = %v, want %v", d.Limits, want)
	}
	for k, v := range want {
		if d.Limits[k] != v {
			t.Errorf("limits[%s] = %d, want %d", k, d.Limits[k], v)
		}
	}
}

func TestLookup_NotFound(t *testing.T) {
	db := openTestDB(t)
	for _, domain := range []string{"gmail.com", "example.com", "example.org", ""} {
		_, err := db.Lookup(domain)
		if !errors.Is(err, ErrNotFound) {
			t.Errorf("Lookup(%q) err = %v, want ErrNotFound", domain, err)
		}
	}
}

func TestLookup_NormalizesCase(t *testing.T) {
	db := openTestDB(t)
	d := mustLookup(t, db, "  YagoLivas.COM ")
	if d.OrgID != 1049 || d.Name != "yagolivas.com" {
		t.Errorf("got %+v", d)
	}
}

// boxmailhere.com carries a 193 KB value with 9 798 per-address overrides.
func TestLookup_HugeRecord(t *testing.T) {
	db := openTestDB(t)
	d := mustLookup(t, db, "boxmailhere.com")
	if d.OrgID != 87716 {
		t.Errorf("OrgID = %d, want 87716", d.OrgID)
	}
	if len(d.Limits) != 9804 {
		t.Errorf("len(Limits) = %d, want 9804 (6 base + 9798 overrides)", len(d.Limits))
	}
	checkLimit(t, d.Limits, AddressLevel, FlowIn, "gangstaelmo7", 500, true)
	checkLimit(t, d.Limits, AddressLevel, FlowOut, "gangstaelmo7", 30, true)
	checkLimit(t, d.Limits, AddressLevel, FlowIn, "nobody-with-override", 250000, true)
	checkLimit(t, d.Limits, OrganizationLevel, FlowOut, "", 4000, true)
}

func checkLimit(t *testing.T, l Limits, level, flow, local string, want int64, wantOK bool) {
	t.Helper()
	got, ok := l.Limit(level, flow, local)
	if got != want || ok != wantOK {
		t.Errorf("Limit(%s, %s, %q) = (%d, %v), want (%d, %v)", level, flow, local, got, ok, want, wantOK)
	}
}

// crisdanshopping.com.br: A/i=500,A/o=100 plus contato@ and admin@ overrides.
func TestLimit_AddressOverride(t *testing.T) {
	db := openTestDB(t)
	d := mustLookup(t, db, "crisdanshopping.com.br")
	if d.OrgID != 56 {
		t.Fatalf("OrgID = %d", d.OrgID)
	}
	l := d.Limits
	checkLimit(t, l, AddressLevel, FlowIn, "contato", 300, true)
	checkLimit(t, l, AddressLevel, FlowOut, "contato", 150, true)
	checkLimit(t, l, AddressLevel, FlowIn, "admin", 300, true)
	checkLimit(t, l, AddressLevel, FlowOut, "admin", 150, true)
	checkLimit(t, l, AddressLevel, FlowIn, "someone", 500, true) // falls back to A/i
	checkLimit(t, l, AddressLevel, FlowOut, "someone", 100, true)
	checkLimit(t, l, DomainLevel, FlowIn, "", 1000, true)
	checkLimit(t, l, DomainLevel, FlowOut, "", 120, true)
	checkLimit(t, l, OrganizationLevel, FlowIn, "", 1000, true)
	checkLimit(t, l, OrganizationLevel, FlowOut, "", 120, true)
	// local part is irrelevant for D and O
	checkLimit(t, l, DomainLevel, FlowIn, "contato", 1000, true)
}

// meu-servidor.com has overrides whose local parts contain dots.
func TestLimit_DottedLocalPart(t *testing.T) {
	db := openTestDB(t)
	d := mustLookup(t, db, "meu-servidor.com")
	if d.OrgID != 726 {
		t.Fatalf("OrgID = %d", d.OrgID)
	}
	checkLimit(t, d.Limits, AddressLevel, FlowIn, "sexprime.xxx", 100, true)
	checkLimit(t, d.Limits, AddressLevel, FlowIn, "cf-tgamesandroid.com", 1000, true)
	checkLimit(t, d.Limits, AddressLevel, FlowOut, "cf-tgamesandroid.com", 100, true)
	checkLimit(t, d.Limits, AddressLevel, FlowIn, "alerta", 1000, true)
	checkLimit(t, d.Limits, AddressLevel, FlowOut, "alerta", 100, true)
	checkLimit(t, d.Limits, AddressLevel, FlowIn, "nobody", 500, true)
}

func TestLimit_MigaduSelf(t *testing.T) {
	db := openTestDB(t)
	d := mustLookup(t, db, "migadu.com")
	if d.OrgID != 3743 {
		t.Fatalf("OrgID = %d", d.OrgID)
	}
	checkLimit(t, d.Limits, AddressLevel, FlowIn, "admin", 10000, true)
	checkLimit(t, d.Limits, AddressLevel, FlowIn, "someone", 100000, true)
}

func TestLimit_Unlimited(t *testing.T) {
	l := ParseLimits("O/i=5")
	checkLimit(t, l, OrganizationLevel, FlowIn, "", 5, true)
	checkLimit(t, l, OrganizationLevel, FlowOut, "", 0, false)
	checkLimit(t, l, DomainLevel, FlowIn, "", 0, false)
	checkLimit(t, l, AddressLevel, FlowIn, "x", 0, false)
	checkLimit(t, Limits(nil), AddressLevel, FlowIn, "x", 0, false)
	checkLimit(t, l, "Z", FlowIn, "", 0, false) // unknown level
}

func TestParseLimits(t *testing.T) {
	l := ParseLimits(" O/i = 5000,O/o=1000, broken ,x=notanumber,=3,A/o=,,alice@/o=7")
	want := Limits{"O/i": 5000, "O/o": 1000, "alice@/o": 7}
	if len(l) != len(want) {
		t.Fatalf("got %v, want %v", l, want)
	}
	for k, v := range want {
		if l[k] != v {
			t.Errorf("%s = %d, want %d", k, l[k], v)
		}
	}
	if got := ParseLimits(""); len(got) != 0 {
		t.Errorf("ParseLimits(\"\") = %v, want empty", got)
	}
}

func TestParseRecord(t *testing.T) {
	tests := []struct {
		name, value string
		wantOrg     int64
		wantLimits  int
		wantErr     bool
	}{
		{"normal", "1049|O/i=200,O/o=20", 1049, 2, false},
		{"spaces", " 7 | A/i=3 ", 7, 1, false},
		{"no pipe", "42", 42, 0, false}, // must not panic (the Lua plugin did)
		{"empty limits", "42|", 42, 0, false},
		{"non-numeric org", "abc|O/i=1", 0, 0, true},
		{"empty", "", 0, 0, true},
		{"negative org", "-1|O/i=1", 0, 0, true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			d, err := ParseRecord("x.com", tt.value)
			if tt.wantErr {
				if err == nil {
					t.Fatalf("expected error, got %+v", d)
				}
				return
			}
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}
			if d.Name != "x.com" || d.OrgID != tt.wantOrg || len(d.Limits) != tt.wantLimits {
				t.Errorf("got %+v, want org=%d limits=%d", d, tt.wantOrg, tt.wantLimits)
			}
		})
	}
}

// String renders the Lua-like limits_map used by the transaction log, sorted.
func TestLimits_String(t *testing.T) {
	l := ParseLimits("O/o=20,A/i=200,A/o=20,O/i=200,D/i=200,D/o=20,info@/o=5")
	want := "{A/i=200, A/o=20, D/i=200, D/o=20, O/i=200, O/o=20, info@/o=5}"
	if got := l.String(); got != want {
		t.Errorf("String() = %q, want %q", got, want)
	}
	if got := (Limits{}).String(); got != "{}" {
		t.Errorf("empty String() = %q", got)
	}
}

func TestDB_Ping(t *testing.T) {
	db, err := Open(testCDB)
	if err != nil {
		t.Fatal(err)
	}
	if err := db.Ping(); err != nil {
		t.Errorf("Ping() = %v", err)
	}
	db.Close()
	if err := db.Ping(); err == nil {
		t.Error("Ping() after Close() should fail")
	}
}

// A CDB written with colinmarc/cdb's Writer (as tests and tooling do) must be
// readable through the same path as the agent-built file.
func TestLookup_WriterBuiltCDB(t *testing.T) {
	path := filepath.Join(t.TempDir(), "small.cdb")
	w, err := cdb.Create(path)
	if err != nil {
		t.Fatal(err)
	}
	if err := w.Put([]byte("newdomain.test"), []byte("5|A/i=1,A/o=1,D/i=2,D/o=2,O/i=3,O/o=3")); err != nil {
		t.Fatal(err)
	}
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}
	db, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	if db.Count() != 1 {
		t.Errorf("Count() = %d, want 1", db.Count())
	}
	d, err := db.Lookup("newdomain.test")
	if err != nil {
		t.Fatalf("Lookup: %v", err)
	}
	if d.OrgID != 5 {
		t.Errorf("OrgID = %d", d.OrgID)
	}
	if _, err := db.Lookup("other.test"); !errors.Is(err, ErrNotFound) {
		t.Errorf("Lookup(other) err = %v, want ErrNotFound", err)
	}
}

// Parsed records are cached per loaded table: a second Lookup of the same
// domain returns the very same *Domain (no re-parse), and is safe because
// callers never mutate it.
func TestLookup_CachesParsedRecord(t *testing.T) {
	db := openTestDB(t)
	first := mustLookup(t, db, "boxmailhere.com")
	second := mustLookup(t, db, "BoxMailHere.com ")
	if first != second {
		t.Error("second lookup re-parsed the record instead of serving the cached *Domain")
	}
	if db.Cached() != 1 {
		t.Errorf("Cached() = %d, want 1", db.Cached())
	}
	if _, err := db.Lookup("nobody.invalid"); !errors.Is(err, ErrNotFound) {
		t.Fatal(err)
	}
	if db.Cached() != 1 {
		t.Errorf("not-found lookups must not be cached: Cached() = %d", db.Cached())
	}
}

// The cache is bounded: beyond max entries lookups still work but parse
// without caching; 0 disables the cache entirely.
func TestLookup_CacheBounded(t *testing.T) {
	db, err := OpenWithCache(testCDB, 1)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	a1 := mustLookup(t, db, "yagolivas.com")
	a2 := mustLookup(t, db, "yagolivas.com")
	b1 := mustLookup(t, db, "boxmailhere.com")
	b2 := mustLookup(t, db, "boxmailhere.com")
	if a1 != a2 {
		t.Error("first domain should be cached")
	}
	if b1 == b2 {
		t.Error("second domain exceeds max_cached_domains=1 and must not be cached")
	}
	if b1.OrgID != 87716 || db.Cached() != 1 {
		t.Errorf("uncached lookup still parses: org=%d cached=%d", b1.OrgID, db.Cached())
	}

	off, err := OpenWithCache(testCDB, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer off.Close()
	if mustLookup(t, off, "yagolivas.com") == mustLookup(t, off, "yagolivas.com") {
		t.Error("cache disabled (0) must parse every time")
	}
}

func benchLookup(b *testing.B, dom string) {
	db, err := Open(testCDB)
	if err != nil {
		b.Fatal(err)
	}
	defer db.Close()
	if _, err := db.Lookup(dom); err != nil { // warm the cache
		b.Fatal(err)
	}
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if _, err := db.Lookup(dom); err != nil {
			b.Fatal(err)
		}
	}
}

// With the cache warm, the 193 KB boxmailhere.com record costs the same as a
// small one (review finding #5: 0.9 ms and 1.4 MB per lookup before).
func BenchmarkLookupSmall(b *testing.B) { benchLookup(b, "yagolivas.com") }
func BenchmarkLookupHuge(b *testing.B)  { benchLookup(b, "boxmailhere.com") }

// A row with a non-numeric org id is reported as malformed, which callers
// treat like "not hosted" (review #8) — it must never fail the request.
func TestLookup_MalformedRow(t *testing.T) {
	db, err := Open("../../review/_proof_tests/b/bad.cdb")
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	_, err = db.Lookup("bad.com")
	if !errors.Is(err, ErrMalformed) {
		t.Fatalf("Lookup(bad.com) err = %v, want ErrMalformed", err)
	}
	if !strings.Contains(err.Error(), `invalid org id "x12"`) {
		t.Errorf("error should name the problem: %v", err)
	}
	if d, err := db.Lookup("good.com"); err != nil || d.OrgID != 7 {
		t.Errorf("good.com = %+v, %v", d, err)
	}
	if db.Cached() != 1 {
		t.Errorf("malformed rows must not be cached: Cached() = %d", db.Cached())
	}
}

func TestMailbox(t *testing.T) {
	for in, want := range map[string]string{
		"b+1@example.com":        "b@example.com",
		"b+tag+more@example.com": "b@example.com",
		"b@example.com":          "b@example.com",
		"+lead@example.com":      "+lead@example.com", // empty local part would be nonsense: keep as is
		"nodomain":               "nodomain",
	} {
		if got := Mailbox(in); got != want {
			t.Errorf("Mailbox(%q) = %q, want %q", in, got, want)
		}
	}
}

func TestParseLimits_OverrideKeysLowercased(t *testing.T) {
	l := ParseLimits("A/o=20,Info@/o=5,O/o=9")
	if n, ok := l.Limit(AddressLevel, FlowOut, "info"); !ok || n != 5 {
		t.Errorf("override for info = %d, %v; want 5", n, ok)
	}
	if n, ok := l.Limit(OrganizationLevel, FlowOut, ""); !ok || n != 9 {
		t.Errorf("O/o = %d, %v; level keys must keep their case", n, ok)
	}
}

// A file cut off inside its hash tables passes a walk over the records, but
// lookups would then fail: it must be refused at open.
func TestOpen_RejectsTruncatedHashTables(t *testing.T) {
	path := filepath.Join(t.TempDir(), "t.cdb")
	w, err := cdb.Create(path)
	if err != nil {
		t.Fatal(err)
	}
	for _, k := range []string{"a.example", "b.example", "c.example", "d.example"} {
		if err := w.Put([]byte(k), []byte("1|A/o=1")); err != nil {
			t.Fatal(err)
		}
	}
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}
	fi, err := os.Stat(path)
	if err != nil {
		t.Fatal(err)
	}
	if err := os.Truncate(path, fi.Size()-8); err != nil {
		t.Fatal(err)
	}
	if db, err := Open(path); err == nil {
		db.Close()
		t.Fatal("truncated CDB was accepted")
	}
}

// Lengths that overflow uint32 make the cdb library panic; a bad file must be
// an error, because reloads run on a goroutine with no recover.
func TestOpen_CorruptLengthsAreAnError(t *testing.T) {
	buf := make([]byte, 2048+32)
	binary.LittleEndian.PutUint32(buf[0:], 2048+32) // table 0 starts after one bogus record
	binary.LittleEndian.PutUint32(buf[2048:], 0xFFFFFFFF)
	binary.LittleEndian.PutUint32(buf[2052:], 2)
	path := filepath.Join(t.TempDir(), "bad.cdb")
	if err := os.WriteFile(path, buf, 0o600); err != nil {
		t.Fatal(err)
	}
	if db, err := Open(path); err == nil {
		db.Close()
		t.Fatal("corrupt CDB was accepted")
	}
}