counters/internal/counting/counting.go

// Package counting turns one ingested recipient into the counter increments
// it causes, and executes them against the store.
//
// Rules (see plans/2026-10-13-migadu-counters.md):
//
//   - /out (authenticated submission): the sender's A/D/O outgoing counters
//     +1 when the sender's domain is hosted here, and the recipient's A/D/O
//     incoming counters +1 when the recipient's domain is hosted here.
//   - /in (inbound from outside): the recipient's A/D/O incoming counters +1
//     when hosted here; the sender is never counted; junk is skipped unless
//     configured otherwise.
//   - Mail to or from an exempt address (support@, contact@) never counts.
//   - A domain that is not in the domains CDB is not counted.
//
// The Planner is pure (no I/O besides the domain lookup) so the rules are
// testable without Redis; Counter runs a Plan's increments in one pipeline.
package counting

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"migadu_counters/internal/domains"
	"migadu_counters/internal/store"
)

// Routes and skip reasons as reported in responses and the transaction log.
const (
	RouteIn  = "in"
	RouteOut = "out"

	SkipExempt        = "exempt"
	SkipJunk          = "junk"
	SkipUnknownDomain = "unknown_domain"
)

// Lookuper resolves a domain to its CDB record; *domains.DB implements it.
type Lookuper interface {
	Lookup(domain string) (*domains.Domain, error)
}

// LookupFunc adapts a function (e.g. one reading an atomic.Pointer to the
// current DB) to Lookuper.
type LookupFunc func(domain string) (*domains.Domain, error)

// Lookup implements Lookuper.
func (f LookupFunc) Lookup(domain string) (*domains.Domain, error) { return f(domain) }

// Planner applies the counting rules.
type Planner struct {
	Lookup    Lookuper
	Exempt    map[string]bool  // lowercased addresses that never count
	CountJunk bool             // count /in requests flagged junk
	Now       func() time.Time // nil = time.Now; the UTC day of Now() is the counter day
	// WarnPercent is the share of a limit that triggers an early warning; 0 = off.
	WarnPercent int64
}

// Job is one validated request handed to the queue: the fields the counting
// rules need and nothing else. A side that failed validation is empty, with
// the reason in SenderUnusable / RecipientUnusable for the log.
type Job struct {
	Route    string // RouteIn or RouteOut
	TraceID  string
	Received time.Time // when the request arrived (queue wait metric)

	AuthUser   string // /out: normalised SMTP AUTH user, "" when not usable
	TokenLogin bool   // /out: a user@dom@SUFFIX login
	MailFrom   string // normalised, or the sanitised input when not an address
	// MailFromUsable tells whether MailFrom is a normalised address (and may
	// match the exempt list); an unusable value is never exempt.
	MailFromUsable bool
	Recipient      string // normalised, "" when not usable
	Junk           bool   // /in only

	SenderUnusable    Unusable // /out: auth_user not usable (Value == "" when usable)
	RecipientUnusable Unusable // recipient not usable
}

// Unusable records an address side that could not be counted.
type Unusable struct {
	Role   string // "sender" or "recipient"
	Value  string // sanitised input, "" when the field was missing
	Reason string // "missing" or "invalid e-mail address"
}

// Set reports whether a side was recorded as unusable.
func (u Unusable) Set() bool { return u.Role != "" }

// Lookup records one domain resolution for the transaction log.
type Lookup struct {
	Role       string // "sender" or "recipient"
	Address    string
	DomainName string
	Domain     *domains.Domain // nil when not found
	Err        error           // domains.ErrNotFound when not hosted here
}

// Increment is one counter the request bumps.
type Increment struct {
	Level   string // A, D or O
	ID      string // address, domain name or organisation id
	Flow    string // i or o
	Key     string // Redis key for the day
	By      string // address that caused it (sender for o, recipient for i); $ADDRESS in notifications
	Limit   int64  // configured limit, 0 when unlimited
	Limited bool   // false when no limit is configured for this level/flow
	WarnAt  int64  // counter value that triggers the early warning, 0 = none
	Domain  *domains.Domain
}

