counters/internal/domains/domains.go

// Package domains reads the rspamd domains CDB (rspamd_domains_db) that holds
// the organisation id and the daily limits of every Migadu-hosted domain.
//
// The file is built by the agent's extract.sh with tinycdb (`cdb -m`) from
// lines of the form
//
//	example.com 1049|O/i=200,O/o=20,D/i=200,D/o=20,A/i=200,A/o=20,info@/o=5
//
// i.e. key = lowercase domain (no NUL terminator), value =
// "<org_id>|<comma separated limits>". Limit keys are "<level>/<flow>" with
// level O (organisation), D (domain) or A (address default) and flow i
// (incoming) or o (outgoing); "<localpart>@/<flow>" overrides the A default for
// one mailbox. A level/flow without a key is unlimited. A domain that is not
// in the file is not hosted by Migadu and is never counted.
package domains

import (
	"errors"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/colinmarc/cdb"
)

// DefaultMaxCached is the number of parsed records a DB caches when opened
// with Open; the daemon passes its max_cached_domains instead.
const DefaultMaxCached = 20000

// Limit levels and flows, as used in the CDB values and the Redis keys.
const (
	OrganizationLevel = "O"
	DomainLevel       = "D"
	AddressLevel      = "A"

	FlowIn  = "i"
	FlowOut = "o"
)

// ErrNotFound is returned by Lookup when the domain is not in the CDB.
var ErrNotFound = errors.New("domain not found in domains db")

// ErrMalformed wraps a Lookup error for a row that exists but cannot be
// parsed (non-numeric org id). Callers treat it like "not hosted" and warn:
// one bad generated row must never fail a request (review #8).
var ErrMalformed = errors.New("malformed row in domains db")

// Mailbox strips a plus-address tag from the local part: "b+tag@d" → "b@d".
// Counters are kept per mailbox, which is what the web app reads and what
// rcptd resolves; a tag never names a different mailbox (review #7). An
// address without "@", or whose local part would become empty, is returned
// unchanged.
func Mailbox(addr string) string {
	at := strings.LastIndex(addr, "@")
	if at <= 0 {
		return addr
	}
	plus := strings.Index(addr[:at], "+")
	if plus <= 0 {
		return addr
	}
	return addr[:plus] + addr[at:]
}

// Limits maps limit keys ("O/i", "A/o", "alice@/o", ...) to their values.
type Limits map[string]int64

// ParseLimits parses "O/i=5000, A/o=200, alice@/o=7" into a Limits map.
// Keys are trimmed; pairs without "=", with an empty key or with a
// non-numeric value are skipped rather than failing the whole record — the
// file is generated, and one bad pair must not disable counting for a domain.
func ParseLimits(s string) Limits {
	limits := Limits{}
	for _, pair := range strings.Split(s, ",") {
		k, v, ok := strings.Cut(pair, "=")
		if !ok {
			continue
		}
		k = strings.TrimSpace(k)
		if k == "" {
			continue
		}
		n, err := strconv.ParseInt(strings.TrimSpace(v), 10, 64)
		if err != nil {
			continue
		}
		limits[k] = n
	}
	return limits
}

// Limit resolves the configured limit for a level and flow, or (0, false)
// when unlimited. At address level the per-mailbox override "<local>@/<flow>"
// wins over the domain-wide default "A/<flow>" — same rule as the rspamd
// plugin. localPart is ignored for the D and O levels.
func (l Limits) Limit(level, flow, localPart string) (int64, bool) {
	switch level {
	case AddressLevel:
		if n, ok := l[localPart+"@/"+flow]; ok {
			return n, true
		}
		n, ok := l[AddressLevel+"/"+flow]
		return n, ok
	case DomainLevel, OrganizationLevel:
		n, ok := l[level+"/"+flow]
		return n, ok
	}
	return 0, false
}

// String renders the limits sorted by key, e.g. "{A/i=200, A/o=20, O/i=200}",
// the limits_map notation used in the counting transaction log.
func (l Limits) String() string {
	keys := make([]string, 0, len(l))
	for k := range l {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	var b strings.Builder
	b.WriteByte('{')
	for i, k := range keys {
		if i > 0 {
			b.WriteString(", ")
		}
		fmt.Fprintf(&b, "%s=%d", k, l[k])
	}
	b.WriteByte('}')
	return b.String()
}

// Domain is one parsed CDB record.
type Domain struct {
	Name   string
	OrgID  int64
	Limits Limits
}

// ParseRecord parses a CDB value "<org_id>|<limits>" for domain. A missing
// "|" yields a domain without limits; a missing or non-numeric org id is an
// error because the O counters cannot be keyed without it.
func ParseRecord(domain, value string) (*Domain, error) {
	orgPart, limitsPart, _ := strings.Cut(value, "|")
	orgPart = strings.TrimSpace(orgPart)
	orgID, err := strconv.ParseInt(orgPart, 10, 64)
	if err != nil || orgID < 0 {
		return nil, fmt.Errorf("%w: domain %s: invalid org id %q", ErrMalformed, domain, orgPart)
	}
	return &Domain{Name: domain, OrgID: orgID, Limits: ParseLimits(limitsPart)}, nil
}

// DB is an open domains CDB. It is safe for concurrent use.
//
// Parsed records are cached per DB (domain → *Domain) up to maxCached
// entries, because some values are huge (boxmailhere.com: 193 KB, 9 798
// overrides; ~1 ms and 1.4 MB to parse). The cache is never invalidated: a
// reload replaces the whole DB object, and the cache goes with it. The cached
// *Domain is shared read-only; callers must not mutate it.
type DB struct {
	path      string
	db        *cdb.CDB
	count     int
	maxCached int
	cache     sync.Map // domain -> *Domain
	cached    atomic.Int64
}

// Open opens the CDB at path with the default cache size.
func Open(path string) (*DB, error) { return OpenWithCache(path, DefaultMaxCached) }

// OpenWithCache opens the CDB at path and counts its records (a few
// milliseconds for ~100k records; the count feeds the cdb_domains gauge and
// the reload log). maxCached bounds the parsed-record cache; 0 disables it.
func OpenWithCache(path string, maxCached int) (*DB, error) {
	c, err := cdb.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open domains db %s: %w", path, err)
	}
	n := 0
	it := c.Iter()
	for it.Next() {
		n++
	}
	if err := it.Err(); err != nil {
		c.Close()
		return nil, fmt.Errorf("read domains db %s: %w", path, err)
	}
	if maxCached < 0 {
		maxCached = 0
	}
	return &DB{path: path, db: c, count: n, maxCached: maxCached}, nil
}

// Lookup returns the record for domain (trimmed and lowercased first), or
// ErrNotFound when the domain is not hosted here. The returned *Domain may
// be shared with other callers and must be treated as read-only.
func (d *DB) Lookup(domain string) (*Domain, error) {
	domain = strings.ToLower(strings.TrimSpace(domain))
	if domain == "" {
		return nil, ErrNotFound
	}
	if hit, ok := d.cache.Load(domain); ok {
		return hit.(*Domain), nil
	}
	v, err := d.db.Get([]byte(domain))
	if err != nil {
		return nil, fmt.Errorf("read domains db %s: %w", d.path, err)
	}
	if v == nil {
		return nil, ErrNotFound
	}
	rec, err := ParseRecord(domain, string(v))
	if err != nil {
		return nil, err
	}
	if d.maxCached > 0 && d.cached.Load() < int64(d.maxCached) {
		if prev, loaded := d.cache.LoadOrStore(domain, rec); loaded {
			return prev.(*Domain), nil
		}
		d.cached.Add(1)
	}
	return rec, nil
}

// Cached returns the number of parsed records currently cached.
func (d *DB) Cached() int { return int(d.cached.Load()) }

// Ping reads the first record and reports an error when the file can no
// longer be read or holds no domains at all (nothing would ever be counted).
func (d *DB) Ping() error {
	it := d.db.Iter()
	if it.Next() {
		return nil
	}
	if err := it.Err(); err != nil {
		return fmt.Errorf("read domains db %s: %w", d.path, err)
	}
	return fmt.Errorf("domains db %s is empty", d.path)
}

// Path returns the file the DB was opened from.
func (d *DB) Path() string { return d.path }

// Count returns the number of records counted at Open.
func (d *DB) Count() int { return d.count }

// Close closes the underlying file; in-flight Lookups then fail with an error.
func (d *DB) Close() error { return d.db.Close() }