// Plan is the result of applying the rules to one request.
type Plan struct {
	Route      string
	TraceID    string
	Date       string // UTC day used in the keys
	Skipped    string // "" or one of the Skip* reasons
	AuthUser   string
	TokenLogin bool
	MailFrom   string
	Recipient  string
	Junk       bool // /in only: the message was flagged junk by mizu

	Unusable   []Unusable // sides that were not counted because they failed validation
	Lookups    []Lookup
	Increments []Increment
}

func (p *Planner) now() time.Time {
	if p.Now != nil {
		return p.Now()
	}
	return time.Now()
}

func (p *Planner) isExempt(addr string) bool { return addr != "" && p.Exempt[addr] }

func normalize(addr string) string { return strings.ToLower(strings.TrimSpace(addr)) }

// splitAddress returns local part and domain; ok is false without an "@".
func splitAddress(addr string) (local, domain string, ok bool) {
	i := strings.LastIndex(addr, "@")
	if i <= 0 || i == len(addr)-1 {
		return "", "", false
	}
	return addr[:i], addr[i+1:], true
}

// PlanOut applies the /out rules for one authenticated message recipient
// whose fields all validated.
func (p *Planner) PlanOut(traceID, authUser, mailFrom, recipient string) (*Plan, error) {
	return p.Plan(Job{Route: RouteOut, TraceID: traceID, AuthUser: authUser, MailFrom: mailFrom, MailFromUsable: mailFrom != "", Recipient: recipient})
}

// PlanIn applies the /in rules for one inbound recipient. mailFrom may be
// empty (bounce sender <>).
func (p *Planner) PlanIn(traceID, mailFrom, recipient string, junk bool) (*Plan, error) {
	return p.Plan(Job{Route: RouteIn, TraceID: traceID, MailFrom: mailFrom, MailFromUsable: mailFrom != "", Recipient: recipient, Junk: junk})
}

// Plan applies the counting rules to a validated Job. Unusable sides are
// recorded for the log and skipped; the usable side is counted.
func (p *Planner) Plan(job Job) (*Plan, error) {
	authUser, mailFrom, recipient := normalize(job.AuthUser), job.MailFrom, normalize(job.Recipient)
	if job.MailFromUsable {
		mailFrom = normalize(mailFrom)
	}
	counted := job.Received
	if counted.IsZero() {
		counted = p.now()
	}
	plan := &Plan{Route: job.Route, TraceID: job.TraceID, Date: store.DayKey(counted),
		AuthUser: authUser, TokenLogin: job.TokenLogin, MailFrom: mailFrom, Recipient: recipient, Junk: job.Junk}
	if job.SenderUnusable.Set() {
		plan.Unusable = append(plan.Unusable, job.SenderUnusable)
	}
	if job.RecipientUnusable.Set() {
		plan.Unusable = append(plan.Unusable, job.RecipientUnusable)
	}
	if p.isExempt(domains.Mailbox(authUser)) || (job.MailFromUsable && p.isExempt(domains.Mailbox(mailFrom))) || p.isExempt(domains.Mailbox(recipient)) {
		plan.Skipped = SkipExempt
		return plan, nil
	}
	if job.Route == RouteIn && job.Junk && !p.CountJunk {
		plan.Skipped = SkipJunk
		return plan, nil
	}
	if job.Route == RouteOut && authUser != "" {
		if err := p.add(plan, "sender", authUser, domains.FlowOut); err != nil {
			return nil, err
		}
	}
	if recipient != "" {
		if err := p.add(plan, "recipient", recipient, domains.FlowIn); err != nil {
			return nil, err
		}
	}
	if len(plan.Increments) == 0 {
		plan.Skipped = SkipUnknownDomain
	}
	return plan, nil
}

// add looks up the address's domain and, when hosted here, appends its
// A, D and O increments for flow. The address is counted as its mailbox
// (plus-tag stripped). A not-found or malformed domain row is recorded and
// skipped; any other lookup error (unreadable CDB) aborts the plan.
func (p *Planner) add(plan *Plan, role, address, flow string) error {
	address = domains.Mailbox(address)
	local, domain, ok := splitAddress(address)
	lk := Lookup{Role: role, Address: address, DomainName: domain}
	if !ok {
		lk.Err = domains.ErrNotFound
		plan.Lookups = append(plan.Lookups, lk)
		return nil
	}
	d, err := p.Lookup.Lookup(domain)
	if err != nil {
		lk.Err = err
		plan.Lookups = append(plan.Lookups, lk)
		if errors.Is(err, domains.ErrNotFound) || errors.Is(err, domains.ErrMalformed) {
			return nil
		}
		return fmt.Errorf("lookup %s domain %s: %w", role, domain, err)
	}
	lk.Domain = d
	plan.Lookups = append(plan.Lookups, lk)

	levels := []struct{ level, id string }{
		{domains.AddressLevel, address},
		{domains.DomainLevel, d.Name},
		{domains.OrganizationLevel, strconv.FormatInt(d.OrgID, 10)},
	}
	for _, l := range levels {
		limit, limited := d.Limits.Limit(l.level, flow, local)
		plan.Increments = append(plan.Increments, Increment{
			Level: l.level, ID: l.id, Flow: flow,
			Key: store.Key(l.level, l.id, plan.Date, flow),
			By:  address, Limit: limit, Limited: limited, Domain: d,
			WarnAt: warnAt(limit, limited, p.WarnPercent),
		})
	}
	return nil
}

func warnAt(limit int64, limited bool, percent int64) int64 {
	if !limited {
		return 0
	}
	return WarnThreshold(limit, percent)
}

// WarnThreshold returns the counter value at which percent of limit is used,
// rounded up, or 0 when no warning applies: warnings are off, there is no
// positive limit, or the threshold is not below the limit itself.
func WarnThreshold(limit, percent int64) int64 {
	if limit <= 0 || percent <= 0 {
		return 0
	}
	t := limit/100*percent + (limit%100*percent+99)/100
	if t >= limit {
		return 0
	}
	return t
}

// Crossed reports whether an increment of amount that produced newValue is
// the one that reached limit: old < limit <= new. Because INCRBY is atomic
// exactly one increment can cross, so this is the once-per-crossing trigger
// for notifications; a limit of 0 (or unlimited) never crosses.
func Crossed(newValue, amount, limit int64) bool {
	return limit > 0 && newValue-amount < limit && limit <= newValue
}

// Result is an executed Increment.
type Result struct {
	Increment
	Count   int64 // value after the increment (from EXEC, or re-read after a lost reply)
	Reached bool  // this increment crossed the limit and is the one to notify
	Warned  bool  // this increment crossed the warning threshold and is the one to notify
	Err     error // this key's own INCR failed (e.g. it holds a non-number); Count is 0
	// MarkerErr is set when the notification marker could not be read or
	// written. After a lost reply the notification is then withheld.
	MarkerErr error
}

// Outcome is an executed Plan.
type Outcome struct {
	*Plan
	Results []Result
	// Unknown is set when the EXEC reply was lost after the write; Counts
	// were then re-read from the master and Reached means count >= limit.
	Unknown error
}

// Counter plans and executes.
type Counter struct {
	Planner       *Planner
	Store         *store.Store
	ReReadTimeout time.Duration // budget for the master re-read after a lost reply (default 2s)
	// AbortRetryDelay is the pause before the single retry of a transaction
	// Redis rejected, e.g. during a failover (default 200ms).
	AbortRetryDelay time.Duration
}

// CountOut plans and executes an /out request.
func (c *Counter) CountOut(ctx context.Context, traceID, authUser, mailFrom, recipient string) (*Outcome, error) {
	plan, err := c.Planner.PlanOut(traceID, authUser, mailFrom, recipient)
	if err != nil {
		return nil, err
	}
	return c.Execute(ctx, plan)
}

// CountIn plans and executes an /in request.
func (c *Counter) CountIn(ctx context.Context, traceID, mailFrom, recipient string, junk bool) (*Outcome, error) {
	plan, err := c.Planner.PlanIn(traceID, mailFrom, recipient, junk)
	if err != nil {
		return nil, err
	}
	return c.Execute(ctx, plan)
}

// Execute runs all increments of plan in one pipeline. A skipped plan does
// not touch Redis.
//
// Failure handling (review #4): a key whose own INCR failed is reported in
// its Result.Err while the other keys are counted normally. A transaction
// Redis rejected as a whole applied nothing and is retried once. When the
// whole reply was lost (timeout, reset) after the write may have been applied,
// the keys are re-read once from the master and the outcome is flagged
// Unknown. A definite failure (connection refused, cancelled) is returned as
// an error: nothing was counted.
//
// Notifications go out once per counter: a confirmed crossing is exact, so it
// always fires and writes a marker. After a lost reply the crossing cannot be
// pinned to this increment, so it fires only for the caller that wins the
// marker, and a counter that already notified stays silent.
func (c *Counter) Execute(ctx context.Context, plan *Plan) (*Outcome, error) {
	out := &Outcome{Plan: plan}
	if len(plan.Increments) == 0 {
		return out, nil
	}
	keys := make([]string, len(plan.Increments))
	for i, inc := range plan.Increments {
		keys[i] = inc.Key
	}
	out.Results = make([]Result, len(plan.Increments))
	results, err := c.incr(ctx, keys)
	if err == nil {
		for i, inc := range plan.Increments {
			r := Result{Increment: inc, Err: results[i].Err}
			if r.Err == nil {
				r.Count = results[i].Count
				r.Reached = inc.Limited && Crossed(r.Count, 1, inc.Limit)
				r.Warned = Crossed(r.Count, 1, inc.WarnAt)
				c.mark(ctx, plan.Date, &r)
			}
			out.Results[i] = r
		}
		return out, nil
	}
	if !store.OutcomeUnknown(err) {
		return nil, err
	}
	rctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), c.reReadTimeout())
	defer cancel()
	counts, rerr := c.Store.GetMaster(rctx, keys)
	if rerr != nil {
		return nil, fmt.Errorf("%w (outcome unknown, re-read failed: %v)", err, rerr)
	}
	out.Unknown = err
	for i, inc := range plan.Increments {
		r := Result{Increment: inc, Count: counts[i]}
		switch {
		case inc.Limited && inc.Limit > 0 && r.Count >= inc.Limit:
			r.Reached = c.claim(rctx, plan.Date, &r, store.MarkerLimit)
		case inc.WarnAt > 0 && r.Count >= inc.WarnAt:
			r.Warned = c.claim(rctx, plan.Date, &r, store.MarkerWarn)
		}
		out.Results[i] = r
	}
	return out, nil
}

// incr runs the transaction and retries it once when Redis rejected it.
func (c *Counter) incr(ctx context.Context, keys []string) ([]store.IncrResult, error) {
	results, err := c.Store.Incr(ctx, keys)
	if !errors.Is(err, store.ErrNotApplied) {
		return results, err
	}
	select {
	case <-time.After(c.abortRetryDelay()):
	case <-ctx.Done():
		return nil, err
	}
	return c.Store.Incr(ctx, keys)
}

// mark records a confirmed crossing. A failure is reported but the
// notification still goes out, because this crossing is certain.
func (c *Counter) mark(ctx context.Context, date string, r *Result) {
	if !r.Reached && !r.Warned {
		return
	}
	kind := store.MarkerLimit
	if r.Warned {
		kind = store.MarkerWarn
	}
	mctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), c.reReadTimeout())
	defer cancel()
	r.MarkerErr = c.Store.MarkNotified(mctx, store.MarkerKey(r.Level, r.ID, date, r.Flow, kind))
}

// claim takes the marker after a lost reply. If that fails nothing is sent:
// a missed notification is preferred to a repeated one.
func (c *Counter) claim(ctx context.Context, date string, r *Result, kind string) bool {
	won, err := c.Store.ClaimNotified(ctx, store.MarkerKey(r.Level, r.ID, date, r.Flow, kind))
	if err != nil {
		r.MarkerErr = err
		return false
	}
	return won
}

func (c *Counter) abortRetryDelay() time.Duration {
	if c.AbortRetryDelay > 0 {
		return c.AbortRetryDelay
	}
	return 200 * time.Millisecond
}

func (c *Counter) reReadTimeout() time.Duration {
	if c.ReReadTimeout > 0 {
		return c.ReReadTimeout
	}
	return 2 * time.Second
}