dimitrije/counters
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

First commit Go counter service

Browse Source
This commit is contained in:
Dimitrije Dimitrijevic
2026-10-15 02:12:18 +02:00
commit c43ff17627
59 changed files with 11205 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/api/api_test.go
+255
View File
@@ -0,0 +1,255 @@
package api
import (
"encoding/json"
"errors"
"strings"
"testing"
"time"
)
func TestValidateOut(t *testing.T) {
tests := []struct {
name string
in OutRequest
wantErr string // substring; "" = accepted
want ValidatedOut
}{
{"valid", OutRequest{TraceID: "tr", AuthUser: "x@example.com", MailFrom: "x@example.com", Recipient: "y@example.com"}, "",
ValidatedOut{TraceID: "tr", AuthUser: Address{Value: "x@example.com"}, MailFrom: "x@example.com", MailFromUsable: true, Recipient: Address{Value: "y@example.com"}}},
{"normalised", OutRequest{AuthUser: " X@Example.COM ", Recipient: "Y@EXAMPLE.com"}, "",
ValidatedOut{AuthUser: Address{Value: "x@example.com"}, Recipient: Address{Value: "y@example.com"}}},
// Support/token logins (review #3): mizu forwards the raw SASL login
// user@dom@SUFFIX; the base address is counted, the suffix logged.
{"login suffix", OutRequest{AuthUser: "x@yagolivas.com@SUPPORTTOKEN", Recipient: "y@example.com"}, "",
ValidatedOut{AuthUser: Address{Value: "x@yagolivas.com"}, LoginSuffix: "SUPPORTTOKEN", Recipient: Address{Value: "y@example.com"}}},
// mail_from never rejects (review #2): kept normalised when possible,
// otherwise as a sanitised string for the log.
{"mail_from empty", OutRequest{AuthUser: "x@example.com", Recipient: "y@example.com", MailFrom: ""}, "",
ValidatedOut{AuthUser: Address{Value: "x@example.com"}, Recipient: Address{Value: "y@example.com"}}},
{"mail_from garbage kept for log", OutRequest{AuthUser: "x@example.com", Recipient: "y@example.com", MailFrom: "john smith@x.com\x1b[1m"}, "",
ValidatedOut{AuthUser: Address{Value: "x@example.com"}, MailFrom: "john smith@x.com[1m", Recipient: Address{Value: "y@example.com"}}},
// One unusable side: the other is still counted.
{"auth_user missing, recipient ok", OutRequest{Recipient: "y@example.com"}, "",
ValidatedOut{AuthUser: Address{Reason: "missing"}, Recipient: Address{Value: "y@example.com"}}},
{"auth_user bad, recipient ok", OutRequest{AuthUser: "nobody", Recipient: "y@example.com"}, "",
ValidatedOut{AuthUser: Address{Raw: "nobody", Reason: "invalid e-mail address"}, Recipient: Address{Value: "y@example.com"}}},
{"recipient literal, sender ok", OutRequest{AuthUser: "x@example.com", Recipient: "u@[192.0.2.1]"}, "",
ValidatedOut{AuthUser: Address{Value: "x@example.com"}, Recipient: Address{Raw: "u@[192.0.2.1]", Reason: "invalid e-mail address"}}},
{"recipient underscore host, sender ok", OutRequest{AuthUser: "x@example.com", Recipient: "u@host_name.example.com"}, "",
ValidatedOut{AuthUser: Address{Value: "x@example.com"}, Recipient: Address{Raw: "u@host_name.example.com", Reason: "invalid e-mail address"}}},
// IDN domains are counted under their punycode CDB key.
{"idn recipient", OutRequest{AuthUser: "x@example.com", Recipient: "info@Bücher.de"}, "",
ValidatedOut{AuthUser: Address{Value: "x@example.com"}, Recipient: Address{Value: "info@xn--bcher-kva.de"}}},
{"trailing dot stripped", OutRequest{AuthUser: "x@example.com", Recipient: "u@example.com."}, "",
ValidatedOut{AuthUser: Address{Value: "x@example.com"}, Recipient: Address{Value: "u@example.com"}}},
// Neither side usable: 400 naming both.
{"both missing", OutRequest{}, "auth_user: missing; recipient: missing", ValidatedOut{}},
{"both bad", OutRequest{AuthUser: "nobody", Recipient: " "}, "auth_user: invalid e-mail address; recipient: missing", ValidatedOut{}},
{"trace sanitised", OutRequest{TraceID: "ab\ncd\x1b[31m" + strings.Repeat("z", 200), AuthUser: "x@example.com", Recipient: "y@example.com"}, "",
ValidatedOut{TraceID: "abcd[31m" + strings.Repeat("z", 120), AuthUser: Address{Value: "x@example.com"}, Recipient: Address{Value: "y@example.com"}}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := ValidateOut(tt.in)
if tt.wantErr != "" {
if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
t.Fatalf("err = %v, want containing %q", err, tt.wantErr)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if got != tt.want {
t.Errorf("got %+v\nwant %+v", got, tt.want)
}
})
}
}
func TestValidateIn(t *testing.T) {
tests := []struct {
name string
in InRequest
wantErr string
want ValidatedIn
}{
{"bounce sender", InRequest{TraceID: "t", MailFrom: "", Recipient: "Y@Example.com", Junk: true}, "",
ValidatedIn{TraceID: "t", Recipient: "y@example.com", Junk: true}},
{"valid sender", InRequest{MailFrom: "S@Gmail.com", Recipient: "y@example.com"}, "",
ValidatedIn{MailFrom: "s@gmail.com", MailFromUsable: true, Recipient: "y@example.com"}},
// Every odd-but-accepted-by-mizu sender keeps the recipient counted.
{"sender with space", InRequest{MailFrom: "john smith@example.com", Recipient: "y@example.com"}, "",
ValidatedIn{MailFrom: "john smith@example.com", Recipient: "y@example.com"}},
{"sender two ats", InRequest{MailFrom: "a@b@example.com", Recipient: "y@example.com"}, "",
ValidatedIn{MailFrom: "a@b@example.com", Recipient: "y@example.com"}},
{"sender idn", InRequest{MailFrom: "info@bücher.de", Recipient: "y@example.com"}, "",
ValidatedIn{MailFrom: "info@xn--bcher-kva.de", MailFromUsable: true, Recipient: "y@example.com"}},
{"sender literal", InRequest{MailFrom: "postmaster@[192.0.2.1]", Recipient: "y@example.com"}, "",
ValidatedIn{MailFrom: "postmaster@[192.0.2.1]", Recipient: "y@example.com"}},
{"sender too long", InRequest{MailFrom: strings.Repeat("v", 250) + "@example.com", Recipient: "y@example.com"}, "",
ValidatedIn{MailFrom: strings.Repeat("v", 250) + "@exa", Recipient: "y@example.com"}},
{"sender control chars stripped", InRequest{MailFrom: "a\r\nb@x.com", Recipient: "y@example.com"}, "",
ValidatedIn{MailFrom: "ab@x.com", Recipient: "y@example.com"}}, // sanitised for the log, not repaired into an address
{"recipient idn", InRequest{Recipient: "info@Bücher.de"}, "", ValidatedIn{Recipient: "info@xn--bcher-kva.de"}},
{"recipient missing", InRequest{MailFrom: "s@gmail.com"}, "recipient: missing", ValidatedIn{}},
{"recipient bad", InRequest{Recipient: "nobody"}, "recipient: invalid e-mail address", ValidatedIn{}},
{"recipient literal", InRequest{Recipient: "u@[192.0.2.1]"}, "recipient: invalid e-mail address", ValidatedIn{}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := ValidateIn(tt.in)
if tt.wantErr != "" {
if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
t.Fatalf("err = %v, want containing %q", err, tt.wantErr)
}
var fe *FieldError
if !errors.As(err, &fe) {
t.Errorf("error should be a *FieldError, got %T", err)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if got != tt.want {
t.Errorf("got %+v\nwant %+v", got, tt.want)
}
})
}
}
func TestNormalizeAuthUser(t *testing.T) {
tests := []struct {
in, addr, suffix, wantErr string
}{
{"x@yagolivas.com", "x@yagolivas.com", "", ""},
{" X@Yagolivas.com@SUPPORTTOKEN ", "x@yagolivas.com", "SUPPORTTOKEN", ""},
{"x@yagolivas.com@tok@en", "x@yagolivas.com", "tok@en", ""},
{"x@yagolivas.com@", "x@yagolivas.com", "", ""},
{"nobody", "", "", "exactly one @"},
{"@dom@SUFFIX", "", "", "exactly one @"},
{"", "", "", "empty"},
}
for _, tt := range tests {
addr, suffix, err := NormalizeAuthUser(tt.in)
if tt.wantErr != "" {
if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
t.Errorf("NormalizeAuthUser(%q) err = %v, want %q", tt.in, err, tt.wantErr)
}
continue
}
if err != nil || addr != tt.addr || suffix != tt.suffix {
t.Errorf("NormalizeAuthUser(%q) = %q, %q, %v; want %q, %q", tt.in, addr, suffix, err, tt.addr, tt.suffix)
}
}
}
func TestNormalizeEmail(t *testing.T) {
good := map[string]string{
" A.B+tag@Sub.Example.COM ": "a.b+tag@sub.example.com",
"x@example.com\t": "x@example.com",
"x@xn--exmple-cua.com": "x@xn--exmple-cua.com",
"jürgen@example.com": "jürgen@example.com", // SMTPUTF8 local part
"x@exämple.com": "x@xn--exmple-cua.com",
"u@münchen.de": "u@xn--mnchen-3ya.de",
"u@日本.jp": "u@xn--wgv71a.jp",
"u@example.com.": "u@example.com",
}
for in, want := range good {
if got, err := NormalizeEmail(in); err != nil || got != want {
t.Errorf("NormalizeEmail(%q) = %q, %v; want %q", in, got, err, want)
}
}
bad := []string{"", "nobody", "@x.com", "x@", "a@b@c.com", "x@exam ple.com", "x y@example.com", "x\x01@example.com",
strings.Repeat("a", 250) + "@e.com", "x@.example.com", "x@a..b.com", "x@[192.0.2.1]", "x@host_name.example.com", "x@-bad.com", "x@example.com.."}
for _, in := range bad {
if got, err := NormalizeEmail(in); err == nil {
t.Errorf("NormalizeEmail(%q) = %q, want error", in, got)
}
}
}
func TestSanitizeForLog(t *testing.T) {
if got := SanitizeForLog("a\r\nb\x1b[31m" + strings.Repeat("z", 300)); !strings.HasPrefix(got, "ab[31m") || len(got) != MaxEmailLen {
t.Errorf("got %q (len %d)", got, len(got))
}
}
func TestSanitizeTraceID(t *testing.T) {
if got := SanitizeTraceID(" a-b_c.d:e "); got != "a-b_c.d:e" {
t.Errorf("got %q", got)
}
if got := SanitizeTraceID("x\r\ny\x00z\x7f"); got != "xyz" {
t.Errorf("got %q", got)
}
if got := SanitizeTraceID(""); got != "" {
t.Errorf("got %q", got)
}
}
func TestParseDate(t *testing.T) {
now := time.Date(2026, 10, 13, 23, 30, 0, 0, time.FixedZone("CEST", 2*3600)) // 21:30 UTC on the 13th
tests := []struct {
in, want, wantErr string
}{
{"", "2026-10-13", ""},
{"2026-10-13", "2026-10-13", ""},
{"2026-10-12", "2026-10-12", ""},
{"2026-10-14", "", "future"},
{"2026-1-5", "", "YYYY-MM-DD"},
{"2026-13-01", "", "YYYY-MM-DD"},
{"13.10.2026", "", "YYYY-MM-DD"},
{"2026-10-13T00:00:00Z", "", "YYYY-MM-DD"},
}
for _, tt := range tests {
got, err := ParseDate(tt.in, now)
if tt.wantErr != "" {
if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
t.Errorf("ParseDate(%q) err = %v, want containing %q", tt.in, err, tt.wantErr)
}
continue
}
if err != nil || got != tt.want {
t.Errorf("ParseDate(%q) = %q, %v; want %q", tt.in, got, err, tt.want)
}
}
}
func TestCountResponseJSON(t *testing.T) {
b, err := json.Marshal(CountResponse{Accepted: true})
if err != nil {
t.Fatal(err)
}
if string(b) != `{"accepted":true}` {
t.Errorf("JSON = %s", b)
}
}
func TestCountersResponseJSON(t *testing.T) {
org := int64(1234)
in := int64(500)
resp := CountersResponse{Query: "me@example.com", Mode: "address", Date: "2026-10-13", Managed: true,
Domain: "example.com", OrgID: &org,
Levels: []CountersLevel{{Level: "A", ID: "me@example.com", In: 12, InLimit: &in, Out: 3, OutLimit: nil}}}
b, _ := json.Marshal(resp)
s := string(b)
for _, want := range []string{`"query":"me@example.com"`, `"mode":"address"`, `"date":"2026-10-13"`, `"managed":true`, `"exempt":false`,
`"domain":"example.com"`, `"org_id":1234`, `"levels":[`, `"in":12`, `"in_limit":500`, `"out":3`, `"out_limit":null`} {
if !strings.Contains(s, want) {
t.Errorf("JSON %s lacks %s", s, want)
}
}
b, _ = json.Marshal(CountersResponse{Query: "x@nowhere.tld", Mode: "address", Date: "2026-10-13", Levels: []CountersLevel{}})
if strings.Contains(string(b), "org_id") || strings.Contains(string(b), `"domain"`) {
t.Errorf("unmanaged response should omit org_id and domain: %s", b)
}
}
func TestErrorResponseJSON(t *testing.T) {
b, _ := json.Marshal(ErrorResponse{Error: "recipient: missing"})
if string(b) != `{"error":"recipient: missing"}` {
t.Errorf("got %s", b)
}
}
internal/api/types.go
+55
View File
@@ -0,0 +1,55 @@
// Package api holds the JSON contract between mailqueuer, the daemon and
// counterctl, plus the request validators (the "fields" step of the HTTP
// validation layer).
package api
// OutRequest is the body of POST /out: one recipient of an authenticated
// submission.
type OutRequest struct {
TraceID string `json:"trace_id"` // mizu's X-Trace-Id; log correlation only
AuthUser string `json:"auth_user"` // SMTP AUTH user (mizu's X-Auth-User)
MailFrom string `json:"mail_from"` // envelope sender, optional
Recipient string `json:"recipient"`
}
// InRequest is the body of POST /in: one recipient of an inbound message.
type InRequest struct {
TraceID string `json:"trace_id"`
MailFrom string `json:"mail_from"` // "" for a bounce sender <>
Recipient string `json:"recipient"`
Junk bool `json:"junk"` // mizu's X-Junk: yes
}
// CountResponse answers /in and /out with 202: the request was validated and
// queued; counting happens off the request path. The counted values appear in
// the counting log and in GET /counters.
type CountResponse struct {
Accepted bool `json:"accepted"`
}
// ErrorResponse is the body of every non-2xx answer.
type ErrorResponse struct {
Error string `json:"error"`
}
// CountersLevel is one row of GET /counters.
type CountersLevel struct {
Level string `json:"level"`
ID string `json:"id"`
In int64 `json:"in"`
InLimit *int64 `json:"in_limit"` // null when unlimited or unmanaged
Out int64 `json:"out"`
OutLimit *int64 `json:"out_limit"`
}
// CountersResponse answers GET /counters?address=|domain=[&date=].
type CountersResponse struct {
Query string `json:"query"`
Mode string `json:"mode"` // address | domain
Date string `json:"date"` // UTC day
Managed bool `json:"managed"`
Exempt bool `json:"exempt"`
Domain string `json:"domain,omitempty"`
OrgID *int64 `json:"org_id,omitempty"`
Levels []CountersLevel `json:"levels"`
}
internal/api/validate.go
+223
View File
@@ -0,0 +1,223 @@
package api
import (
"errors"
"fmt"
"strings"
"time"
"unicode"
"golang.org/x/net/idna"
)
// Limits applied by the validators.
const (
MaxEmailLen = 254
MaxTraceIDLen = 128
)
// FieldError names the offending field: "recipient: missing".
type FieldError struct {
Field string
Reason string
}
func (e *FieldError) Error() string { return e.Field + ": " + e.Reason }
// Address is one validated address side of a request. A side that is not
// usable is logged (Raw, Reason) but never rejects the request on its own:
// the other side is still counted.
type Address struct {
Value string // normalised address; "" when not usable
Raw string // sanitised input, for the log, when not usable
Reason string // "" when usable; "missing" or "invalid e-mail address"
}
// Usable reports whether the side can be counted.
func (a Address) Usable() bool { return a.Value != "" }
// ValidatedOut is an accepted /out body.
type ValidatedOut struct {
TraceID string
AuthUser Address
LoginSuffix string // "SUPPORTTOKEN" of a user@dom@SUPPORTTOKEN login
MailFrom string // normalised, or the sanitised input when not an address
MailFromUsable bool
Recipient Address
}
// ValidatedIn is an accepted /in body.
type ValidatedIn struct {
TraceID string
MailFrom string // normalised, or the sanitised input when not an address
MailFromUsable bool
Recipient string
Junk bool
}
// ValidateOut checks an /out body. auth_user and recipient are validated
// independently and the request is rejected only when neither is usable;
// mail_from only feeds the exempt check and never rejects.
func ValidateOut(r OutRequest) (ValidatedOut, error) {
v := ValidatedOut{TraceID: SanitizeTraceID(r.TraceID)}
v.AuthUser, v.LoginSuffix = validateAuthUser(r.AuthUser)
v.Recipient = validateAddress(r.Recipient)
v.MailFrom, v.MailFromUsable = lenient(r.MailFrom)
if !v.AuthUser.Usable() && !v.Recipient.Usable() {
return ValidatedOut{}, fmt.Errorf("%w; %w", &FieldError{"auth_user", v.AuthUser.Reason}, &FieldError{"recipient", v.Recipient.Reason})
}
return v, nil
}
// ValidateIn checks an /in body: recipient is required and must be an
// address; mail_from never rejects (it may be "" for a bounce sender).
func ValidateIn(r InRequest) (ValidatedIn, error) {
v := ValidatedIn{TraceID: SanitizeTraceID(r.TraceID), Junk: r.Junk}
rcpt := validateAddress(r.Recipient)
if !rcpt.Usable() {
return ValidatedIn{}, &FieldError{"recipient", rcpt.Reason}
}
v.Recipient = rcpt.Value
v.MailFrom, v.MailFromUsable = lenient(r.MailFrom)
return v, nil
}
func validateAddress(s string) Address {
if strings.TrimSpace(s) == "" {
return Address{Reason: "missing"}
}
n, err := NormalizeEmail(s)
if err != nil {
return Address{Raw: SanitizeForLog(s), Reason: "invalid e-mail address"}
}
return Address{Value: n}
}
func validateAuthUser(s string) (Address, string) {
if strings.TrimSpace(s) == "" {
return Address{Reason: "missing"}, ""
}
n, suffix, err := NormalizeAuthUser(s)
if err != nil {
return Address{Raw: SanitizeForLog(s), Reason: "invalid e-mail address"}, ""
}
return Address{Value: n}, suffix
}
// lenient normalises when possible and otherwise keeps a sanitised copy for
// the log; "" stays "" (bounce sender).
func lenient(s string) (string, bool) {
if strings.TrimSpace(s) == "" {
return "", false
}
if n, err := NormalizeEmail(s); err == nil {
return n, true
}
return SanitizeForLog(s), false
}
// NormalizeAuthUser normalises an SMTP AUTH login. mizu forwards the raw SASL
// user name, and support/token logins have the form user@domain@SUFFIX: the
// base address is everything up to the second "@", the rest is returned as
// the suffix for the log.
func NormalizeAuthUser(s string) (addr, suffix string, err error) {
s = strings.TrimSpace(s)
if first := strings.Index(s, "@"); first >= 0 {
if second := strings.Index(s[first+1:], "@"); second >= 0 {
base := s[:first+1+second]
suffix = s[first+1+second+1:]
s = base
}
}
addr, err = NormalizeEmail(s)
if err != nil {
return "", "", err
}
return addr, suffix, nil
}
// NormalizeEmail trims, lowercases and syntax-checks an address. The check
// is deliberately permissive — the domains CDB decides what is counted — it
// only rejects garbage: no "@" (or more than one), empty parts, whitespace or
// control characters, over-long input. Non-ASCII local parts are accepted
// (SMTPUTF8); a Unicode domain is converted to its punycode form, which is
// how the CDB keys it. A trailing dot on the domain is dropped. Address
// literals ([192.0.2.1]) and underscores in the domain are rejected: such
// domains are never hosted here.
func NormalizeEmail(s string) (string, error) {
s = strings.ToLower(strings.TrimSpace(s))
if s == "" || len(s) > MaxEmailLen {
return "", errors.New("empty or longer than 254 characters")
}
at := strings.Index(s, "@")
if at <= 0 || at != strings.LastIndex(s, "@") || at == len(s)-1 {
return "", errors.New("must be local@domain with exactly one @")
}
local, domain := s[:at], s[at+1:]
for _, r := range local {
if r <= ' ' || r == 0x7f || unicode.IsSpace(r) || unicode.IsControl(r) {
return "", errors.New("local part contains whitespace or control characters")
}
}
domain = strings.TrimSuffix(domain, ".")
if domain == "" || strings.HasPrefix(domain, ".") || strings.HasSuffix(domain, ".") || strings.Contains(domain, "..") {
return "", errors.New("domain has an empty label")
}
ascii, err := idna.Lookup.ToASCII(domain)
if err != nil {
return "", fmt.Errorf("domain is not a valid host name: %v", err)
}
for i := 0; i < len(ascii); i++ {
c := ascii[i]
if !(c >= 'a' && c <= 'z' || c >= '0' && c <= '9' || c == '-' || c == '.') {
return "", errors.New("domain contains characters other than a-z, 0-9, '-' and '.'")
}
}
out := local + "@" + ascii
if len(out) > MaxEmailLen {
return "", errors.New("longer than 254 characters")
}
return out, nil
}
// SanitizeTraceID trims the id, strips control characters (it ends up in log
// lines) and caps it at MaxTraceIDLen bytes.
func SanitizeTraceID(s string) string { return sanitize(s, MaxTraceIDLen) }
// SanitizeForLog strips control characters from an arbitrary field value and
// caps it at MaxEmailLen bytes, so unusable addresses can be logged safely.
func SanitizeForLog(s string) string { return sanitize(s, MaxEmailLen) }
func sanitize(s string, max int) string {
var b strings.Builder
for _, r := range strings.TrimSpace(s) {
if r < 0x20 || r == 0x7f || unicode.IsControl(r) {
continue
}
if b.Len()+len(string(r)) > max {
break
}
b.WriteRune(r)
}
return b.String()
}
// ParseDate validates a /counters or counterctl date: "" means today (UTC),
// otherwise an ISO 8601 calendar date YYYY-MM-DD that is not in the future.
// It returns the normalised YYYY-MM-DD string.
func ParseDate(s string, now time.Time) (string, error) {
y, m, d := now.UTC().Date()
today := time.Date(y, m, d, 0, 0, 0, 0, time.UTC)
s = strings.TrimSpace(s)
if s == "" {
return today.Format("2006-01-02"), nil
}
day, err := time.Parse("2006-01-02", s)
if err != nil {
return "", fmt.Errorf("invalid date %q: use ISO 8601 YYYY-MM-DD", s)
}
if day.After(today) {
return "", fmt.Errorf("date %s is in the future (today is %s UTC)", s, today.Format("2006-01-02"))
}
return day.Format("2006-01-02"), nil
}
internal/config/config.go
+369
View File
@@ -0,0 +1,369 @@
// Package config loads and validates the migadu_counters TOML configuration,
// shared by the daemon (cmd/migadu_counters) and the CLI (cmd/counterctl).
package config
import (
"errors"
"fmt"
"net"
"net/netip"
"net/url"
"os"
"path/filepath"
"strings"
"time"
"github.com/BurntSushi/toml"
)
// DefaultPaths are tried in order when no -config flag is given.
var DefaultPaths = []string{
"/etc/migadu_counters/config.toml",
"/usr/local/etc/migadu_counters/config.toml",
}
// Environment variables that override the secrets from the file.
const (
EnvAuthToken = "MIGADU_COUNTERS_AUTH_TOKEN"
EnvRedisPassword = "MIGADU_COUNTERS_REDIS_PASSWORD"
EnvNotifyPassword = "MIGADU_COUNTERS_NOTIFY_PASSWORD"
)
// NotifyPlaceholders are the template variables every notify URL is expected
// to carry: the web app's /api/limits action pattern-matches all six query
// parameters and answers 500 when one is missing.
var NotifyPlaceholders = []string{"$LEVEL", "$ID", "$FLOW", "$PERIOD", "$ADDRESS", "$COUNT"}
// Config is the full daemon configuration.
type Config struct {
Addr string `toml:"addr"`
AdminAddr string `toml:"admin_addr"`
AuthToken string `toml:"auth_token"`
AllowedCIDRs []string `toml:"allowed_cidrs"`
DomainsCDB string `toml:"domains_cdb"`
Exempt []string `toml:"exempt"`
CountJunk bool `toml:"count_junk"`
// Counting runs off the request path: /in and /out enqueue a job and
// answer 202; Workers goroutines drain a queue of QueueSize jobs.
QueueSize int `toml:"queue_size"`
Workers int `toml:"workers"`
// Parsed CDB records are cached per loaded table up to this many
// domains; beyond it lookups parse without caching. 0 disables the cache.
MaxCachedDomains int `toml:"max_cached_domains"`
// HTTP server timeouts in seconds.
ReadHeaderTimeout int `toml:"read_header_timeout"`
WriteTimeout int `toml:"write_timeout"`
IdleTimeout int `toml:"idle_timeout"`
HandlerTimeout int `toml:"handler_timeout"`
Logging LoggingConfig `toml:"logging"`
Redis RedisConfig `toml:"redis"`
Notify NotifyConfig `toml:"notify"`
prefixes []netip.Prefix
exempt map[string]bool
}
// LoggingConfig holds the application log and the counting transaction log
// settings (the first three fields match the rcptd/mailqueuer/mizu pattern).
type LoggingConfig struct {
Output string `toml:"output"` // "stderr", "stdout", or a file path
Format string `toml:"format"` // "console" or "json"
Level string `toml:"level"` // "debug", "info", "warn", "error"
CountingOutput string `toml:"counting_output"` // default: same as Output; "off" disables
CountingFormat string `toml:"counting_format"` // "text" (Lua-like lines) or "json"
}
// CountingEnabled reports whether the counting transaction log is written.
func (l LoggingConfig) CountingEnabled() bool { return l.CountingOutput != "off" }
// RedisConfig describes the write (master) and read (replica) sides.
// Sentinel mode is used when Sentinels is non-empty; otherwise the static
// WriteAddr/ReadAddrs are used (development, tests).
type RedisConfig struct {
DB int `toml:"db"`
Password string `toml:"password"`
PoolSize int `toml:"pool_size"`
TimeoutMS int `toml:"timeout_ms"`
MasterName string `toml:"master_name"`
Sentinels []string `toml:"sentinels"`
SentinelPassword string `toml:"sentinel_password"`
WriteAddr string `toml:"write_addr"`
ReadAddrs []string `toml:"read_addrs"`
}
// SentinelMode reports whether addresses are discovered through Sentinel.
func (r RedisConfig) SentinelMode() bool { return len(r.Sentinels) > 0 }
// Timeout is the dial/read/write timeout for every Redis command.
func (r RedisConfig) Timeout() time.Duration { return time.Duration(r.TimeoutMS) * time.Millisecond }
// NotifyConfig holds the per-level web app notification URLs. An empty URL
// disables notifications for that level.
type NotifyConfig struct {
OrgURL string `toml:"org_url"`
AddressURL string `toml:"address_url"`
DomainURL string `toml:"domain_url"`
User string `toml:"user"`
Password string `toml:"password"`
TimeoutSeconds int `toml:"timeout_seconds"`
}
// Timeout is the per-attempt HTTP timeout for a notification.
func (n NotifyConfig) Timeout() time.Duration { return time.Duration(n.TimeoutSeconds) * time.Second }
// Default returns the configuration with every default applied.
func Default() Config {
return Config{
Addr: ":5555",
AdminAddr: ":9095",
Exempt: []string{"support@migadu.com", "contact@migadu.com"},
QueueSize: 10000,
Workers: 8,
MaxCachedDomains: 20000,
ReadHeaderTimeout: 5,
WriteTimeout: 30,
IdleTimeout: 60,
HandlerTimeout: 4,
Logging: LoggingConfig{Output: "stderr", Format: "console", Level: "info", CountingFormat: "text"},
Redis: RedisConfig{PoolSize: 32, TimeoutMS: 500},
Notify: NotifyConfig{TimeoutSeconds: 5},
}
}
// Discover returns flagPath when given, otherwise the first existing entry of
// DefaultPaths, or an error naming every path that was tried.
func Discover(flagPath string) (string, error) {
if flagPath != "" {
return flagPath, nil
}
for _, p := range DefaultPaths {
if _, err := os.Stat(p); err == nil {
return p, nil
}
}
return "", fmt.Errorf("no config found at %s; specify -config PATH", strings.Join(DefaultPaths, " or "))
}
// Load reads the TOML file at path, applies defaults, environment overrides
// and normalisation, and validates the result. Unknown keys are fatal: a typo
// such as "auth-token" must not start the daemon with the default in place.
func Load(path string) (*Config, error) {
cfg := Default()
md, err := toml.DecodeFile(path, &cfg)
if err != nil {
return nil, fmt.Errorf("read config %s: %w", path, err)
}
var unknown []error
for _, k := range md.Undecoded() {
unknown = append(unknown, fmt.Errorf("unknown config key %s", k.String()))
}
if len(unknown) > 0 {
return nil, fmt.Errorf("config %s: %w", path, errors.Join(unknown...))
}
if v := os.Getenv(EnvAuthToken); v != "" {
cfg.AuthToken = v
}
if v := os.Getenv(EnvRedisPassword); v != "" {
cfg.Redis.Password = v
}
if v := os.Getenv(EnvNotifyPassword); v != "" {
cfg.Notify.Password = v
}
// A relative CDB path is resolved against the config file's directory,
// like rcptd's cdb_tables_dir.
if cfg.DomainsCDB != "" && !filepath.IsAbs(cfg.DomainsCDB) {
cfg.DomainsCDB = filepath.Join(filepath.Dir(path), cfg.DomainsCDB)
}
cfg.normalize()
if err := cfg.Validate(); err != nil {
return nil, fmt.Errorf("config %s: %w", path, err)
}
return &cfg, nil
}
func (c *Config) normalize() {
c.exempt = make(map[string]bool, len(c.Exempt))
for _, a := range c.Exempt {
if a = strings.ToLower(strings.TrimSpace(a)); a != "" {
c.exempt[a] = true
}
}
c.Logging.Format = strings.ToLower(strings.TrimSpace(c.Logging.Format))
c.Logging.Level = strings.ToLower(strings.TrimSpace(c.Logging.Level))
c.Logging.CountingFormat = strings.ToLower(strings.TrimSpace(c.Logging.CountingFormat))
c.Logging.CountingOutput = strings.TrimSpace(c.Logging.CountingOutput)
switch {
case c.Logging.CountingOutput == "":
c.Logging.CountingOutput = c.Logging.Output
case strings.EqualFold(c.Logging.CountingOutput, "off"):
c.Logging.CountingOutput = "off"
}
if !c.Redis.SentinelMode() && len(c.Redis.ReadAddrs) == 0 && c.Redis.WriteAddr != "" {
c.Redis.ReadAddrs = []string{c.Redis.WriteAddr}
}
}
// Validate checks the configuration and returns every problem found, joined.
func (c *Config) Validate() error {
var errs []error
fail := func(format string, args ...any) { errs = append(errs, fmt.Errorf(format, args...)) }
if c.DomainsCDB == "" {
fail("domains_cdb is required")
}
if strings.TrimSpace(c.AuthToken) == "" {
fail("auth_token is required; set it in the config or %s", EnvAuthToken)
}
if c.QueueSize <= 0 {
fail("queue_size must be > 0")
}
if c.Workers <= 0 {
fail("workers must be > 0")
}
if c.MaxCachedDomains < 0 {
fail("max_cached_domains must be >= 0")
}
for _, a := range []struct{ name, addr string }{{"addr", c.Addr}, {"admin_addr", c.AdminAddr}} {
if _, _, err := net.SplitHostPort(a.addr); err != nil {
fail("%s: invalid listen address %q: %v", a.name, a.addr, err)
}
}
for _, d := range []struct {
name string
v int
}{{"read_header_timeout", c.ReadHeaderTimeout}, {"write_timeout", c.WriteTimeout}, {"idle_timeout", c.IdleTimeout}, {"handler_timeout", c.HandlerTimeout}} {
if d.v <= 0 {
fail("%s must be > 0 (seconds)", d.name)
}
}
c.prefixes = c.prefixes[:0]
for _, s := range c.AllowedCIDRs {
p, err := parsePrefix(s)
if err != nil {
fail("allowed_cidrs: invalid entry %q: %v", s, err)
continue
}
c.prefixes = append(c.prefixes, p)
}
switch c.Logging.Format {
case "console", "json":
default:
fail("logging.format: %q (must be console or json)", c.Logging.Format)
}
if _, err := parseLogLevel(c.Logging.Level); err != nil {
fail("logging.level: %v", err)
}
switch c.Logging.CountingFormat {
case "text", "json":
default:
fail("logging.counting_format: %q (must be text or json)", c.Logging.CountingFormat)
}
if c.Redis.DB != 0 {
fail("redis.db must be 0: migadu_web reads the message counters from DB 0 (got %d)", c.Redis.DB)
}
if c.Redis.PoolSize <= 0 {
fail("redis.pool_size must be > 0")
}
if c.Redis.TimeoutMS <= 0 {
fail("redis.timeout_ms must be > 0")
}
if c.Redis.SentinelMode() {
if c.Redis.MasterName == "" {
fail("redis.master_name is required when redis.sentinels is set")
}
for _, s := range c.Redis.Sentinels {
if _, _, err := net.SplitHostPort(s); err != nil {
fail("redis.sentinels: invalid address %q: %v", s, err)
}
}
} else {
if c.Redis.WriteAddr == "" {
fail("redis.write_addr is required (or configure redis.sentinels)")
} else if _, _, err := net.SplitHostPort(c.Redis.WriteAddr); err != nil {
fail("redis.write_addr: invalid address %q: %v", c.Redis.WriteAddr, err)
}
for _, s := range c.Redis.ReadAddrs {
if _, _, err := net.SplitHostPort(s); err != nil {
fail("redis.read_addrs: invalid address %q: %v", s, err)
}
}
}
for _, u := range c.notifyURLs() {
if u.url == "" {
continue
}
parsed, err := url.Parse(u.url)
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Host == "" {
fail("notify.%s: must be an absolute http(s) URL (got %q)", u.name, u.url)
}
}
if c.Notify.TimeoutSeconds <= 0 {
fail("notify.timeout_seconds must be > 0")
}
return errors.Join(errs...)
}
// Warnings returns non-fatal findings worth logging at startup.
func (c *Config) Warnings() []string {
var w []string
for _, u := range c.notifyURLs() {
if u.url == "" {
continue
}
for _, ph := range NotifyPlaceholders {
if !strings.Contains(u.url, ph) {
w = append(w, fmt.Sprintf("notify.%s: missing %s placeholder (the web app requires all six query parameters)", u.name, ph))
}
}
}
return w
}
func (c *Config) notifyURLs() []struct{ name, url string } {
return []struct{ name, url string }{
{"org_url", c.Notify.OrgURL},
{"address_url", c.Notify.AddressURL},
{"domain_url", c.Notify.DomainURL},
}
}
// ExemptSet returns the normalised exempt addresses.
func (c *Config) ExemptSet() map[string]bool { return c.exempt }
// AllowedPrefixes returns the parsed allowed_cidrs; empty means any source.
func (c *Config) AllowedPrefixes() []netip.Prefix { return c.prefixes }
// HandlerTimeoutDuration is the http.TimeoutHandler budget.
func (c *Config) HandlerTimeoutDuration() time.Duration {
return time.Duration(c.HandlerTimeout) * time.Second
}
// parsePrefix accepts "10.0.0.0/8" as well as a bare address ("192.168.1.5",
// "::1"), which becomes a single-host prefix.
func parsePrefix(s string) (netip.Prefix, error) {
s = strings.TrimSpace(s)
if p, err := netip.ParsePrefix(s); err == nil {
return p.Masked(), nil
}
a, err := netip.ParseAddr(s)
if err != nil {
return netip.Prefix{}, errors.New("not a CIDR prefix or IP address")
}
return netip.PrefixFrom(a, a.BitLen()), nil
}
internal/config/config_test.go
+288
View File
@@ -0,0 +1,288 @@
package config
import (
"os"
"path/filepath"
"strings"
"testing"
"time"
)
// writeConfig writes body to <tmpdir>/config.toml and returns its path.
func writeConfig(t *testing.T, body string) string {
t.Helper()
path := filepath.Join(t.TempDir(), "config.toml")
if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
t.Fatal(err)
}
return path
}
// TOML: top-level keys must precede any [table], and a table may be defined
// only once — fixtures are therefore assembled as top + ... + redisStatic.
const (
top = "domains_cdb = \"/tmp/rspamd_domains_db\"\nauth_token = \"t\"\n"
redisStatic = "[redis]\nwrite_addr = \"127.0.0.1:6379\"\n"
minimal = top + redisStatic
)
func mustLoad(t *testing.T, body string) *Config {
t.Helper()
cfg, err := Load(writeConfig(t, body))
if err != nil {
t.Fatalf("Load: %v\n--- config ---\n%s", err, body)
}
return cfg
}
func TestLoad_Defaults(t *testing.T) {
cfg := mustLoad(t, minimal)
checks := []struct {
name string
got any
want any
}{
{"addr", cfg.Addr, ":5555"},
{"admin_addr", cfg.AdminAddr, ":9095"},
{"auth_token", cfg.AuthToken, "t"},
{"queue_size", cfg.QueueSize, 10000},
{"workers", cfg.Workers, 8},
{"max_cached_domains", cfg.MaxCachedDomains, 20000},
{"count_junk", cfg.CountJunk, false},
{"read_header_timeout", cfg.ReadHeaderTimeout, 5},
{"write_timeout", cfg.WriteTimeout, 30},
{"idle_timeout", cfg.IdleTimeout, 60},
{"handler_timeout", cfg.HandlerTimeout, 4},
{"logging.output", cfg.Logging.Output, "stderr"},
{"logging.format", cfg.Logging.Format, "console"},
{"logging.level", cfg.Logging.Level, "info"},
{"logging.counting_output", cfg.Logging.CountingOutput, "stderr"},
{"logging.counting_format", cfg.Logging.CountingFormat, "text"},
{"redis.db", cfg.Redis.DB, 0},
{"redis.pool_size", cfg.Redis.PoolSize, 32},
{"redis.timeout_ms", cfg.Redis.TimeoutMS, 500},
{"redis.sentinel_mode", cfg.Redis.SentinelMode(), false},
{"notify.timeout_seconds", cfg.Notify.TimeoutSeconds, 5},
{"redis timeout", cfg.Redis.Timeout(), 500 * time.Millisecond},
{"notify timeout", cfg.Notify.Timeout(), 5 * time.Second},
{"handler timeout", cfg.HandlerTimeoutDuration(), 4 * time.Second},
}
for _, c := range checks {
if c.got != c.want {
t.Errorf("%s = %v, want %v", c.name, c.got, c.want)
}
}
if got := cfg.Redis.ReadAddrs; len(got) != 1 || got[0] != "127.0.0.1:6379" {
t.Errorf("read_addrs = %v, want [write_addr]", got)
}
if got := cfg.ExemptSet(); len(got) != 2 || !got["support@migadu.com"] || !got["contact@migadu.com"] {
t.Errorf("ExemptSet() = %v, want support@ and contact@", got)
}
if len(cfg.AllowedPrefixes()) != 0 {
t.Errorf("AllowedPrefixes() = %v, want none", cfg.AllowedPrefixes())
}
if !cfg.Logging.CountingEnabled() {
t.Error("counting log should be enabled by default")
}
}
// The shipped example must always load, validate and produce no warnings.
func TestLoad_Example(t *testing.T) {
cfg, err := Load("../../config.toml.example")
if err != nil {
t.Fatalf("Load(config.toml.example): %v", err)
}
if !cfg.Redis.SentinelMode() || cfg.Redis.MasterName != "mymaster" || len(cfg.Redis.Sentinels) != 3 {
t.Errorf("example should be in sentinel mode with 3 sentinels: %+v", cfg.Redis)
}
if cfg.Redis.DB != 0 {
t.Errorf("example redis.db = %d, must be 0", cfg.Redis.DB)
}
if cfg.Notify.OrgURL == "" || cfg.Notify.AddressURL == "" {
t.Error("example should set org_url and address_url")
}
if cfg.Notify.DomainURL != "" {
t.Error("example should leave domain_url empty (web app D branch is a TODO)")
}
for _, w := range cfg.Warnings() {
t.Errorf("example produced warning: %s", w)
}
}
func TestLoad_MissingFile(t *testing.T) {
if _, err := Load(filepath.Join(t.TempDir(), "nope.toml")); err == nil {
t.Fatal("expected error")
}
}
func TestLoad_BadTOML(t *testing.T) {
if _, err := Load(writeConfig(t, "addr = [unterminated")); err == nil {
t.Fatal("expected error")
}
}
func TestLoad_RelativeDomainsCDBAnchoredToConfigDir(t *testing.T) {
path := writeConfig(t, "domains_cdb = \"db/rspamd_domains_db\"\nauth_token = \"t\"\n"+redisStatic)
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
want := filepath.Join(filepath.Dir(path), "db/rspamd_domains_db")
if cfg.DomainsCDB != want {
t.Errorf("domains_cdb = %q, want %q", cfg.DomainsCDB, want)
}
}
func TestLoad_EnvOverrides(t *testing.T) {
t.Setenv("MIGADU_COUNTERS_AUTH_TOKEN", "env-token")
t.Setenv("MIGADU_COUNTERS_REDIS_PASSWORD", "env-redis")
t.Setenv("MIGADU_COUNTERS_NOTIFY_PASSWORD", "env-notify")
cfg := mustLoad(t, "domains_cdb = \"/tmp/x\"\nauth_token = \"file-token\"\n"+redisStatic+"password = \"file-redis\"\n[notify]\npassword = \"file-notify\"\n")
if cfg.AuthToken != "env-token" || cfg.Redis.Password != "env-redis" || cfg.Notify.Password != "env-notify" {
t.Errorf("env overrides not applied: %q %q %q", cfg.AuthToken, cfg.Redis.Password, cfg.Notify.Password)
}
}
func TestLoad_ExemptNormalized(t *testing.T) {
cfg := mustLoad(t, top+"exempt = [\" Support@Migadu.com \", \"\", \"x@y.z\"]\n"+redisStatic)
got := cfg.ExemptSet()
if len(got) != 2 || !got["support@migadu.com"] || !got["x@y.z"] {
t.Errorf("ExemptSet() = %v, want only the configured entries, normalised", got)
}
// An explicitly empty list disables the exemption entirely.
if got := mustLoad(t, top+"exempt = []\n"+redisStatic).ExemptSet(); len(got) != 0 {
t.Errorf("exempt = [] gave %v, want none", got)
}
}
func TestLoad_AllowedCIDRs(t *testing.T) {
cfg := mustLoad(t, top+"allowed_cidrs = [\"10.0.0.0/8\", \"192.168.1.5\", \"::1\"]\n"+redisStatic)
p := cfg.AllowedPrefixes()
if len(p) != 3 || p[0].Bits() != 8 || p[1].Bits() != 32 || p[2].Bits() != 128 {
t.Errorf("AllowedPrefixes() = %v", p)
}
}
func TestLoad_SentinelMode(t *testing.T) {
cfg := mustLoad(t, top+"[redis]\nmaster_name = \"mymaster\"\nsentinels = [\"10.0.0.1:26379\", \"10.0.0.2:26379\"]\n")
if !cfg.Redis.SentinelMode() || len(cfg.Redis.ReadAddrs) != 0 || cfg.Redis.WriteAddr != "" {
t.Errorf("sentinel mode config = %+v", cfg.Redis)
}
}
func TestValidate_Errors(t *testing.T) {
tests := []struct {
name, toml, wantErr string
}{
{"redis db must be 0", top + redisStatic + "db = 1\n", "redis.db"},
{"domains_cdb required", "auth_token = \"t\"\n" + redisStatic, "domains_cdb"},
{"auth_token required", "domains_cdb = \"/tmp/x\"\n" + redisStatic, "auth_token is required"},
{"auth_token blank", "domains_cdb = \"/tmp/x\"\nauth_token = \" \"\n" + redisStatic, "auth_token is required"},
{"typo auth-token", "domains_cdb = \"/tmp/x\"\nauth-token = \"s3cret\"\n" + redisStatic, "unknown config key auth-token"},
{"typo allowed_cidr", top + "allowed_cidr = [\"10.0.0.0/8\"]\n" + redisStatic, "unknown config key allowed_cidr"},
{"typo in table", top + redisStatic + "pool = 3\n", "unknown config key redis.pool"},
{"queue_size", top + "queue_size = 0\n" + redisStatic, "queue_size"},
{"workers", top + "workers = 0\n" + redisStatic, "workers"},
{"max_cached_domains", top + "max_cached_domains = -1\n" + redisStatic, "max_cached_domains"},
{"sentinel needs master_name", top + "[redis]\nsentinels = [\"10.0.0.1:26379\"]\n", "master_name"},
{"sentinel bad address", top + "[redis]\nmaster_name = \"m\"\nsentinels = [\"10.0.0.1\"]\n", "sentinels"},
{"static needs write_addr", top, "write_addr"},
{"static bad read_addr", top + redisStatic + "read_addrs = [\"nohost\"]\n", "read_addrs"},
{"bad cidr mask", top + "allowed_cidrs = [\"10.0.0.0/33\"]\n" + redisStatic, "allowed_cidrs"},
{"bad cidr text", top + "allowed_cidrs = [\"nonsense\"]\n" + redisStatic, "allowed_cidrs"},
{"notify scheme", minimal + "[notify]\norg_url = \"ftp://x/limits\"\n", "org_url"},
{"notify no scheme", minimal + "[notify]\naddress_url = \"not a url\"\n", "address_url"},
{"notify no host", minimal + "[notify]\ndomain_url = \"https:///path\"\n", "domain_url"},
{"logging format", minimal + "[logging]\nformat = \"xml\"\n", "logging.format"},
{"logging level", minimal + "[logging]\nlevel = \"loud\"\n", "logging.level"},
{"counting format", minimal + "[logging]\ncounting_format = \"csv\"\n", "counting_format"},
{"pool size", top + redisStatic + "pool_size = -1\n", "pool_size"},
{"timeout_ms", top + redisStatic + "timeout_ms = 0\n", "timeout_ms"},
{"handler_timeout", top + "handler_timeout = 0\n" + redisStatic, "handler_timeout"},
{"addr", top + "addr = \"no-port\"\n" + redisStatic, "addr"},
{"admin_addr", top + "admin_addr = \"no-port\"\n" + redisStatic, "admin_addr"},
{"notify timeout", minimal + "[notify]\ntimeout_seconds = 0\n", "timeout_seconds"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
_, err := Load(writeConfig(t, tt.toml))
if err == nil {
t.Fatalf("expected error mentioning %q", tt.wantErr)
}
if !strings.Contains(err.Error(), tt.wantErr) {
t.Errorf("error %q does not mention %q", err, tt.wantErr)
}
})
}
}
func TestLogging_CountingOutput(t *testing.T) {
tests := []struct {
toml string
wantOut string
wantEnabled bool
}{
{"", "stderr", true},
{"[logging]\noutput = \"/var/log/a.log\"\n", "/var/log/a.log", true}, // follows output
{"[logging]\noutput = \"/var/log/a.log\"\ncounting_output = \"/var/log/c.log\"\n", "/var/log/c.log", true},
{"[logging]\ncounting_output = \"off\"\n", "off", false},
{"[logging]\ncounting_output = \"OFF\"\n", "off", false},
}
for _, tt := range tests {
cfg := mustLoad(t, minimal+tt.toml)
if cfg.Logging.CountingOutput != tt.wantOut || cfg.Logging.CountingEnabled() != tt.wantEnabled {
t.Errorf("%q: counting_output=%q enabled=%v, want %q %v", tt.toml,
cfg.Logging.CountingOutput, cfg.Logging.CountingEnabled(), tt.wantOut, tt.wantEnabled)
}
}
}
func TestWarnings_MissingPlaceholders(t *testing.T) {
cfg := mustLoad(t, top+redisStatic+`
[notify]
org_url = "https://admin.example/api/limits?level=$LEVEL&id=$ID&flow=$FLOW&period=$PERIOD&by=$ADDRESS"
address_url = "https://admin.example/api/limits?level=$LEVEL&id=$ID&flow=$FLOW&period=$PERIOD&by=$ADDRESS&count=$COUNT"
`)
w := cfg.Warnings()
if len(w) != 1 || !strings.Contains(w[0], "org_url") || !strings.Contains(w[0], "$COUNT") {
t.Errorf("Warnings() = %v, want exactly one about org_url missing $COUNT", w)
}
}
// The env var satisfies the mandatory token like the file does; a typo in the
// key never starts the daemon unauthenticated.
func TestLoad_TokenFromEnvOnly(t *testing.T) {
t.Setenv("MIGADU_COUNTERS_AUTH_TOKEN", "env-token")
cfg := mustLoad(t, "domains_cdb = \"/tmp/x\"\n"+redisStatic)
if cfg.AuthToken != "env-token" {
t.Errorf("auth_token = %q", cfg.AuthToken)
}
if _, err := Load(writeConfig(t, "domains_cdb = \"/tmp/x\"\nauth-token = \"x\"\n"+redisStatic)); err == nil {
t.Error("typo key must be fatal even when the env var supplies a token")
}
}
func TestLoad_QueueAndCache(t *testing.T) {
cfg := mustLoad(t, top+"queue_size = 50\nworkers = 2\nmax_cached_domains = 0\n"+redisStatic)
if cfg.QueueSize != 50 || cfg.Workers != 2 || cfg.MaxCachedDomains != 0 {
t.Errorf("queue/cache = %d %d %d", cfg.QueueSize, cfg.Workers, cfg.MaxCachedDomains)
}
}
func TestDiscover(t *testing.T) {
if got, err := Discover("/explicit/path.toml"); err != nil || got != "/explicit/path.toml" {
t.Errorf("Discover(explicit) = %q, %v", got, err)
}
// With no default file present on the test machine, Discover must explain where it looked.
paths := DefaultPaths
DefaultPaths = []string{filepath.Join(t.TempDir(), "missing.toml")}
defer func() { DefaultPaths = paths }()
if _, err := Discover(""); err == nil || !strings.Contains(err.Error(), "missing.toml") {
t.Errorf("Discover(\"\") err = %v, want one listing the searched paths", err)
}
found := writeConfig(t, minimal)
DefaultPaths = []string{DefaultPaths[0], found}
if got, err := Discover(""); err != nil || got != found {
t.Errorf("Discover(\"\") = %q, %v; want %q", got, err, found)
}
}
internal/config/logger.go
+116
View File
@@ -0,0 +1,116 @@
package config
import (
"fmt"
"log/slog"
"os"
"strings"
"sync"
)
// LogWriter wraps a log destination and supports reopening for log rotation.
// When the output is a file, Reopen closes the old handle and opens a new one
// so that newsyslog rotation works after SIGHUP. Shared by the application log
// and the counting transaction log. (Same implementation as rcptd.)
type LogWriter struct {
mu sync.Mutex
path string // empty for stdout/stderr
file *os.File // current output file (or os.Stdout/os.Stderr)
}
func (w *LogWriter) Write(p []byte) (n int, err error) {
w.mu.Lock()
defer w.mu.Unlock()
return w.file.Write(p)
}
// Reopen closes and reopens the log file. No-op for stdout/stderr.
func (w *LogWriter) Reopen() error {
if w.path == "" {
return nil
}
w.mu.Lock()
defer w.mu.Unlock()
newFile, err := os.OpenFile(w.path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
if err != nil {
return fmt.Errorf("reopen log file %s: %w", w.path, err)
}
w.file.Close()
w.file = newFile
return nil
}
// Close closes the log file. No-op for stdout/stderr.
func (w *LogWriter) Close() error {
if w.path == "" {
return nil
}
w.mu.Lock()
defer w.mu.Unlock()
return w.file.Close()
}
// NewWriter opens a log destination: "stderr", "stdout" or a file path.
func NewWriter(output string) (*LogWriter, error) { return newLogWriter(output) }
// NewLogger creates the application slog.Logger from cfg and installs it as
// the slog default. The returned LogWriter must be closed by the caller.
func NewLogger(cfg LoggingConfig) (*slog.Logger, *LogWriter, error) {
level, err := parseLogLevel(cfg.Level)
if err != nil {
return nil, nil, fmt.Errorf("invalid log level: %w", err)
}
writer, err := newLogWriter(cfg.Output)
if err != nil {
return nil, nil, fmt.Errorf("failed to create log writer: %w", err)
}
var handler slog.Handler
opts := &slog.HandlerOptions{Level: level}
switch strings.ToLower(cfg.Format) {
case "json":
handler = slog.NewJSONHandler(writer, opts)
case "console", "":
handler = slog.NewTextHandler(writer, opts)
default:
writer.Close()
return nil, nil, fmt.Errorf("invalid log format: %s (must be 'console' or 'json')", cfg.Format)
}
logger := slog.New(handler)
slog.SetDefault(logger)
return logger, writer, nil
}
func parseLogLevel(level string) (slog.Level, error) {
switch strings.ToLower(level) {
case "debug":
return slog.LevelDebug, nil
case "info", "":
return slog.LevelInfo, nil
case "warn", "warning":
return slog.LevelWarn, nil
case "error":
return slog.LevelError, nil
default:
return slog.LevelInfo, fmt.Errorf("unknown level: %s (must be debug, info, warn, or error)", level)
}
}
func newLogWriter(output string) (*LogWriter, error) {
switch strings.ToLower(output) {
case "stderr", "":
return &LogWriter{file: os.Stderr}, nil
case "stdout":
return &LogWriter{file: os.Stdout}, nil
default:
file, err := os.OpenFile(output, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
if err != nil {
return nil, fmt.Errorf("failed to open log file %s: %w", output, err)
}
return &LogWriter{path: output, file: file}, nil
}
}
internal/config/logger_test.go
+73
View File
@@ -0,0 +1,73 @@
package config
import (
"log/slog"
"os"
"path/filepath"
"strings"
"testing"
)
func TestNewLogger_Invalid(t *testing.T) {
if _, _, err := NewLogger(LoggingConfig{Format: "xml"}); err == nil {
t.Error("expected error for format xml")
}
if _, _, err := NewLogger(LoggingConfig{Level: "loud"}); err == nil {
t.Error("expected error for level loud")
}
}
func TestNewLogger_Levels(t *testing.T) {
for _, tt := range []struct {
in string
want slog.Level
}{{"", slog.LevelInfo}, {"debug", slog.LevelDebug}, {"INFO", slog.LevelInfo}, {"warn", slog.LevelWarn}, {"warning", slog.LevelWarn}, {"error", slog.LevelError}} {
got, err := parseLogLevel(tt.in)
if err != nil || got != tt.want {
t.Errorf("parseLogLevel(%q) = %v, %v; want %v", tt.in, got, err, tt.want)
}
}
}
// Rotation: after the file is renamed away, Reopen must create a fresh file at
// the configured path and subsequent writes must land there.
func TestLogWriter_Reopen(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "app.log")
w, err := newLogWriter(path)
if err != nil {
t.Fatal(err)
}
defer w.Close()
if _, err := w.Write([]byte("one\n")); err != nil {
t.Fatal(err)
}
rotated := filepath.Join(dir, "app.log.0")
if err := os.Rename(path, rotated); err != nil {
t.Fatal(err)
}
if err := w.Reopen(); err != nil {
t.Fatalf("Reopen: %v", err)
}
if _, err := w.Write([]byte("two\n")); err != nil {
t.Fatal(err)
}
old, _ := os.ReadFile(rotated)
fresh, _ := os.ReadFile(path)
if strings.TrimSpace(string(old)) != "one" || strings.TrimSpace(string(fresh)) != "two" {
t.Errorf("rotated=%q fresh=%q", old, fresh)
}
}
func TestLogWriter_StderrReopenIsNoop(t *testing.T) {
w, err := newLogWriter("stderr")
if err != nil {
t.Fatal(err)
}
if err := w.Reopen(); err != nil {
t.Errorf("Reopen on stderr: %v", err)
}
if err := w.Close(); err != nil {
t.Errorf("Close on stderr: %v", err)
}
}
internal/counting/counting.go
+345
View File
@@ -0,0 +1,345 @@
// 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
}
// 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
LoginSuffix string // /out: "SUPPORTTOKEN" of a user@dom@SUPPORTTOKEN 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
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
LoginSuffix string
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)
}
plan := &Plan{Route: job.Route, TraceID: job.TraceID, Date: store.DayKey(p.now()),
AuthUser: authUser, LoginSuffix: job.LoginSuffix, 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,
})
}
return nil
}
// 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
Err error // this key's own INCR failed (e.g. it holds a non-number); Count is 0
}
// 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)
}
// 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. 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;
// crossing detection then uses count >= limit so a crossing that happened in
// the lost transaction is still notified. A definite failure (connection
// refused, cancelled) is returned as an error: nothing was counted.
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.Store.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)
}
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 {
out.Results[i] = Result{Increment: inc, Count: counts[i], Reached: inc.Limited && counts[i] >= inc.Limit}
}
return out, nil
}
func (c *Counter) reReadTimeout() time.Duration {
if c.ReReadTimeout > 0 {
return c.ReReadTimeout
}
return 2 * time.Second
}
internal/counting/counting_test.go
+677
View File
@@ -0,0 +1,677 @@
package counting
import (
"bytes"
"context"
"errors"
"fmt"
"net"
"sync/atomic"
"testing"
"time"
"github.com/alicebob/miniredis/v2"
"github.com/redis/go-redis/v9"
"migadu_counters/internal/domains"
"migadu_counters/internal/store"
)
// table is a fake domains lookup: domain -> record. Unknown domains return
// domains.ErrNotFound; the special domain "broken.tld" returns a read error.
type table map[string]*domains.Domain
func (t table) Lookup(domain string) (*domains.Domain, error) {
if domain == "broken.tld" {
return nil, errors.New("cdb read error")
}
if d, ok := t[domain]; ok {
return d, nil
}
return nil, domains.ErrNotFound
}
var (
fixedNow = time.Date(2026, 10, 13, 12, 0, 0, 0, time.UTC)
day = "2026-10-13"
exempt = map[string]bool{"support@migadu.com": true, "contact@migadu.com": true}
)
func newTable() table {
return table{
"example.com": {Name: "example.com", OrgID: 7,
Limits: domains.ParseLimits("O/i=10000,O/o=2200,D/i=10000,D/o=2200,A/i=500,A/o=100,noreply@/o=5")},
"other.org": {Name: "other.org", OrgID: 9,
Limits: domains.ParseLimits("O/i=200,O/o=20,D/i=200,D/o=20,A/i=200,A/o=20")},
"nolimits.net": {Name: "nolimits.net", OrgID: 11, Limits: domains.Limits{}},
}
}
func newPlanner(countJunk bool) *Planner {
return &Planner{Lookup: newTable(), Exempt: exempt, CountJunk: countJunk, Now: func() time.Time { return fixedNow }}
}
type want struct {
level, id, flow, by string
limit int64
limited bool
}
func checkIncrements(t *testing.T, got []Increment, wants []want) {
t.Helper()
if len(got) != len(wants) {
t.Fatalf("got %d increments, want %d:\n%s", len(got), len(wants), dump(got))
}
for i, w := range wants {
g := got[i]
wantKey := store.Key(w.level, w.id, day, w.flow)
if g.Level != w.level || g.ID != w.id || g.Flow != w.flow || g.By != w.by || g.Limit != w.limit || g.Limited != w.limited || g.Key != wantKey {
t.Errorf("increment %d = %+v,\n want level=%s id=%s flow=%s by=%s limit=%d limited=%v key=%s",
i, g, w.level, w.id, w.flow, w.by, w.limit, w.limited, wantKey)
}
}
}
func dump(incs []Increment) string {
s := ""
for i, in := range incs {
s += fmt.Sprintf(" %d: %+v\n", i, in)
}
return s
}
// x@example.com -> y@example.com: sender outgoing on A/D/O, recipient incoming on A/D/O.
func TestPlanOut_LocalRecipient(t *testing.T) {
p := newPlanner(false)
plan, err := p.PlanOut("tr1", "x@example.com", "x@example.com", "y@example.com")
if err != nil {
t.Fatal(err)
}
if plan.Skipped != "" || plan.Route != "out" || plan.TraceID != "tr1" || plan.Date != day {
t.Errorf("plan header = %+v", plan)
}
checkIncrements(t, plan.Increments, []want{
{"A", "x@example.com", "o", "x@example.com", 100, true},
{"D", "example.com", "o", "x@example.com", 2200, true},
{"O", "7", "o", "x@example.com", 2200, true},
{"A", "y@example.com", "i", "y@example.com", 500, true},
{"D", "example.com", "i", "y@example.com", 10000, true},
{"O", "7", "i", "y@example.com", 10000, true},
})
if len(plan.Lookups) != 2 || plan.Lookups[0].Role != "sender" || plan.Lookups[0].Domain == nil ||
plan.Lookups[1].Role != "recipient" || plan.Lookups[1].Domain == nil {
t.Errorf("lookups = %+v", plan.Lookups)
}
}
// x@example.com -> s@gmail.com: only the sender side; the recipient lookup is recorded as not found.
func TestPlanOut_ExternalRecipient(t *testing.T) {
plan, err := newPlanner(false).PlanOut("tr", "x@example.com", "x@example.com", "s@gmail.com")
if err != nil {
t.Fatal(err)
}
checkIncrements(t, plan.Increments, []want{
{"A", "x@example.com", "o", "x@example.com", 100, true},
{"D", "example.com", "o", "x@example.com", 2200, true},
{"O", "7", "o", "x@example.com", 2200, true},
})
if len(plan.Lookups) != 2 || plan.Lookups[1].Domain != nil || !errors.Is(plan.Lookups[1].Err, domains.ErrNotFound) {
t.Errorf("recipient lookup should be recorded as not found: %+v", plan.Lookups)
}
}
// Recipient in another hosted org: its own domain row decides D/O ids and limits.
func TestPlanOut_RecipientInOtherOrg(t *testing.T) {
plan, err := newPlanner(false).PlanOut("tr", "x@example.com", "", "z@other.org")
if err != nil {
t.Fatal(err)
}
checkIncrements(t, plan.Increments[3:], []want{
{"A", "z@other.org", "i", "z@other.org", 200, true},
{"D", "other.org", "i", "z@other.org", 200, true},
{"O", "9", "i", "z@other.org", 200, true},
})
}
func TestPlanOut_SenderDomainUnknown(t *testing.T) {
plan, err := newPlanner(false).PlanOut("tr", "staff@unknown.tld", "", "y@example.com")
if err != nil {
t.Fatal(err)
}
if plan.Skipped != "" {
t.Errorf("skipped = %q, want counting of the local recipient", plan.Skipped)
}
checkIncrements(t, plan.Increments, []want{
{"A", "y@example.com", "i", "y@example.com", 500, true},
{"D", "example.com", "i", "y@example.com", 10000, true},
{"O", "7", "i", "y@example.com", 10000, true},
})
}
func TestPlanOut_BothUnknown(t *testing.T) {
plan, err := newPlanner(false).PlanOut("tr", "a@unknown.tld", "", "b@gmail.com")
if err != nil {
t.Fatal(err)
}
if plan.Skipped != SkipUnknownDomain || len(plan.Increments) != 0 {
t.Errorf("plan = %+v", plan)
}
}
func TestPlanOut_Exempt(t *testing.T) {
cases := []struct{ authUser, mailFrom, recipient string }{
{"support@migadu.com", "support@migadu.com", "y@example.com"},
{"x@example.com", "x@example.com", "support@migadu.com"},
{"x@example.com", "contact@migadu.com", "y@example.com"},
{"x@example.com", "x@example.com", "Contact@Migadu.com"},
}
for _, c := range cases {
plan, err := newPlanner(false).PlanOut("tr", c.authUser, c.mailFrom, c.recipient)
if err != nil {
t.Fatal(err)
}
if plan.Skipped != SkipExempt || len(plan.Increments) != 0 || len(plan.Lookups) != 0 {
t.Errorf("%+v: plan = %+v, want exempt with no lookups", c, plan)
}
}
}
func TestPlanOut_AddressOverrideAndUnlimited(t *testing.T) {
plan, err := newPlanner(false).PlanOut("tr", "noreply@example.com", "", "q@nolimits.net")
if err != nil {
t.Fatal(err)
}
checkIncrements(t, plan.Increments, []want{
{"A", "noreply@example.com", "o", "noreply@example.com", 5, true}, // noreply@/o override
{"D", "example.com", "o", "noreply@example.com", 2200, true},
{"O", "7", "o", "noreply@example.com", 2200, true},
{"A", "q@nolimits.net", "i", "q@nolimits.net", 0, false},
{"D", "nolimits.net", "i", "q@nolimits.net", 0, false},
{"O", "11", "i", "q@nolimits.net", 0, false},
})
}
func TestPlanOut_NormalizesAddresses(t *testing.T) {
plan, err := newPlanner(false).PlanOut("tr", " X@Example.COM ", "", "Y@EXAMPLE.com")
if err != nil {
t.Fatal(err)
}
if plan.Increments[0].ID != "x@example.com" || plan.Increments[3].ID != "y@example.com" {
t.Errorf("ids not normalised: %s", dump(plan.Increments))
}
}
func TestPlanIn_LocalRecipient(t *testing.T) {
plan, err := newPlanner(false).PlanIn("tr9", "s@gmail.com", "y@example.com", false)
if err != nil {
t.Fatal(err)
}
if plan.Route != "in" || plan.TraceID != "tr9" || plan.Skipped != "" {
t.Errorf("plan header = %+v", plan)
}
checkIncrements(t, plan.Increments, []want{
{"A", "y@example.com", "i", "y@example.com", 500, true},
{"D", "example.com", "i", "y@example.com", 10000, true},
{"O", "7", "i", "y@example.com", 10000, true},
})
if len(plan.Lookups) != 1 || plan.Lookups[0].Role != "recipient" {
t.Errorf("lookups = %+v", plan.Lookups)
}
}
func TestPlanIn_SenderNeverCounted(t *testing.T) {
// A hosted sender on /in (e.g. mail relayed back in) still only counts the recipient.
plan, err := newPlanner(false).PlanIn("tr", "x@example.com", "z@other.org", false)
if err != nil {
t.Fatal(err)
}
for _, in := range plan.Increments {
if in.Flow != "i" || in.By != "z@other.org" {
t.Errorf("unexpected increment %+v", in)
}
}
}
func TestPlanIn_UnknownDomain(t *testing.T) {
plan, err := newPlanner(false).PlanIn("tr", "s@gmail.com", "nobody@gmail.com", false)
if err != nil {
t.Fatal(err)
}
if plan.Skipped != SkipUnknownDomain || len(plan.Increments) != 0 {
t.Errorf("plan = %+v", plan)
}
}
func TestPlanIn_Junk(t *testing.T) {
plan, err := newPlanner(false).PlanIn("tr", "s@gmail.com", "y@example.com", true)
if err != nil {
t.Fatal(err)
}
if plan.Skipped != SkipJunk || len(plan.Increments) != 0 || len(plan.Lookups) != 0 {
t.Errorf("junk not skipped: %+v", plan)
}
plan, err = newPlanner(true).PlanIn("tr", "s@gmail.com", "y@example.com", true)
if err != nil {
t.Fatal(err)
}
if plan.Skipped != "" || len(plan.Increments) != 3 {
t.Errorf("count_junk=true should count: %+v", plan)
}
}
func TestPlanIn_Exempt(t *testing.T) {
for _, c := range []struct{ from, rcpt string }{{"support@migadu.com", "y@example.com"}, {"s@gmail.com", "contact@migadu.com"}} {
plan, err := newPlanner(false).PlanIn("tr", c.from, c.rcpt, false)
if err != nil {
t.Fatal(err)
}
if plan.Skipped != SkipExempt || len(plan.Increments) != 0 {
t.Errorf("%+v: plan = %+v", c, plan)
}
}
}
func TestPlanIn_EmptyMailFromIsBounce(t *testing.T) {
plan, err := newPlanner(false).PlanIn("tr", "", "y@example.com", false)
if err != nil || plan.Skipped != "" || len(plan.Increments) != 3 {
t.Errorf("bounce (<>) must still count the recipient: %+v, %v", plan, err)
}
}
// A CDB read error (not "not found") must surface so the server answers 503.
func TestPlan_LookupError(t *testing.T) {
if _, err := newPlanner(false).PlanIn("tr", "s@gmail.com", "y@broken.tld", false); err == nil {
t.Error("PlanIn: expected error")
}
if _, err := newPlanner(false).PlanOut("tr", "x@broken.tld", "", "y@example.com"); err == nil {
t.Error("PlanOut: expected error")
}
}
// Against the real CDB: yagolivas.com is org 1049 with A/o=20, D/o=20, O/o=20.
func TestPlanOut_RealCDB(t *testing.T) {
db, err := domains.Open("../../test_files/rspamd_domains_db.cdb")
if err != nil {
t.Fatal(err)
}
defer db.Close()
p := &Planner{Lookup: db, Exempt: exempt, Now: func() time.Time { return fixedNow }}
plan, err := p.PlanOut("tr", "someone@yagolivas.com", "", "ext@gmail.com")
if err != nil {
t.Fatal(err)
}
checkIncrements(t, plan.Increments, []want{
{"A", "someone@yagolivas.com", "o", "someone@yagolivas.com", 20, true},
{"D", "yagolivas.com", "o", "someone@yagolivas.com", 20, true},
{"O", "1049", "o", "someone@yagolivas.com", 20, true},
})
}
func TestCrossed(t *testing.T) {
tests := []struct {
newValue, amount, limit int64
want bool
}{
{1, 1, 1, true},
{5, 1, 5, true},
{4, 1, 5, false},
{6, 1, 5, false}, // already over: no second notification
{5, 3, 5, true}, // 2 -> 5 crosses
{7, 3, 5, true}, // 4 -> 7 crosses
{8, 3, 5, false}, // 5 -> 8 was already at the limit
{1, 1, 0, false}, // 0 / unlimited never crosses
{1, 1, -1, false},
}
for _, tt := range tests {
if got := Crossed(tt.newValue, tt.amount, tt.limit); got != tt.want {
t.Errorf("Crossed(%d, %d, %d) = %v, want %v", tt.newValue, tt.amount, tt.limit, got, tt.want)
}
}
}
func newCounter(t *testing.T) (*Counter, *miniredis.Miniredis) {
t.Helper()
m := miniredis.RunT(t)
st := store.NewWithClients(redis.NewClient(&redis.Options{Addr: m.Addr()}), nil)
t.Cleanup(func() { st.Close() })
return &Counter{Planner: newPlanner(false), Store: st}, m
}
// Executes the plan: counts come back from Redis and the crossing increment
// is flagged exactly once.
func TestCounter_Out(t *testing.T) {
c, m := newCounter(t)
ctx := context.Background()
// noreply@example.com has A/o=5: the fifth message reaches it.
var last *Outcome
for i := 1; i <= 6; i++ {
out, err := c.CountOut(ctx, "tr", "noreply@example.com", "", "ext@gmail.com")
if err != nil {
t.Fatal(err)
}
if len(out.Results) != 3 || out.Results[0].Count != int64(i) {
t.Fatalf("message %d: results = %+v", i, out.Results)
}
if reached := out.Results[0].Reached; reached != (i == 5) {
t.Errorf("message %d: A reached = %v", i, reached)
}
last = out
}
if v, _ := m.Get(store.Key("A", "noreply@example.com", day, "o")); v != "6" {
t.Errorf("redis A/o = %q, want 6", v)
}
if v, _ := m.Get(store.Key("O", "7", day, "o")); v != "6" {
t.Errorf("redis O/o = %q, want 6", v)
}
if last.Skipped != "" || last.Route != "out" {
t.Errorf("outcome = %+v", last)
}
}
// After the web app resets a counter (DEL), the next crossing fires again.
func TestCounter_ReachedAgainAfterReset(t *testing.T) {
c, m := newCounter(t)
ctx := context.Background()
for i := 0; i < 5; i++ {
if _, err := c.CountOut(ctx, "tr", "noreply@example.com", "", "ext@gmail.com"); err != nil {
t.Fatal(err)
}
}
m.Del(store.Key("A", "noreply@example.com", day, "o"))
var reachedAt int
for i := 1; i <= 5; i++ {
out, err := c.CountOut(ctx, "tr", "noreply@example.com", "", "ext@gmail.com")
if err != nil {
t.Fatal(err)
}
if out.Results[0].Reached {
reachedAt = i
}
}
if reachedAt != 5 {
t.Errorf("after reset the limit was reached at message %d, want 5", reachedAt)
}
}
func TestCounter_In(t *testing.T) {
c, m := newCounter(t)
out, err := c.CountIn(context.Background(), "tr", "s@gmail.com", "y@example.com", false)
if err != nil {
t.Fatal(err)
}
if len(out.Results) != 3 || out.Results[2].Count != 1 || out.Results[2].Level != "O" {
t.Errorf("results = %+v", out.Results)
}
if v, _ := m.Get(store.Key("D", "example.com", day, "i")); v != "1" {
t.Errorf("redis D/i = %q, want 1", v)
}
}
// A skipped plan must not touch Redis at all.
func TestCounter_SkippedDoesNotTouchRedis(t *testing.T) {
c, m := newCounter(t)
out, err := c.CountIn(context.Background(), "tr", "support@migadu.com", "y@example.com", false)
if err != nil || out.Skipped != SkipExempt || len(out.Results) != 0 {
t.Fatalf("outcome = %+v, %v", out, err)
}
if keys := m.Keys(); len(keys) != 0 {
t.Errorf("redis keys = %v, want none", keys)
}
}
func TestCounter_RedisDown(t *testing.T) {
c, m := newCounter(t)
m.Close()
if _, err := c.CountIn(context.Background(), "tr", "s@gmail.com", "y@example.com", false); err == nil {
t.Error("expected error when redis is down")
}
}
// --- Jobs with unusable sides (review #2, #3) --------------------------------
func TestPlan_OutSenderUnusableCountsRecipient(t *testing.T) {
job := Job{Route: RouteOut, TraceID: "u1", Recipient: "y@example.com",
SenderUnusable: Unusable{Role: "sender", Value: "nobody", Reason: "invalid e-mail address"}}
plan, err := newPlanner(false).Plan(job)
if err != nil {
t.Fatal(err)
}
checkIncrements(t, plan.Increments, []want{
{"A", "y@example.com", "i", "y@example.com", 500, true},
{"D", "example.com", "i", "y@example.com", 10000, true},
{"O", "7", "i", "y@example.com", 10000, true},
})
if len(plan.Unusable) != 1 || plan.Unusable[0].Role != "sender" || plan.Unusable[0].Value != "nobody" {
t.Errorf("unusable = %+v", plan.Unusable)
}
if len(plan.Lookups) != 1 || plan.Lookups[0].Role != "recipient" {
t.Errorf("only the recipient should be looked up: %+v", plan.Lookups)
}
}
func TestPlan_OutRecipientUnusableCountsSender(t *testing.T) {
job := Job{Route: RouteOut, TraceID: "u2", AuthUser: "x@example.com", LoginSuffix: "SUPPORTTOKEN", MailFrom: "x@example.com",
RecipientUnusable: Unusable{Role: "recipient", Value: "u@[192.0.2.1]", Reason: "invalid e-mail address"}}
plan, err := newPlanner(false).Plan(job)
if err != nil {
t.Fatal(err)
}
checkIncrements(t, plan.Increments, []want{
{"A", "x@example.com", "o", "x@example.com", 100, true},
{"D", "example.com", "o", "x@example.com", 2200, true},
{"O", "7", "o", "x@example.com", 2200, true},
})
if plan.LoginSuffix != "SUPPORTTOKEN" || plan.AuthUser != "x@example.com" {
t.Errorf("login suffix not carried: %+v", plan)
}
if len(plan.Unusable) != 1 || plan.Unusable[0].Role != "recipient" {
t.Errorf("unusable = %+v", plan.Unusable)
}
}
// An unusable mail_from on /in is logged, treated as non-exempt, and the
// recipient is counted.
func TestPlan_InMailFromUnusable(t *testing.T) {
job := Job{Route: RouteIn, TraceID: "u3", MailFrom: "john smith@x.com", Recipient: "y@example.com"}
plan, err := newPlanner(false).Plan(job)
if err != nil {
t.Fatal(err)
}
if plan.Skipped != "" || len(plan.Increments) != 3 || plan.MailFrom != "john smith@x.com" {
t.Errorf("plan = %+v", plan)
}
}
// The exempt check still applies to usable sides of a Job.
func TestPlan_ExemptViaJob(t *testing.T) {
plan, err := newPlanner(false).Plan(Job{Route: RouteIn, MailFrom: "support@migadu.com", MailFromUsable: true, Recipient: "y@example.com"})
if err != nil || plan.Skipped != SkipExempt {
t.Errorf("plan = %+v, %v", plan, err)
}
}
// --- Execute: partial EXEC failure and lost replies (review #4) --------------
// One key holding garbage fails its own INCR; the others are counted and the
// crossing among them is still detected.
func TestCounter_PartialExecFailure(t *testing.T) {
c, m := newCounter(t)
m.Set(store.Key("O", "7", day, "o"), "garbage")
m.Set(store.Key("A", "noreply@example.com", day, "o"), "4") // A/o=5: this increment crosses
out, err := c.CountOut(context.Background(), "p1", "noreply@example.com", "", "ext@gmail.com")
if err != nil {
t.Fatalf("a per-key failure must not fail the request: %v", err)
}
if len(out.Results) != 3 {
t.Fatalf("results = %+v", out.Results)
}
a, d, o := out.Results[0], out.Results[1], out.Results[2]
if a.Err != nil || a.Count != 5 || !a.Reached {
t.Errorf("A = %+v, want count 5 reached", a)
}
if d.Err != nil || d.Count != 1 {
t.Errorf("D = %+v", d)
}
if o.Err == nil || o.Reached {
t.Errorf("O = %+v, want its own error and no crossing", o)
}
if out.Unknown != nil {
t.Errorf("outcome is known: %v", out.Unknown)
}
}
// delayProxy forwards TCP to backend and, once it has seen EXEC in the
// request stream, holds back the next `delayed` replies by `delay`, so the
// client's read timeout fires after Redis has applied the transaction (the
// connection handshake replies pass undelayed).
func delayProxy(t *testing.T, backend string, delayed int, delay time.Duration) string {
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { ln.Close() })
var remaining atomic.Int64
remaining.Store(int64(delayed))
go func() {
for {
c, err := ln.Accept()
if err != nil {
return
}
b, err := net.Dial("tcp", backend)
if err != nil {
c.Close()
continue
}
var armed atomic.Bool
go func() {
defer b.Close()
buf := make([]byte, 32*1024)
for {
n, err := c.Read(buf)
if n > 0 {
if bytes.Contains(bytes.ToLower(buf[:n]), []byte("exec")) {
armed.Store(true)
}
if _, werr := b.Write(buf[:n]); werr != nil {
return
}
}
if err != nil {
return
}
}
}()
go func() {
defer c.Close()
buf := make([]byte, 32*1024)
for {
n, err := b.Read(buf)
if n > 0 {
if armed.Load() && remaining.Add(-1) >= 0 {
time.Sleep(delay)
}
if _, werr := c.Write(buf[:n]); werr != nil {
return
}
}
if err != nil {
return
}
}
}()
}
}()
return ln.Addr().String()
}
// The EXEC reply is lost after Redis applied the increments: the counter is
// re-read from the master and the crossing is still reported.
func TestCounter_LostReplyReReadsMaster(t *testing.T) {
m := miniredis.RunT(t)
proxy := delayProxy(t, m.Addr(), 1, 800*time.Millisecond)
st := store.NewWithClients(redis.NewClient(&redis.Options{Addr: proxy, ReadTimeout: 200 * time.Millisecond, PoolSize: 2}), nil)
t.Cleanup(func() { st.Close() })
c := &Counter{Planner: newPlanner(false), Store: st}
m.Set(store.Key("A", "noreply@example.com", day, "o"), "4") // A/o=5
out, err := c.CountOut(context.Background(), "lost", "noreply@example.com", "", "ext@gmail.com")
if err != nil {
t.Fatalf("lost reply must be resolved by a re-read, got error: %v", err)
}
if out.Unknown == nil {
t.Fatal("outcome should be flagged unknown (reply lost)")
}
if v, _ := m.Get(store.Key("A", "noreply@example.com", day, "o")); v != "5" {
t.Fatalf("redis A/o = %q, the increment should have been applied", v)
}
if a := out.Results[0]; a.Count != 5 || !a.Reached {
t.Errorf("re-read A = %+v, want count 5 reached", a)
}
if o := out.Results[2]; o.Count != 1 || o.Reached {
t.Errorf("re-read O = %+v", o)
}
}
// A definite failure (connection refused) is still an error: nothing counted.
func TestCounter_RefusedIsNotReRead(t *testing.T) {
c, m := newCounter(t)
m.Close()
out, err := c.CountIn(context.Background(), "tr", "s@gmail.com", "y@example.com", false)
if err == nil {
t.Errorf("expected error, got outcome %+v", out)
}
}
// --- review #7: plus-addressing counts against the mailbox -------------------
func TestPlan_PlusTagCountsMailbox(t *testing.T) {
p := newPlanner(false)
plan, err := p.PlanOut("pt", "noreply+news@example.com", "", "y+1@example.com")
if err != nil {
t.Fatal(err)
}
checkIncrements(t, plan.Increments, []want{
{"A", "noreply@example.com", "o", "noreply@example.com", 5, true}, // override noreply@/o still applies
{"D", "example.com", "o", "noreply@example.com", 2200, true},
{"O", "7", "o", "noreply@example.com", 2200, true},
{"A", "y@example.com", "i", "y@example.com", 500, true},
{"D", "example.com", "i", "y@example.com", 10000, true},
{"O", "7", "i", "y@example.com", 10000, true},
})
if plan.Recipient != "y+1@example.com" || plan.AuthUser != "noreply+news@example.com" {
t.Errorf("the request line keeps the addresses as sent: %+v", plan)
}
// Exempt applies to the mailbox too.
plan, _ = p.PlanIn("pt", "s@gmail.com", "support+ticket@migadu.com", false)
if plan.Skipped != SkipExempt {
t.Errorf("support+ticket@ should be exempt: %+v", plan)
}
}
// --- review #8: a malformed CDB row is "not hosted", not a failure -----------
func TestPlan_MalformedRowIsNotHosted(t *testing.T) {
db, err := domains.Open("../../review/_proof_tests/b/bad.cdb")
if err != nil {
t.Fatal(err)
}
defer db.Close()
p := &Planner{Lookup: db, Now: func() time.Time { return fixedNow }}
plan, err := p.PlanOut("m", "a@good.com", "", "b@bad.com")
if err != nil {
t.Fatalf("malformed recipient row must not fail the request: %v", err)
}
checkIncrements(t, plan.Increments, []want{
{"A", "a@good.com", "o", "a@good.com", 5, true},
{"D", "good.com", "o", "a@good.com", 0, false},
{"O", "7", "o", "a@good.com", 0, false},
})
if len(plan.Lookups) != 2 || plan.Lookups[1].Domain != nil || !errors.Is(plan.Lookups[1].Err, domains.ErrMalformed) {
t.Errorf("recipient lookup should record the malformed row: %+v", plan.Lookups)
}
plan, err = p.PlanIn("m", "s@gmail.com", "b@bad.com", false)
if err != nil || plan.Skipped != SkipUnknownDomain || len(plan.Increments) != 0 {
t.Errorf("/in to a malformed row: plan=%+v err=%v, want skipped unknown_domain", plan, err)
}
}
internal/counting/queue.go
+138
View File
@@ -0,0 +1,138 @@
package counting
import (
"context"
"fmt"
"runtime/debug"
"sync"
"sync/atomic"
"time"
)
// QueueConfig sizes the counting queue.
type QueueConfig struct {
Size int // queued jobs before Enqueue refuses (default 10 000)
Workers int // concurrent workers (default 8)
Timeout time.Duration // budget per job: plan, Redis, log (default 4s)
}
// Observer receives queue events for metrics and logging; nil funcs are ignored.
type Observer struct {
Depth func(n int) // queue depth after every enqueue/dequeue
Wait func(d time.Duration) // time a job waited in the queue
Panic func(job Job, p any, stack []byte) // a handler panicked (recovered)
}
// Queue decouples counting from the HTTP request: /in and /out validate,
// enqueue a Job and answer 202; workers run the handler with their own
// bounded context, so a client that hangs up or times out can no longer
// cancel a count (review #1).
type Queue struct {
cfg QueueConfig
jobs chan Job
handler func(ctx context.Context, job Job)
obs Observer
pending atomic.Int64 // enqueued and not yet finished
mu sync.RWMutex // Enqueue vs Drain (sending on a closed channel panics)
closed bool
wg sync.WaitGroup
}
// NewQueue applies defaults and starts the workers.
func NewQueue(cfg QueueConfig, handler func(ctx context.Context, job Job), obs Observer) *Queue {
if cfg.Size <= 0 {
cfg.Size = 10000
}
if cfg.Workers <= 0 {
cfg.Workers = 8
}
if cfg.Timeout <= 0 {
cfg.Timeout = 4 * time.Second
}
q := &Queue{cfg: cfg, jobs: make(chan Job, cfg.Size), handler: handler, obs: obs}
for i := 0; i < cfg.Workers; i++ {
q.wg.Add(1)
go q.worker()
}
return q
}
// Enqueue hands a job to the workers without blocking. It returns false when
// the queue is full or draining; the caller answers 503 and logs the drop.
func (q *Queue) Enqueue(job Job) bool {
if job.Received.IsZero() {
job.Received = time.Now()
}
q.mu.RLock()
defer q.mu.RUnlock()
if q.closed {
return false
}
select {
case q.jobs <- job:
q.pending.Add(1)
q.observeDepth()
return true
default:
return false
}
}
// Depth is the number of jobs waiting for a worker.
func (q *Queue) Depth() int { return len(q.jobs) }
// Pending is the number of jobs enqueued and not yet finished (waiting or
// in flight); tests and the shutdown log use it.
func (q *Queue) Pending() int { return int(q.pending.Load()) }
// Drain stops accepting, lets the workers finish every queued job and waits
// for them, bounded by ctx. On an expired ctx it reports how many jobs were
// still pending.
func (q *Queue) Drain(ctx context.Context) error {
q.mu.Lock()
if !q.closed {
q.closed = true
close(q.jobs)
}
q.mu.Unlock()
done := make(chan struct{})
go func() { q.wg.Wait(); close(done) }()
select {
case <-done:
return nil
case <-ctx.Done():
return fmt.Errorf("counting queue: %d jobs still pending: %w", q.Pending(), ctx.Err())
}
}
func (q *Queue) worker() {
defer q.wg.Done()
for job := range q.jobs {
if q.obs.Wait != nil {
q.obs.Wait(time.Since(job.Received))
}
q.run(job)
q.pending.Add(-1)
q.observeDepth()
}
}
// run executes one job under the queue's own timeout, never a request
// context, and keeps the worker alive when the handler panics.
func (q *Queue) run(job Job) {
ctx, cancel := context.WithTimeout(context.Background(), q.cfg.Timeout)
defer cancel()
defer func() {
if p := recover(); p != nil && q.obs.Panic != nil {
q.obs.Panic(job, p, debug.Stack())
}
}()
q.handler(ctx, job)
}
func (q *Queue) observeDepth() {
if q.obs.Depth != nil {
q.obs.Depth(len(q.jobs))
}
}
internal/counting/queue_test.go
+179
View File
@@ -0,0 +1,179 @@
package counting
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
)
type recorder struct {
mu sync.Mutex
jobs []Job
waits []time.Duration
depths []int
panics int
}
func (r *recorder) handle(ctx context.Context, job Job) {
r.mu.Lock()
r.jobs = append(r.jobs, job)
r.mu.Unlock()
}
func (r *recorder) observer() Observer {
return Observer{
Depth: func(n int) { r.mu.Lock(); r.depths = append(r.depths, n); r.mu.Unlock() },
Wait: func(d time.Duration) { r.mu.Lock(); r.waits = append(r.waits, d); r.mu.Unlock() },
Panic: func(job Job, p any, stack []byte) { r.mu.Lock(); r.panics++; r.mu.Unlock() },
}
}
func (r *recorder) count() int {
r.mu.Lock()
defer r.mu.Unlock()
return len(r.jobs)
}
func settle(t *testing.T, q *Queue) {
t.Helper()
deadline := time.Now().Add(5 * time.Second)
for q.Pending() > 0 {
if time.Now().After(deadline) {
t.Fatalf("queue did not settle: pending=%d", q.Pending())
}
time.Sleep(time.Millisecond)
}
}
func TestQueue_ProcessesJobs(t *testing.T) {
r := &recorder{}
q := NewQueue(QueueConfig{Size: 10, Workers: 2, Timeout: time.Second}, r.handle, r.observer())
defer q.Drain(context.Background())
for i := 0; i < 5; i++ {
if !q.Enqueue(Job{Route: RouteIn, TraceID: "t", Recipient: "y@example.com", Received: time.Now()}) {
t.Fatalf("enqueue %d refused", i)
}
}
settle(t, q)
if r.count() != 5 {
t.Errorf("handled %d jobs, want 5", r.count())
}
r.mu.Lock()
defer r.mu.Unlock()
if len(r.waits) != 5 || len(r.depths) == 0 {
t.Errorf("observer: waits=%d depths=%d", len(r.waits), len(r.depths))
}
for _, w := range r.waits {
if w < 0 || w > time.Second {
t.Errorf("queue wait %v out of range", w)
}
}
}
// Workers run the handler with their own bounded context, never a request's.
func TestQueue_HandlerContextHasTimeout(t *testing.T) {
var got atomic.Value
q := NewQueue(QueueConfig{Size: 1, Workers: 1, Timeout: 250 * time.Millisecond}, func(ctx context.Context, job Job) {
dl, ok := ctx.Deadline()
got.Store([2]any{ok, time.Until(dl)})
}, Observer{})
defer q.Drain(context.Background())
q.Enqueue(Job{})
settle(t, q)
v := got.Load().([2]any)
if !v[0].(bool) || v[1].(time.Duration) > 250*time.Millisecond || v[1].(time.Duration) < 100*time.Millisecond {
t.Errorf("deadline = %v", v)
}
}
// A full queue refuses the job instead of blocking the HTTP handler.
func TestQueue_FullRefuses(t *testing.T) {
release := make(chan struct{})
q := NewQueue(QueueConfig{Size: 2, Workers: 1, Timeout: time.Second}, func(ctx context.Context, job Job) { <-release }, Observer{})
// One job occupies the worker ...
if !q.Enqueue(Job{}) {
t.Fatal("first enqueue refused")
}
deadline := time.Now().Add(time.Second)
for q.Depth() != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
// ... two fill the channel.
for i := 0; i < 2; i++ {
if !q.Enqueue(Job{}) {
t.Fatalf("enqueue %d should be accepted", i)
}
}
if q.Enqueue(Job{}) {
t.Error("enqueue on a full queue should be refused")
}
if q.Depth() != 2 || q.Pending() != 3 {
t.Errorf("depth=%d pending=%d, want 2/3", q.Depth(), q.Pending())
}
close(release)
if err := q.Drain(context.Background()); err != nil {
t.Fatal(err)
}
if q.Enqueue(Job{}) {
t.Error("enqueue after Drain must be refused, not panic")
}
}
// Drain waits for queued and in-flight jobs; a context that expires first
// reports what was left behind.
func TestQueue_Drain(t *testing.T) {
r := &recorder{}
slow := func(ctx context.Context, job Job) { time.Sleep(50 * time.Millisecond); r.handle(ctx, job) }
q := NewQueue(QueueConfig{Size: 10, Workers: 1, Timeout: time.Second}, slow, Observer{})
for i := 0; i < 4; i++ {
q.Enqueue(Job{})
}
start := time.Now()
if err := q.Drain(context.Background()); err != nil {
t.Fatal(err)
}
if r.count() != 4 || time.Since(start) < 150*time.Millisecond {
t.Errorf("Drain returned after %v with %d jobs handled, want all 4", time.Since(start), r.count())
}
block := make(chan struct{})
q2 := NewQueue(QueueConfig{Size: 10, Workers: 1, Timeout: time.Second}, func(ctx context.Context, job Job) { <-block }, Observer{})
q2.Enqueue(Job{})
q2.Enqueue(Job{})
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
err := q2.Drain(ctx)
if err == nil || !errors.Is(err, context.DeadlineExceeded) {
t.Errorf("Drain with a stuck worker should report the deadline: %v", err)
}
close(block)
}
// A panicking handler is recovered and reported; the worker keeps serving.
func TestQueue_RecoversPanic(t *testing.T) {
r := &recorder{}
q := NewQueue(QueueConfig{Size: 10, Workers: 1, Timeout: time.Second}, func(ctx context.Context, job Job) {
if job.TraceID == "boom" {
panic("boom")
}
r.handle(ctx, job)
}, r.observer())
defer q.Drain(context.Background())
q.Enqueue(Job{TraceID: "boom"})
q.Enqueue(Job{TraceID: "ok"})
settle(t, q)
if r.count() != 1 || r.panics != 1 {
t.Errorf("handled=%d panics=%d, want 1/1", r.count(), r.panics)
}
}
func TestQueue_Defaults(t *testing.T) {
q := NewQueue(QueueConfig{}, func(context.Context, Job) {}, Observer{})
defer q.Drain(context.Background())
if q.cfg.Size != 10000 || q.cfg.Workers != 8 || q.cfg.Timeout != 4*time.Second {
t.Errorf("defaults = %+v", q.cfg)
}
}
internal/domains/domains.go
+250
View File
@@ -0,0 +1,250 @@
// 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() }
internal/domains/domains_test.go
+389
View File
@@ -0,0 +1,389 @@
package domains
import (
"errors"
"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)
}
}
}
internal/metrics/metrics.go
+121
View File
@@ -0,0 +1,121 @@
// Package metrics defines the Prometheus metrics served on the admin port.
package metrics
import (
"github.com/prometheus/client_golang/prometheus"
)
// Metrics holds every collector; fields are exported for tests.
type Metrics struct {
Requests *prometheus.CounterVec // route, status
Duration *prometheus.HistogramVec // route
Increments *prometheus.CounterVec // level, flow
LimitReached *prometheus.CounterVec // level, flow
Notifications *prometheus.CounterVec // level, result (sent|failed|disabled)
RedisErrors *prometheus.CounterVec // op
CDBDomains prometheus.Gauge
CDBReloads *prometheus.CounterVec // result (success|failure)
Dropped *prometheus.CounterVec // route: requests refused because the counting queue was full
QueueDepth prometheus.Gauge
QueueWait prometheus.Histogram
}
// New creates the collectors and registers them with reg.
func New(reg prometheus.Registerer) *Metrics {
m := &Metrics{
Requests: prometheus.NewCounterVec(prometheus.CounterOpts{
Name: "migadu_counters_requests_total",
Help: "HTTP requests by route and status code.",
}, []string{"route", "status"}),
Duration: prometheus.NewHistogramVec(prometheus.HistogramOpts{
Name: "migadu_counters_request_duration_seconds",
Help: "Request latency by route.",
Buckets: []float64{.0005, .001, .0025, .005, .01, .025, .05, .1, .25, .5, 1, 2.5},
}, []string{"route"}),
Increments: prometheus.NewCounterVec(prometheus.CounterOpts{
Name: "migadu_counters_increments_total",
Help: "Counter increments by level (A, D, O) and flow (i, o).",
}, []string{"level", "flow"}),
LimitReached: prometheus.NewCounterVec(prometheus.CounterOpts{
Name: "migadu_counters_limit_reached_total",
Help: "Increments that crossed a configured limit, by level and flow.",
}, []string{"level", "flow"}),
Notifications: prometheus.NewCounterVec(prometheus.CounterOpts{
Name: "migadu_counters_notifications_total",
Help: "Web app notifications by level and result (sent, failed, disabled).",
}, []string{"level", "result"}),
RedisErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
Name: "migadu_counters_redis_errors_total",
Help: "Failed Redis operations by op (incr, mget, ping).",
}, []string{"op"}),
CDBDomains: prometheus.NewGauge(prometheus.GaugeOpts{
Name: "migadu_counters_cdb_domains",
Help: "Domains in the currently loaded domains CDB.",
}),
CDBReloads: prometheus.NewCounterVec(prometheus.CounterOpts{
Name: "migadu_counters_cdb_reloads_total",
Help: "Domains CDB reload attempts by result (success, failure).",
}, []string{"result"}),
Dropped: prometheus.NewCounterVec(prometheus.CounterOpts{
Name: "migadu_counters_dropped_total",
Help: "Validated requests refused with 503 because the counting queue was full, by route.",
}, []string{"route"}),
QueueDepth: prometheus.NewGauge(prometheus.GaugeOpts{
Name: "migadu_counters_queue_depth",
Help: "Jobs waiting in the counting queue.",
}),
QueueWait: prometheus.NewHistogram(prometheus.HistogramOpts{
Name: "migadu_counters_queue_wait_seconds",
Help: "Time a job waited in the counting queue before a worker took it.",
Buckets: []float64{.0001, .0005, .001, .0025, .005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5},
}),
}
reg.MustRegister(m.Requests, m.Duration, m.Increments, m.LimitReached, m.Notifications, m.RedisErrors, m.CDBDomains, m.CDBReloads,
m.Dropped, m.QueueDepth, m.QueueWait)
return m
}
// RecordDropped counts a request refused because the queue was full.
func (m *Metrics) RecordDropped(route string) { m.Dropped.WithLabelValues(route).Inc() }
// SetQueueDepth publishes the current queue depth.
func (m *Metrics) SetQueueDepth(n int) { m.QueueDepth.Set(float64(n)) }
// ObserveQueueWait records how long a job waited for a worker.
func (m *Metrics) ObserveQueueWait(seconds float64) { m.QueueWait.Observe(seconds) }
// RecordRequest counts one HTTP request and observes its latency.
func (m *Metrics) RecordRequest(route, status string, seconds float64) {
m.Requests.WithLabelValues(route, status).Inc()
m.Duration.WithLabelValues(route).Observe(seconds)
}
// RecordIncrement counts one counter increment.
func (m *Metrics) RecordIncrement(level, flow string) {
m.Increments.WithLabelValues(level, flow).Inc()
}
// RecordLimitReached counts one limit crossing.
func (m *Metrics) RecordLimitReached(level, flow string) {
m.LimitReached.WithLabelValues(level, flow).Inc()
}
// RecordNotification counts one notification outcome.
func (m *Metrics) RecordNotification(level, result string) {
m.Notifications.WithLabelValues(level, result).Inc()
}
// RecordRedisError counts one failed Redis operation.
func (m *Metrics) RecordRedisError(op string) { m.RedisErrors.WithLabelValues(op).Inc() }
// SetCDBDomains publishes the record count of the loaded CDB.
func (m *Metrics) SetCDBDomains(n int) { m.CDBDomains.Set(float64(n)) }
// RecordReload counts one CDB reload attempt.
func (m *Metrics) RecordReload(ok bool) {
result := "failure"
if ok {
result = "success"
}
m.CDBReloads.WithLabelValues(result).Inc()
}
internal/metrics/metrics_test.go
+111
View File
@@ -0,0 +1,111 @@
package metrics
import (
"testing"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/testutil"
)
func TestMetrics(t *testing.T) {
reg := prometheus.NewRegistry()
m := New(reg)
m.RecordRequest("out", "200", 0.001)
m.RecordRequest("out", "200", 0.002)
m.RecordRequest("in", "401", 0.0005)
m.RecordIncrement("A", "o")
m.RecordIncrement("A", "o")
m.RecordIncrement("O", "i")
m.RecordLimitReached("O", "o")
m.RecordNotification("O", "sent")
m.RecordNotification("A", "failed")
m.RecordRedisError("incr")
m.SetCDBDomains(93549)
m.RecordReload(true)
m.RecordReload(false)
m.RecordReload(false)
m.RecordDropped("/in")
m.RecordDropped("/in")
m.SetQueueDepth(7)
m.ObserveQueueWait(0.01)
checks := []struct {
name string
c prometheus.Collector
want float64
}{
{"requests out 200", m.Requests.WithLabelValues("out", "200"), 2},
{"requests in 401", m.Requests.WithLabelValues("in", "401"), 1},
{"increments A o", m.Increments.WithLabelValues("A", "o"), 2},
{"increments O i", m.Increments.WithLabelValues("O", "i"), 1},
{"limit reached O o", m.LimitReached.WithLabelValues("O", "o"), 1},
{"notifications O sent", m.Notifications.WithLabelValues("O", "sent"), 1},
{"notifications A failed", m.Notifications.WithLabelValues("A", "failed"), 1},
{"redis errors incr", m.RedisErrors.WithLabelValues("incr"), 1},
{"cdb domains", m.CDBDomains, 93549},
{"reload success", m.CDBReloads.WithLabelValues("success"), 1},
{"reload failure", m.CDBReloads.WithLabelValues("failure"), 2},
{"dropped /in", m.Dropped.WithLabelValues("/in"), 2},
{"queue depth", m.QueueDepth, 7},
}
for _, c := range checks {
if got := testutil.ToFloat64(c.c); got != c.want {
t.Errorf("%s = %v, want %v", c.name, got, c.want)
}
}
families, err := reg.Gather()
if err != nil {
t.Fatal(err)
}
seen := map[string]bool{}
for _, f := range families {
seen[f.GetName()] = true
if f.GetName() == "migadu_counters_request_duration_seconds" {
for _, metric := range f.GetMetric() {
for _, l := range metric.GetLabel() {
if l.GetName() == "route" && l.GetValue() == "out" && metric.GetHistogram().GetSampleCount() != 2 {
t.Errorf("duration histogram for out has %d samples, want 2", metric.GetHistogram().GetSampleCount())
}
}
}
}
}
for _, name := range []string{
"migadu_counters_requests_total",
"migadu_counters_request_duration_seconds",
"migadu_counters_increments_total",
"migadu_counters_limit_reached_total",
"migadu_counters_notifications_total",
"migadu_counters_redis_errors_total",
"migadu_counters_cdb_domains",
"migadu_counters_cdb_reloads_total",
"migadu_counters_dropped_total",
"migadu_counters_queue_depth",
"migadu_counters_queue_wait_seconds",
} {
if !seen[name] {
t.Errorf("metric %s not registered (have %v)", name, keys(seen))
}
}
}
func keys(m map[string]bool) []string {
out := make([]string, 0, len(m))
for k := range m {
out = append(out, k)
}
return out
}
// Two daemons' worth of metrics (tests, counterctl) must not collide: each
// registry gets its own collectors.
func TestNew_SeparateRegistries(t *testing.T) {
a := New(prometheus.NewRegistry())
b := New(prometheus.NewRegistry())
a.RecordRequest("in", "200", 0.1)
if testutil.ToFloat64(b.Requests.WithLabelValues("in", "200")) != 0 {
t.Error("registries share state")
}
}
internal/notify/notify.go
+236
View File
@@ -0,0 +1,236 @@
// Package notify tells the web app that a counter reached its limit.
//
// One URL template per level (org_url, address_url, domain_url) is filled
// with the event and fetched with GET + basic auth, in the background, with
// retries. The web app (migadu_web GET /api/limits) decides what happens to
// the customer; this package only delivers the event.
package notify
import (
"context"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strconv"
"strings"
"sync"
"time"
)
// Config holds the notification settings.
type Config struct {
OrgURL string // $LEVEL=O
AddressURL string // $LEVEL=A
DomainURL string // $LEVEL=D
User string // basic auth; empty = no Authorization header
Password string
Timeout time.Duration // per attempt (default 5s)
Attempts int // total attempts (default 5)
Backoff time.Duration // first retry delay, doubled each retry (default 1s)
MaxConcurrent int // notifications in flight (default 8)
}
// URLFor returns the template for a level, "" when that level is disabled.
func (c Config) URLFor(level string) string {
switch level {
case "O":
return c.OrgURL
case "A":
return c.AddressURL
case "D":
return c.DomainURL
}
return ""
}
// URLName returns the config key that serves a level, for log lines.
func URLName(level string) string {
switch level {
case "O":
return "org_url"
case "A":
return "address_url"
case "D":
return "domain_url"
}
return ""
}
// Event is one "limit reached" occurrence.
type Event struct {
TraceID string
Route string // in / out, for log correlation
Level string // A, D or O
ID string // address, domain or organisation id
Flow string // i or o
Address string // the address that triggered it ($ADDRESS / by=)
Count int64 // counter value after the increment ($COUNT)
Limit int64 // configured limit ($LIMIT)
}
// BuildURL fills the template placeholders with URL-escaped event values.
// $PERIOD is always "d" (per day).
func BuildURL(tmpl string, ev Event) string {
esc := url.QueryEscape
return strings.NewReplacer(
"$LEVEL", esc(ev.Level),
"$LIMIT", strconv.FormatInt(ev.Limit, 10),
"$ID", esc(ev.ID),
"$FLOW", esc(ev.Flow),
"$PERIOD", "d",
"$ADDRESS", esc(ev.Address),
"$COUNT", strconv.FormatInt(ev.Count, 10),
).Replace(tmpl)
}
// Hooks receive the outcome of every dispatched notification (transaction
// log, metrics). Nil hooks are ignored.
type Hooks struct {
Sent func(ev Event, url string, attempt int)
Failed func(ev Event, url string, attempts int, err error)
}
// Notifier sends events in the background.
type Notifier struct {
cfg Config
hooks Hooks
client *http.Client
sem chan struct{}
wg sync.WaitGroup
quit chan struct{}
mu sync.Mutex
closed bool
}
// New applies defaults and returns a ready Notifier.
func New(cfg Config, hooks Hooks) *Notifier {
if cfg.Timeout <= 0 {
cfg.Timeout = 5 * time.Second
}
if cfg.Attempts <= 0 {
cfg.Attempts = 5
}
if cfg.Backoff <= 0 {
cfg.Backoff = time.Second
}
if cfg.MaxConcurrent <= 0 {
cfg.MaxConcurrent = 8
}
return &Notifier{
cfg: cfg,
hooks: hooks,
client: &http.Client{Timeout: cfg.Timeout},
sem: make(chan struct{}, cfg.MaxConcurrent),
quit: make(chan struct{}),
}
}
// Notify dispatches ev in the background. It returns false, without doing
// anything, when no URL is configured for the level or after Shutdown.
func (n *Notifier) Notify(ev Event) bool {
tmpl := n.cfg.URLFor(ev.Level)
if tmpl == "" {
return false
}
n.mu.Lock()
if n.closed {
n.mu.Unlock()
return false
}
n.wg.Add(1)
n.mu.Unlock()
u := BuildURL(tmpl, ev)
go func() {
defer n.wg.Done()
n.sem <- struct{}{}
defer func() { <-n.sem }()
n.send(ev, u)
}()
return true
}
func (n *Notifier) send(ev Event, u string) {
var lastErr error
for attempt := 1; attempt <= n.cfg.Attempts; attempt++ {
err, retry := n.attempt(u)
if err == nil {
if n.hooks.Sent != nil {
n.hooks.Sent(ev, u, attempt)
}
return
}
lastErr = err
if !retry {
if n.hooks.Failed != nil {
n.hooks.Failed(ev, u, attempt, err)
}
return
}
if attempt < n.cfg.Attempts {
select {
case <-time.After(n.cfg.Backoff << (attempt - 1)):
case <-n.quit:
lastErr = fmt.Errorf("shutting down after attempt %d: %w", attempt, err)
attempt = n.cfg.Attempts // stop retrying
}
}
}
if n.hooks.Failed != nil {
n.hooks.Failed(ev, u, n.cfg.Attempts, lastErr)
}
}
// attempt performs one GET. retry is false for 4xx: wrong credentials or a
// wrong URL will not fix themselves.
func (n *Notifier) attempt(u string) (err error, retry bool) {
req, err := http.NewRequest(http.MethodGet, u, nil)
if err != nil {
return err, false
}
if n.cfg.User != "" {
req.SetBasicAuth(n.cfg.User, n.cfg.Password)
}
req.Header.Set("User-Agent", "migadu_counters")
resp, err := n.client.Do(req)
if err != nil {
return err, true
}
defer resp.Body.Close()
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4096))
switch {
case resp.StatusCode >= 200 && resp.StatusCode < 300:
return nil, false
case resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500:
return fmt.Errorf("unexpected status %d", resp.StatusCode), true
default:
return fmt.Errorf("unexpected status %d", resp.StatusCode), false
}
}
// Shutdown refuses new events, stops pending retries and waits for in-flight
// requests until ctx expires. Safe to call more than once.
func (n *Notifier) Shutdown(ctx context.Context) error {
n.mu.Lock()
if !n.closed {
n.closed = true
close(n.quit)
}
n.mu.Unlock()
done := make(chan struct{})
go func() {
n.wg.Wait()
close(done)
}()
select {
case <-done:
return nil
case <-ctx.Done():
return errors.Join(ctx.Err(), errors.New("notifications still in flight"))
}
}
internal/notify/notify_test.go
+332
View File
@@ -0,0 +1,332 @@
package notify
import (
"context"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
)
const tmpl = "https://admin.example/api/limits?level=$LEVEL&id=$ID&flow=$FLOW&period=$PERIOD&by=$ADDRESS&count=$COUNT&limit=$LIMIT"
func TestBuildURL(t *testing.T) {
ev := Event{Level: "A", ID: "me@example.com", Flow: "o", Address: "me@example.com", Count: 100, Limit: 100}
got := BuildURL(tmpl, ev)
want := "https://admin.example/api/limits?level=A&id=me%40example.com&flow=o&period=d&by=me%40example.com&count=100&limit=100"
if got != want {
t.Errorf("BuildURL =\n %s\nwant\n %s", got, want)
}
// Dangerous characters in values are escaped, templates without placeholders pass through.
ev.ID = "a b&c=d@x.y"
u, err := url.Parse(BuildURL(tmpl, ev))
if err != nil {
t.Fatal(err)
}
if u.Query().Get("id") != "a b&c=d@x.y" {
t.Errorf("id round-trip = %q", u.Query().Get("id"))
}
if got := BuildURL("https://x/static", ev); got != "https://x/static" {
t.Errorf("static template changed: %s", got)
}
}
func TestURLFor(t *testing.T) {
cfg := Config{OrgURL: "o", AddressURL: "a", DomainURL: "d"}
for level, want := range map[string]string{"O": "o", "A": "a", "D": "d", "X": ""} {
if got := cfg.URLFor(level); got != want {
t.Errorf("URLFor(%s) = %q, want %q", level, got, want)
}
}
}
// capture is an httptest handler recording requests.
type capture struct {
mu sync.Mutex
requests []*http.Request
status []int // status to answer per call; last entry repeats
calls atomic.Int32
delay time.Duration
inflight atomic.Int32
maxSeen atomic.Int32
}
func (c *capture) ServeHTTP(w http.ResponseWriter, r *http.Request) {
n := c.inflight.Add(1)
for {
m := c.maxSeen.Load()
if n <= m || c.maxSeen.CompareAndSwap(m, n) {
break
}
}
defer c.inflight.Add(-1)
if c.delay > 0 {
time.Sleep(c.delay)
}
i := int(c.calls.Add(1)) - 1
c.mu.Lock()
c.requests = append(c.requests, r.Clone(context.Background()))
c.mu.Unlock()
status := http.StatusOK
if len(c.status) > 0 {
if i < len(c.status) {
status = c.status[i]
} else {
status = c.status[len(c.status)-1]
}
}
w.WriteHeader(status)
}
type hooks struct {
mu sync.Mutex
sent []sentCall
failed []failedCall
done chan struct{}
}
type sentCall struct {
ev Event
url string
attempt int
}
type failedCall struct {
ev Event
url string
attempts int
err error
}
func newHooks() *hooks { return &hooks{done: make(chan struct{}, 64)} }
func (h *hooks) Hooks() Hooks {
return Hooks{
Sent: func(ev Event, url string, attempt int) {
h.mu.Lock()
h.sent = append(h.sent, sentCall{ev, url, attempt})
h.mu.Unlock()
h.done <- struct{}{}
},
Failed: func(ev Event, url string, attempts int, err error) {
h.mu.Lock()
h.failed = append(h.failed, failedCall{ev, url, attempts, err})
h.mu.Unlock()
h.done <- struct{}{}
},
}
}
func (h *hooks) wait(t *testing.T, n int) {
t.Helper()
for i := 0; i < n; i++ {
select {
case <-h.done:
case <-time.After(5 * time.Second):
t.Fatalf("timeout waiting for notification outcome %d/%d", i+1, n)
}
}
}
func fastConfig(orgURL, addrURL, domURL string) Config {
return Config{
OrgURL: orgURL, AddressURL: addrURL, DomainURL: domURL,
User: "api", Password: "secret",
Timeout: 2 * time.Second, Attempts: 5, Backoff: time.Millisecond, MaxConcurrent: 8,
}
}
func TestNotify_SendsToLevelURLWithBasicAuth(t *testing.T) {
org, addr, dom := &capture{}, &capture{}, &capture{}
so, sa, sd := httptest.NewServer(org), httptest.NewServer(addr), httptest.NewServer(dom)
defer so.Close()
defer sa.Close()
defer sd.Close()
q := "/api/limits?level=$LEVEL&id=$ID&flow=$FLOW&period=$PERIOD&by=$ADDRESS&count=$COUNT"
h := newHooks()
n := New(fastConfig(so.URL+q, sa.URL+q, sd.URL+q), h.Hooks())
defer n.Shutdown(context.Background())
ev := Event{TraceID: "tr", Level: "A", ID: "me@example.com", Flow: "o", Address: "me@example.com", Count: 100, Limit: 100}
if !n.Notify(ev) {
t.Fatal("Notify returned false for a configured level")
}
h.wait(t, 1)
if org.calls.Load() != 0 || dom.calls.Load() != 0 || addr.calls.Load() != 1 {
t.Fatalf("calls org=%d addr=%d dom=%d, want only addr", org.calls.Load(), addr.calls.Load(), dom.calls.Load())
}
r := addr.requests[0]
if r.Method != http.MethodGet {
t.Errorf("method = %s", r.Method)
}
if u, p, ok := r.BasicAuth(); !ok || u != "api" || p != "secret" {
t.Errorf("basic auth = %q %q %v", u, p, ok)
}
qv := r.URL.Query()
for k, want := range map[string]string{"level": "A", "id": "me@example.com", "flow": "o", "period": "d", "by": "me@example.com", "count": "100"} {
if qv.Get(k) != want {
t.Errorf("query %s = %q, want %q", k, qv.Get(k), want)
}
}
if len(h.sent) != 1 || h.sent[0].attempt != 1 || h.sent[0].ev != ev || !strings.HasPrefix(h.sent[0].url, sa.URL) {
t.Errorf("sent hook = %+v", h.sent)
}
// Org and domain levels route to their own URLs.
n.Notify(Event{Level: "O", ID: "7"})
n.Notify(Event{Level: "D", ID: "example.com"})
h.wait(t, 2)
if org.calls.Load() != 1 || dom.calls.Load() != 1 {
t.Errorf("calls org=%d dom=%d, want 1 each", org.calls.Load(), dom.calls.Load())
}
}
func TestNotify_EmptyURLSkips(t *testing.T) {
h := newHooks()
n := New(fastConfig("", "", ""), h.Hooks())
if n.Notify(Event{Level: "A"}) || n.Notify(Event{Level: "Z"}) {
t.Error("Notify should return false when no URL is configured for the level")
}
if err := n.Shutdown(context.Background()); err != nil {
t.Fatal(err)
}
if len(h.sent)+len(h.failed) != 0 {
t.Errorf("no hooks expected, got sent=%v failed=%v", h.sent, h.failed)
}
}
func TestNotify_NoBasicAuthWithoutUser(t *testing.T) {
c := &capture{}
s := httptest.NewServer(c)
defer s.Close()
h := newHooks()
cfg := fastConfig(s.URL+"/l?level=$LEVEL", "", "")
cfg.User, cfg.Password = "", ""
n := New(cfg, h.Hooks())
n.Notify(Event{Level: "O"})
h.wait(t, 1)
n.Shutdown(context.Background())
if _, _, ok := c.requests[0].BasicAuth(); ok {
t.Error("Authorization header sent although no user is configured")
}
}
func TestNotify_RetriesThenSucceeds(t *testing.T) {
c := &capture{status: []int{500, 503, 200}}
s := httptest.NewServer(c)
defer s.Close()
h := newHooks()
n := New(fastConfig(s.URL+"/l?level=$LEVEL", "", ""), h.Hooks())
n.Notify(Event{Level: "O", ID: "7"})
h.wait(t, 1)
n.Shutdown(context.Background())
if c.calls.Load() != 3 || len(h.sent) != 1 || h.sent[0].attempt != 3 || len(h.failed) != 0 {
t.Errorf("calls=%d sent=%+v failed=%+v; want success on attempt 3", c.calls.Load(), h.sent, h.failed)
}
}
func TestNotify_GivesUpAfterAttempts(t *testing.T) {
c := &capture{status: []int{500}}
s := httptest.NewServer(c)
defer s.Close()
h := newHooks()
n := New(fastConfig(s.URL+"/l?level=$LEVEL", "", ""), h.Hooks())
n.Notify(Event{Level: "O", ID: "7"})
h.wait(t, 1)
n.Shutdown(context.Background())
if c.calls.Load() != 5 || len(h.failed) != 1 || h.failed[0].attempts != 5 || len(h.sent) != 0 {
t.Errorf("calls=%d failed=%+v sent=%+v; want 5 attempts then failure", c.calls.Load(), h.failed, h.sent)
}
if !strings.Contains(h.failed[0].err.Error(), "500") {
t.Errorf("failure error should carry the status: %v", h.failed[0].err)
}
}
// 4xx means a configuration problem (bad credentials, wrong URL): retrying is pointless.
func TestNotify_NoRetryOn4xx(t *testing.T) {
c := &capture{status: []int{401}}
s := httptest.NewServer(c)
defer s.Close()
h := newHooks()
n := New(fastConfig(s.URL+"/l?level=$LEVEL", "", ""), h.Hooks())
n.Notify(Event{Level: "O", ID: "7"})
h.wait(t, 1)
n.Shutdown(context.Background())
if c.calls.Load() != 1 || len(h.failed) != 1 || h.failed[0].attempts != 1 {
t.Errorf("calls=%d failed=%+v; want a single attempt", c.calls.Load(), h.failed)
}
}
func TestNotify_TimeoutIsRetried(t *testing.T) {
c := &capture{delay: 300 * time.Millisecond}
s := httptest.NewServer(c)
defer s.Close()
h := newHooks()
cfg := fastConfig(s.URL+"/l?level=$LEVEL", "", "")
cfg.Timeout = 50 * time.Millisecond
cfg.Attempts = 2
n := New(cfg, h.Hooks())
n.Notify(Event{Level: "O", ID: "7"})
h.wait(t, 1)
n.Shutdown(context.Background())
if len(h.failed) != 1 || h.failed[0].attempts != 2 {
t.Errorf("failed=%+v; want failure after 2 timed-out attempts", h.failed)
}
}
func TestNotify_BoundedConcurrency(t *testing.T) {
c := &capture{delay: 50 * time.Millisecond}
s := httptest.NewServer(c)
defer s.Close()
h := newHooks()
cfg := fastConfig(s.URL+"/l?level=$LEVEL", "", "")
cfg.MaxConcurrent = 2
n := New(cfg, h.Hooks())
for i := 0; i < 10; i++ {
n.Notify(Event{Level: "O", ID: "7"})
}
h.wait(t, 10)
n.Shutdown(context.Background())
if c.maxSeen.Load() > 2 {
t.Errorf("max concurrent requests = %d, want <= 2", c.maxSeen.Load())
}
if c.calls.Load() != 10 {
t.Errorf("calls = %d, want 10", c.calls.Load())
}
}
// Shutdown waits for in-flight notifications, bounded by the context.
func TestNotify_ShutdownDrainsAndRespectsContext(t *testing.T) {
c := &capture{delay: 200 * time.Millisecond}
s := httptest.NewServer(c)
defer s.Close()
h := newHooks()
n := New(fastConfig(s.URL+"/l?level=$LEVEL", "", ""), h.Hooks())
n.Notify(Event{Level: "O", ID: "7"})
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer cancel()
if err := n.Shutdown(ctx); err == nil {
t.Error("Shutdown with an expired context should return its error")
}
if err := n.Shutdown(context.Background()); err != nil {
t.Errorf("second Shutdown: %v", err)
}
if c.calls.Load() != 1 || len(h.sent) != 1 {
t.Errorf("in-flight notification was not completed: calls=%d sent=%d", c.calls.Load(), len(h.sent))
}
if n.Notify(Event{Level: "O"}) {
t.Error("Notify after Shutdown should be refused")
}
}
func TestNew_Defaults(t *testing.T) {
n := New(Config{OrgURL: "https://x/?level=$LEVEL"}, Hooks{})
defer n.Shutdown(context.Background())
if n.cfg.Attempts != 5 || n.cfg.Backoff != time.Second || n.cfg.MaxConcurrent != 8 || n.cfg.Timeout != 5*time.Second {
t.Errorf("defaults not applied: %+v", n.cfg)
}
}
internal/server/middleware.go
+212
View File
@@ -0,0 +1,212 @@
package server
import (
"crypto/subtle"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"mime"
"net"
"net/http"
"net/netip"
"runtime/debug"
"strconv"
"strings"
"time"
"migadu_counters/internal/api"
)
// maxBody bounds every request body (the biggest legitimate body is a few
// hundred bytes).
const maxBody = 64 << 10
// httpError is a rejection by the validation layer.
type httpError struct {
status int
msg string
}
func (e *httpError) Error() string { return e.msg }
func writeJSON(w http.ResponseWriter, status int, v any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(v)
}
func writeError(w http.ResponseWriter, status int, msg string) {
writeJSON(w, status, api.ErrorResponse{Error: msg})
}
// reject answers an httpError and logs it: a burst of 4xx from mailqueuer
// means contract drift or a token mismatch.
func (s *Server) reject(w http.ResponseWriter, r *http.Request, route string, e *httpError) {
writeError(w, e.status, e.msg)
s.d.Logger.Warn("request rejected", "route", route, "remote", r.RemoteAddr, "status", e.status, "reason", e.msg)
}
// statusWriter captures the status code for metrics.
type statusWriter struct {
http.ResponseWriter
status int
}
func (w *statusWriter) WriteHeader(code int) {
if w.status == 0 {
w.status = code
}
w.ResponseWriter.WriteHeader(code)
}
func (w *statusWriter) Write(p []byte) (int, error) {
if w.status == 0 {
w.status = http.StatusOK
}
return w.ResponseWriter.Write(p)
}
// recoverer keeps the daemon alive when a handler panics. It sits outside
// http.TimeoutHandler, which re-panics in the caller's goroutine.
func (s *Server) recoverer(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
defer func() {
if p := recover(); p != nil {
s.d.Logger.Error("panic in handler", "panic", fmt.Sprint(p), "path", r.URL.Path, "stack", string(debug.Stack()))
writeError(w, http.StatusInternalServerError, "internal error")
}
}()
next.ServeHTTP(w, r)
})
}
// observe records request metrics per route and status.
func (s *Server) observe(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
sw := &statusWriter{ResponseWriter: w}
next.ServeHTTP(sw, r)
if sw.status == 0 {
sw.status = http.StatusOK
}
route := r.URL.Path
switch route {
case "/in", "/out", "/counters":
default:
route = "other"
}
if s.d.Metrics != nil {
s.d.Metrics.RecordRequest(route, strconv.Itoa(sw.status), time.Since(start).Seconds())
}
})
}
// route applies the validation steps common to every endpoint, in order:
// method, declared body size, bearer token, source IP, content type (POST).
func (s *Server) route(path, method string, h http.HandlerFunc) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != method {
w.Header().Set("Allow", method)
s.reject(w, r, path, &httpError{http.StatusMethodNotAllowed, "method not allowed"})
return
}
if r.ContentLength > maxBody {
s.reject(w, r, path, &httpError{http.StatusRequestEntityTooLarge, "request body too large"})
return
}
if e := s.checkAuth(r); e != nil {
w.Header().Set("WWW-Authenticate", "Bearer")
s.reject(w, r, path, e)
return
}
if e := s.checkSource(r); e != nil {
s.reject(w, r, path, e)
return
}
if method == http.MethodPost {
if e := checkContentType(r); e != nil {
s.reject(w, r, path, e)
return
}
r.Body = http.MaxBytesReader(w, r.Body, maxBody)
}
h(w, r)
})
}
func (s *Server) checkAuth(r *http.Request) *httpError {
if s.d.AuthToken == "" {
return nil // config.Warnings() flags this at startup
}
const prefix = "Bearer "
h := r.Header.Get("Authorization")
if !strings.HasPrefix(h, prefix) {
return &httpError{http.StatusUnauthorized, "unauthorized"}
}
if subtle.ConstantTimeCompare([]byte(strings.TrimSpace(h[len(prefix):])), []byte(s.d.AuthToken)) != 1 {
return &httpError{http.StatusUnauthorized, "unauthorized"}
}
return nil
}
// checkSource matches RemoteAddr (never X-Forwarded-For: nothing proxies
// mailqueuer's calls) against allowed_cidrs when configured.
func (s *Server) checkSource(r *http.Request) *httpError {
if len(s.d.AllowedCIDRs) == 0 {
return nil
}
host, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
host = r.RemoteAddr
}
addr, err := netip.ParseAddr(host)
if err != nil {
return &httpError{http.StatusForbidden, "source address not allowed"}
}
addr = addr.Unmap()
for _, p := range s.d.AllowedCIDRs {
if p.Contains(addr) {
return nil
}
}
return &httpError{http.StatusForbidden, "source address not allowed"}
}
func checkContentType(r *http.Request) *httpError {
mt, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
if err != nil || mt != "application/json" {
return &httpError{http.StatusUnsupportedMediaType, "Content-Type must be application/json"}
}
return nil
}
// decodeJSON reads exactly one JSON object into dst. Unknown fields are
// ignored so mailqueuer can add fields without breaking counting.
func decodeJSON(r *http.Request, dst any) *httpError {
dec := json.NewDecoder(r.Body)
if err := dec.Decode(dst); err != nil {
var tooBig *http.MaxBytesError
if errors.As(err, &tooBig) {
return &httpError{http.StatusRequestEntityTooLarge, "request body too large"}
}
return &httpError{http.StatusBadRequest, "invalid json: " + err.Error()}
}
if err := dec.Decode(&struct{}{}); err != io.EOF {
var tooBig *http.MaxBytesError
if errors.As(err, &tooBig) {
return &httpError{http.StatusRequestEntityTooLarge, "request body too large"}
}
return &httpError{http.StatusBadRequest, "body must contain exactly one JSON object"}
}
return nil
}
// logger returns the configured logger or slog's default.
func (d Deps) logger() *slog.Logger {
if d.Logger != nil {
return d.Logger
}
return slog.Default()
}
internal/server/server.go
+371
View File
@@ -0,0 +1,371 @@
// Package server serves the API (/in, /out, /counters) behind the HTTP
// validation layer, and the admin endpoints (/health, /metrics).
package server
import (
"context"
"errors"
"fmt"
"log/slog"
"net/http"
"net/netip"
"strconv"
"strings"
"time"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promhttp"
"migadu_counters/internal/api"
"migadu_counters/internal/counting"
"migadu_counters/internal/domains"
"migadu_counters/internal/metrics"
"migadu_counters/internal/notify"
"migadu_counters/internal/store"
"migadu_counters/internal/txlog"
)
// Deps wires the server to the rest of the daemon.
type Deps struct {
AuthToken string
AllowedCIDRs []netip.Prefix
HandlerTimeout time.Duration // bounds the synchronous part (validation) and each queued job
QueueSize int // counting queue capacity (default 10 000)
Workers int // counting workers (default 8)
Exempt map[string]bool
Counter *counting.Counter
Store *store.Store
Domains func() *domains.DB // current CDB (swapped on reload); nil = not loaded
Notifier *notify.Notifier // nil = never notify
TxLog *txlog.Logger // nil = transaction log disabled
Metrics *metrics.Metrics
Gatherer prometheus.Gatherer // served on /metrics; nil = no endpoint
Logger *slog.Logger
Now func() time.Time
}
// Server holds the handlers and the counting queue.
type Server struct {
d Deps
queue *counting.Queue
panicHook func() // tests: injected into the synchronous path
}
// New applies defaults and starts the counting workers.
func New(d Deps) *Server {
if d.Now == nil {
d.Now = time.Now
}
if d.HandlerTimeout <= 0 {
d.HandlerTimeout = 4 * time.Second
}
d.Logger = d.logger()
s := &Server{d: d}
obs := counting.Observer{
Panic: func(job counting.Job, p any, stack []byte) {
s.d.Logger.Error("panic while counting", "panic", fmt.Sprint(p), "route", job.Route, "trace_id", job.TraceID, "stack", string(stack))
},
}
if d.Metrics != nil {
obs.Depth = d.Metrics.SetQueueDepth
obs.Wait = func(w time.Duration) { d.Metrics.ObserveQueueWait(w.Seconds()) }
}
s.queue = counting.NewQueue(counting.QueueConfig{Size: d.QueueSize, Workers: d.Workers, Timeout: d.HandlerTimeout}, s.process, obs)
return s
}
// Drain stops accepting jobs and waits for the queued ones, bounded by ctx.
// Call it after the HTTP server stopped and before the notifier shuts down.
func (s *Server) Drain(ctx context.Context) error { return s.queue.Drain(ctx) }
// Pending is the number of accepted jobs not yet counted.
func (s *Server) Pending() int { return s.queue.Pending() }
// Depth is the number of jobs waiting for a worker.
func (s *Server) Depth() int { return s.queue.Depth() }
// NotifyHooks wires notification outcomes into the transaction log and the
// metrics; used by the daemon and by tests.
func NotifyHooks(tx *txlog.Logger, m *metrics.Metrics) notify.Hooks {
return notify.Hooks{
Sent: func(ev notify.Event, url string, attempt int) {
if m != nil {
m.RecordNotification(ev.Level, "sent")
}
if tx != nil {
tx.NotificationSent(ev, url, attempt)
}
},
Failed: func(ev notify.Event, url string, attempts int, err error) {
if m != nil {
m.RecordNotification(ev.Level, "failed")
}
if tx != nil {
tx.NotificationFailed(ev, url, attempts, err)
}
},
}
}
// Handler returns the API handler: recover → metrics → timeout → routes.
func (s *Server) Handler() http.Handler {
mux := http.NewServeMux()
mux.Handle("/out", s.route("/out", http.MethodPost, s.handleOut))
mux.Handle("/in", s.route("/in", http.MethodPost, s.handleIn))
mux.Handle("/counters", s.route("/counters", http.MethodGet, s.handleCounters))
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusNotFound, "not found")
})
timed := http.TimeoutHandler(mux, s.d.HandlerTimeout, `{"error":"request timeout"}`)
return s.recoverer(s.observe(timed))
}
// AdminHandler returns /health and /metrics (no auth; bind admin_addr locally).
func (s *Server) AdminHandler() http.Handler {
mux := http.NewServeMux()
mux.HandleFunc("/health", s.handleHealth)
if s.d.Gatherer != nil {
mux.Handle("/metrics", promhttp.HandlerFor(s.d.Gatherer, promhttp.HandlerOpts{}))
}
return mux
}
func (s *Server) handleOut(w http.ResponseWriter, r *http.Request) {
var req api.OutRequest
if e := decodeJSON(r, &req); e != nil {
s.reject(w, r, "/out", e)
return
}
v, err := api.ValidateOut(req)
if err != nil {
s.reject(w, r, "/out", &httpError{http.StatusBadRequest, err.Error()})
return
}
job := counting.Job{
Route: counting.RouteOut, TraceID: v.TraceID, Received: s.d.Now(),
AuthUser: v.AuthUser.Value, LoginSuffix: v.LoginSuffix,
MailFrom: v.MailFrom, MailFromUsable: v.MailFromUsable,
Recipient: v.Recipient.Value,
}
if !v.AuthUser.Usable() {
job.SenderUnusable = counting.Unusable{Role: "sender", Value: v.AuthUser.Raw, Reason: v.AuthUser.Reason}
}
if !v.Recipient.Usable() {
job.RecipientUnusable = counting.Unusable{Role: "recipient", Value: v.Recipient.Raw, Reason: v.Recipient.Reason}
}
s.enqueue(w, r, "/out", job)
}
func (s *Server) handleIn(w http.ResponseWriter, r *http.Request) {
var req api.InRequest
if e := decodeJSON(r, &req); e != nil {
s.reject(w, r, "/in", e)
return
}
v, err := api.ValidateIn(req)
if err != nil {
s.reject(w, r, "/in", &httpError{http.StatusBadRequest, err.Error()})
return
}
s.enqueue(w, r, "/in", counting.Job{
Route: counting.RouteIn, TraceID: v.TraceID, Received: s.d.Now(),
MailFrom: v.MailFrom, MailFromUsable: v.MailFromUsable, Recipient: v.Recipient, Junk: v.Junk,
})
}
// enqueue hands the validated job to the workers and answers 202; a full
// queue answers 503 and the drop is counted and logged.
func (s *Server) enqueue(w http.ResponseWriter, r *http.Request, route string, job counting.Job) {
if s.panicHook != nil {
s.panicHook()
}
if !s.queue.Enqueue(job) {
if s.d.Metrics != nil {
s.d.Metrics.RecordDropped(route)
}
if s.d.TxLog != nil {
s.d.TxLog.Dropped(job.Route, job.TraceID)
}
s.d.Logger.Error("counting queue full, request dropped", "route", route, "trace_id", job.TraceID, "remote", r.RemoteAddr)
writeError(w, http.StatusServiceUnavailable, "queue full")
return
}
writeJSON(w, http.StatusAccepted, api.CountResponse{Accepted: true})
}
// process is the worker body: plan, execute, dispatch notifications for
// crossed limits and write the transaction log. ctx is the queue's own
// bounded context, never a request's.
func (s *Server) process(ctx context.Context, job counting.Job) {
plan, err := s.d.Counter.Planner.Plan(job)
if err != nil {
s.d.Logger.Error("domain lookup failed", "route", job.Route, "trace_id", job.TraceID, "error", err)
if s.d.TxLog != nil {
s.d.TxLog.Error(job.Route, job.TraceID, err)
}
return
}
for _, lk := range plan.Lookups {
if errors.Is(lk.Err, domains.ErrMalformed) {
s.d.Logger.Warn("malformed row in domains db, domain treated as not hosted", "domain", lk.DomainName, "route", job.Route, "trace_id", job.TraceID, "error", lk.Err)
}
}
out, err := s.d.Counter.Execute(ctx, plan)
if err != nil {
if s.d.Metrics != nil {
s.d.Metrics.RecordRedisError("incr")
}
if s.d.TxLog != nil {
s.d.TxLog.Failed(plan, err)
}
s.d.Logger.Error("counting failed", "route", plan.Route, "trace_id", plan.TraceID, "error", err)
return
}
if out.Unknown != nil {
if s.d.Metrics != nil {
s.d.Metrics.RecordRedisError("incr_unknown")
}
s.d.Logger.Warn("redis outcome unknown, counts re-read from master", "route", plan.Route, "trace_id", plan.TraceID, "error", out.Unknown)
}
dispatched := make([]bool, len(out.Results))
for i, res := range out.Results {
if res.Err != nil {
if s.d.Metrics != nil {
s.d.Metrics.RecordRedisError("incr")
}
s.d.Logger.Error("redis error on key", "route", plan.Route, "trace_id", plan.TraceID, "key", res.Key, "error", res.Err)
continue
}
if s.d.Metrics != nil {
s.d.Metrics.RecordIncrement(res.Level, res.Flow)
}
if res.Reached {
if s.d.Metrics != nil {
s.d.Metrics.RecordLimitReached(res.Level, res.Flow)
}
ev := notify.Event{TraceID: plan.TraceID, Route: plan.Route, Level: res.Level, ID: res.ID,
Flow: res.Flow, Address: res.By, Count: res.Count, Limit: res.Limit}
dispatched[i] = s.d.Notifier != nil && s.d.Notifier.Notify(ev)
if !dispatched[i] && s.d.Metrics != nil {
s.d.Metrics.RecordNotification(res.Level, "disabled")
}
}
}
if s.d.TxLog != nil {
s.d.TxLog.Outcome(out, dispatched)
}
}
// handleCounters answers GET /counters?address=|domain=[&date=] from the
// read side.
func (s *Server) handleCounters(w http.ResponseWriter, r *http.Request) {
q := r.URL.Query()
address, domain := strings.TrimSpace(q.Get("address")), strings.ToLower(strings.TrimSpace(q.Get("domain")))
if (address == "") == (domain == "") {
s.reject(w, r, "/counters", &httpError{http.StatusBadRequest, "exactly one of address or domain is required"})
return
}
date, err := api.ParseDate(q.Get("date"), s.d.Now())
if err != nil {
s.reject(w, r, "/counters", &httpError{http.StatusBadRequest, "date: " + err.Error()})
return
}
resp := api.CountersResponse{Date: date, Levels: []api.CountersLevel{}}
var local string
if address != "" {
addr, err := api.NormalizeEmail(address)
if err != nil {
s.reject(w, r, "/counters", &httpError{http.StatusBadRequest, "address: invalid e-mail address"})
return
}
addr = domains.Mailbox(addr)
resp.Query, resp.Mode, resp.Exempt = addr, "address", s.d.Exempt[addr]
at := strings.LastIndex(addr, "@")
local, domain = addr[:at], addr[at+1:]
} else {
resp.Query, resp.Mode = domain, "domain"
}
db := s.d.Domains()
if db == nil {
writeError(w, http.StatusServiceUnavailable, "domains db not loaded")
return
}
d, err := db.Lookup(domain)
if errors.Is(err, domains.ErrMalformed) {
s.d.Logger.Warn("malformed row in domains db, domain treated as not hosted", "domain", domain, "route", "/counters", "error", err)
} else if err != nil && !errors.Is(err, domains.ErrNotFound) {
writeError(w, http.StatusServiceUnavailable, "domain lookup failed: "+err.Error())
return
}
managed := err == nil
if managed {
org := d.OrgID
resp.Managed, resp.Domain, resp.OrgID = true, d.Name, &org
}
type row struct{ level, id string }
var rows []row
if resp.Mode == "address" {
rows = append(rows, row{domains.AddressLevel, resp.Query})
}
rows = append(rows, row{domains.DomainLevel, domain})
if managed {
rows = append(rows, row{domains.OrganizationLevel, strconv.FormatInt(d.OrgID, 10)})
}
keys := make([]string, 0, 2*len(rows))
for _, rw := range rows {
keys = append(keys, store.Key(rw.level, rw.id, date, domains.FlowIn), store.Key(rw.level, rw.id, date, domains.FlowOut))
}
values, err := s.d.Store.Get(r.Context(), keys)
if err != nil {
if s.d.Metrics != nil {
s.d.Metrics.RecordRedisError("mget")
}
writeError(w, http.StatusServiceUnavailable, "redis unavailable: "+err.Error())
return
}
limit := func(level, flow string) *int64 {
if !managed {
return nil
}
if n, ok := d.Limits.Limit(level, flow, local); ok {
return &n
}
return nil
}
for i, rw := range rows {
resp.Levels = append(resp.Levels, api.CountersLevel{
Level: rw.level, ID: rw.id,
In: values[2*i], InLimit: limit(rw.level, domains.FlowIn),
Out: values[2*i+1], OutLimit: limit(rw.level, domains.FlowOut),
})
}
writeJSON(w, http.StatusOK, resp)
}
func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
db := s.d.Domains()
if db == nil {
writeError(w, http.StatusServiceUnavailable, "domains db not loaded")
return
}
if err := db.Ping(); err != nil {
writeError(w, http.StatusServiceUnavailable, "domains db not readable: "+err.Error())
return
}
ctx, cancel := context.WithTimeout(r.Context(), 200*time.Millisecond)
defer cancel()
if err := s.d.Store.Ping(ctx); err != nil {
if s.d.Metrics != nil {
s.d.Metrics.RecordRedisError("ping")
}
writeError(w, http.StatusServiceUnavailable, fmt.Sprintf("redis: %v", err))
return
}
writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}
internal/server/server_test.go
+1122
View File
File diff suppressed because it is too large Load Diff
internal/store/store.go
+238
View File
@@ -0,0 +1,238 @@
// Package store holds the Redis counters. Writes (INCR) go to the master,
// reads (MGET for /counters) go to the replicas; in Sentinel mode go-redis
// discovers both, in static mode the addresses come from the config.
//
// Key format, shared with the rspamd Lua plugin and read by migadu_web
// (apps/core/lib/core/redis_sentinel.ex):
//
// cnt/<A|D|O>/<id>/d/<YYYY-MM-DD>/<i|o>
//
// id is the address, the domain or the organisation id; the day is the UTC
// calendar day; every key lives 31 days so the web app can chart a month.
package store
import (
"context"
"errors"
"fmt"
"io"
"net"
"strconv"
"strings"
"sync/atomic"
"time"
"github.com/redis/go-redis/v9"
"migadu_counters/internal/config"
)
const (
// Prefix of every counter key.
Prefix = "cnt"
// TTL is refreshed on every increment, like the Lua plugin did.
TTL = 31 * 24 * time.Hour
)
// Key builds a counter key: Key("A", "x@example.com", "2026-10-13", "i").
func Key(level, id, date, flow string) string {
return Prefix + "/" + level + "/" + id + "/d/" + date + "/" + flow
}
// DayKey returns the UTC calendar day of t, as used in the keys. The web app
// builds its keys with Date.utc_today(), so this must never use local time.
func DayKey(t time.Time) string { return t.UTC().Format("2006-01-02") }
// Store is a write client plus one or more read clients.
type Store struct {
write *redis.Client
readers []*redis.Client
next atomic.Uint64
}
// New builds the clients from the config. Nothing is dialled until the first
// command, so an unreachable Redis does not prevent the daemon from starting.
func New(cfg config.RedisConfig) (*Store, error) {
if cfg.DB != 0 {
return nil, fmt.Errorf("redis.db must be 0 (got %d): migadu_web reads the message counters from DB 0", cfg.DB)
}
timeout := cfg.Timeout()
if cfg.SentinelMode() {
base := redis.FailoverOptions{
MasterName: cfg.MasterName,
SentinelAddrs: cfg.Sentinels,
SentinelPassword: cfg.SentinelPassword,
Password: cfg.Password,
DB: 0,
PoolSize: cfg.PoolSize,
DialTimeout: timeout,
ReadTimeout: timeout,
WriteTimeout: timeout,
}
write := redis.NewFailoverClient(&base)
// ReplicaOnly picks a random replica and falls back to the master when
// no replica is known.
replica := base
replica.ReplicaOnly = true
read := redis.NewFailoverClient(&replica)
return &Store{write: write, readers: []*redis.Client{read}}, nil
}
if cfg.WriteAddr == "" {
return nil, errors.New("redis.write_addr is required when no sentinels are configured")
}
opts := func(addr string) *redis.Options {
return &redis.Options{
Addr: addr,
Password: cfg.Password,
DB: 0,
PoolSize: cfg.PoolSize,
DialTimeout: timeout,
ReadTimeout: timeout,
WriteTimeout: timeout,
}
}
addrs := cfg.ReadAddrs
if len(addrs) == 0 {
addrs = []string{cfg.WriteAddr}
}
readers := make([]*redis.Client, 0, len(addrs))
for _, a := range addrs {
readers = append(readers, redis.NewClient(opts(a)))
}
return &Store{write: redis.NewClient(opts(cfg.WriteAddr)), readers: readers}, nil
}
// NewWithClients wires pre-built clients (tests). With no readers the write
// client is read from.
func NewWithClients(write *redis.Client, readers []*redis.Client) *Store {
if len(readers) == 0 {
readers = []*redis.Client{write}
}
return &Store{write: write, readers: readers}
}
// reader returns the next read client, round-robin.
func (s *Store) reader() *redis.Client {
n := s.next.Add(1) - 1
return s.readers[n%uint64(len(s.readers))]
}
// IncrResult is the outcome of one key's INCR inside the transaction.
type IncrResult struct {
Count int64 // value after the increment; meaningless when Err != nil
Err error // this key's own failure (e.g. it holds a non-number)
}
// Incr increments every key by one and refreshes its TTL in one MULTI/EXEC,
// returning one result per key in order. EXEC's replies are the
// authoritative counts, so no read-after-write is needed and replica lag
// does not matter. A per-key failure (a key holding a non-number) is reported
// in that key's result; the other keys are still incremented. The returned
// error is set only when the transaction as a whole failed, in which case
// OutcomeUnknown tells whether it may nevertheless have been applied.
func (s *Store) Incr(ctx context.Context, keys []string) ([]IncrResult, error) {
if len(keys) == 0 {
return nil, nil
}
cmds := make([]*redis.IntCmd, len(keys))
_, err := s.write.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
for i, k := range keys {
cmds[i] = pipe.Incr(ctx, k)
pipe.Expire(ctx, k, TTL)
}
return nil
})
out := make([]IncrResult, len(keys))
if err != nil {
// TxPipelined returns the first failed command's error. A Redis reply
// error (redis.Error) means EXEC answered and every command has a
// definite result: value or its own error. Anything else (dial, write,
// read timeout) is a transport failure of the whole transaction.
var rerr redis.Error
if !errors.As(err, &rerr) {
return nil, fmt.Errorf("redis incr: %w", err)
}
}
for i, c := range cmds {
if cerr := c.Err(); cerr != nil {
out[i].Err = fmt.Errorf("redis incr %s: %w", keys[i], cerr)
continue
}
out[i].Count = c.Val()
}
return out, nil
}
// OutcomeUnknown reports whether a failed Incr may still have been applied
// by Redis: the request may have been written but the reply lost (timeout,
// EOF, reset). A refused connection, a dial failure or a cancelled context
// never reached Redis and is a definite failure.
func OutcomeUnknown(err error) bool {
if err == nil || errors.Is(err, context.Canceled) {
return false
}
var opErr *net.OpError
if errors.As(err, &opErr) && opErr.Op == "dial" {
return false
}
if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
return true
}
var nerr net.Error
if errors.As(err, &nerr) && nerr.Timeout() {
return true
}
msg := err.Error()
return strings.Contains(msg, "timeout") || strings.Contains(msg, "reset") || strings.Contains(msg, "broken pipe") || strings.Contains(msg, "EOF")
}
// GetMaster reads the keys from the write side (the master) with one MGET;
// used to establish the outcome of an Incr whose reply was lost, when a
// replica might not have the increment yet. A missing key reads as 0.
func (s *Store) GetMaster(ctx context.Context, keys []string) ([]int64, error) {
return mget(ctx, s.write, keys)
}
// Get reads the keys from a replica with one MGET; a missing key reads as 0.
func (s *Store) Get(ctx context.Context, keys []string) ([]int64, error) {
return mget(ctx, s.reader(), keys)
}
func mget(ctx context.Context, c *redis.Client, keys []string) ([]int64, error) {
if len(keys) == 0 {
return nil, nil
}
vals, err := c.MGet(ctx, keys...).Result()
if err != nil {
return nil, fmt.Errorf("redis mget: %w", err)
}
out := make([]int64, len(keys))
for i, v := range vals {
str, ok := v.(string)
if !ok {
continue // nil = missing key
}
n, err := strconv.ParseInt(str, 10, 64)
if err != nil {
return nil, fmt.Errorf("redis key %s holds non-integer value %q", keys[i], str)
}
out[i] = n
}
return out, nil
}
// Ping checks the write side (the master), which is what /health cares about.
func (s *Store) Ping(ctx context.Context) error { return s.write.Ping(ctx).Err() }
// Close closes every client.
func (s *Store) Close() error {
errs := []error{s.write.Close()}
for _, r := range s.readers {
if r != s.write {
errs = append(errs, r.Close())
}
}
return errors.Join(errs...)
}
internal/store/store_test.go
+314
View File
@@ -0,0 +1,314 @@
package store
import (
"context"
"errors"
"io"
"net"
"strings"
"sync"
"testing"
"time"
"github.com/alicebob/miniredis/v2"
"github.com/redis/go-redis/v9"
"migadu_counters/internal/config"
)
func newMini(t *testing.T) *miniredis.Miniredis {
t.Helper()
m := miniredis.RunT(t)
return m
}
// testStore wires a Store to two separate miniredis instances so tests can
// prove which side (write master vs read replica) a method talks to.
func testStore(t *testing.T) (*Store, *miniredis.Miniredis, *miniredis.Miniredis) {
t.Helper()
w, r := newMini(t), newMini(t)
s := NewWithClients(
redis.NewClient(&redis.Options{Addr: w.Addr()}),
[]*redis.Client{redis.NewClient(&redis.Options{Addr: r.Addr()})},
)
t.Cleanup(func() { s.Close() })
return s, w, r
}
func TestKey(t *testing.T) {
tests := []struct {
level, id, date, flow, want string
}{
{"A", "x@example.com", "2026-10-13", "i", "cnt/A/x@example.com/d/2026-10-13/i"},
{"D", "example.com", "2026-10-13", "o", "cnt/D/example.com/d/2026-10-13/o"},
{"O", "1234", "2026-01-01", "i", "cnt/O/1234/d/2026-01-01/i"},
}
for _, tt := range tests {
if got := Key(tt.level, tt.id, tt.date, tt.flow); got != tt.want {
t.Errorf("Key(%s,%s,%s,%s) = %q, want %q", tt.level, tt.id, tt.date, tt.flow, got, tt.want)
}
}
}
// The day stamp is always the UTC calendar day, matching migadu_web's
// Date.utc_today() — never the server's local day.
func TestDayKey_UTC(t *testing.T) {
cest := time.FixedZone("CEST", 2*3600)
tests := []struct {
in time.Time
want string
}{
{time.Date(2026, 10, 13, 23, 59, 0, 0, cest), "2026-10-13"}, // 21:59 UTC
{time.Date(2026, 10, 14, 1, 30, 0, 0, cest), "2026-10-13"}, // 23:30 UTC the day before
{time.Date(2026, 10, 14, 2, 0, 0, 0, cest), "2026-10-14"}, // 00:00 UTC
{time.Date(2026, 3, 5, 12, 0, 0, 0, time.UTC), "2026-03-05"}, // zero padded
}
for _, tt := range tests {
if got := DayKey(tt.in); got != tt.want {
t.Errorf("DayKey(%v) = %q, want %q", tt.in, got, tt.want)
}
}
}
func TestIncr(t *testing.T) {
s, w, _ := testStore(t)
ctx := context.Background()
keys := []string{"cnt/A/x@example.com/d/2026-10-13/o", "cnt/O/1234/d/2026-10-13/o"}
got, err := s.Incr(ctx, keys)
if err != nil {
t.Fatalf("Incr: %v", err)
}
if len(got) != 2 || got[0].Count != 1 || got[1].Count != 1 || got[0].Err != nil || got[1].Err != nil {
t.Fatalf("first Incr = %+v, want counts [1 1]", got)
}
got, err = s.Incr(ctx, keys)
if err != nil || got[0].Count != 2 || got[1].Count != 2 {
t.Fatalf("second Incr = %+v, %v; want [2 2]", got, err)
}
for _, k := range keys {
if v, _ := w.Get(k); v != "2" {
t.Errorf("stored %s = %q, want \"2\"", k, v)
}
if ttl := w.TTL(k); ttl <= 30*24*time.Hour || ttl > TTL {
t.Errorf("TTL(%s) = %v, want (30d, 31d]", k, ttl)
}
}
if got, err := s.Incr(ctx, nil); err != nil || len(got) != 0 {
t.Errorf("Incr(no keys) = %v, %v; want empty, nil", got, err)
}
}
// One key holding a non-number fails its own INCR inside EXEC; the other
// keys are still incremented and their counts reported (review #4).
func TestIncr_PartialFailure(t *testing.T) {
s, w, _ := testStore(t)
w.Set("bad", "garbage")
got, err := s.Incr(context.Background(), []string{"k1", "bad", "k3"})
if err != nil {
t.Fatalf("a per-key failure must not fail the whole Incr: %v", err)
}
if got[0].Count != 1 || got[0].Err != nil || got[2].Count != 1 || got[2].Err != nil {
t.Errorf("good keys = %+v %+v, want count 1 without error", got[0], got[2])
}
if got[1].Err == nil || !strings.Contains(got[1].Err.Error(), "not an integer") {
t.Errorf("bad key err = %v, want the INCR error", got[1].Err)
}
if v, _ := w.Get("bad"); v != "garbage" {
t.Errorf("bad key = %q, must be untouched", v)
}
if ttl := w.TTL("k1"); ttl <= 30*24*time.Hour {
t.Errorf("TTL(k1) = %v, EXPIRE should still have run", ttl)
}
}
// GetMaster reads from the write side (used after a lost reply, when the
// replica may not have the increment yet).
func TestGetMaster(t *testing.T) {
s, w, r := testStore(t)
w.Set("k", "19")
r.Set("k", "3")
got, err := s.GetMaster(context.Background(), []string{"k", "missing"})
if err != nil || len(got) != 2 || got[0] != 19 || got[1] != 0 {
t.Errorf("GetMaster = %v, %v; want [19 0] from the master", got, err)
}
}
// OutcomeUnknown tells whether a failed Incr may still have been applied: a
// refused connection never reached Redis, a timeout or reset after the write
// might have.
func TestOutcomeUnknown(t *testing.T) {
s, w, _ := testStore(t)
w.Close()
_, err := s.Incr(context.Background(), []string{"k"})
if err == nil {
t.Fatal("expected an error with the write side down")
}
if OutcomeUnknown(err) {
t.Errorf("connection refused must be a known failure, got unknown: %v", err)
}
if OutcomeUnknown(nil) || OutcomeUnknown(context.Canceled) {
t.Error("nil / canceled are not unknown outcomes")
}
for _, e := range []error{context.DeadlineExceeded, io.EOF, &net.OpError{Op: "read", Err: timeoutErr{}}, errors.New("i/o timeout")} {
if !OutcomeUnknown(e) {
t.Errorf("%v should be an unknown outcome", e)
}
}
}
type timeoutErr struct{}
func (timeoutErr) Error() string { return "i/o timeout" }
func (timeoutErr) Timeout() bool { return true }
func (timeoutErr) Temporary() bool { return true }
// Fifty concurrent increments of one key must end at exactly 50 — the
// pipeline relies on INCRBY being atomic, not on read-then-write.
func TestIncr_Concurrent(t *testing.T) {
s, w, _ := testStore(t)
var wg sync.WaitGroup
for i := 0; i < 50; i++ {
wg.Add(1)
go func() {
defer wg.Done()
if _, err := s.Incr(context.Background(), []string{"k"}); err != nil {
t.Error(err)
}
}()
}
wg.Wait()
if v, _ := w.Get("k"); v != "50" {
t.Errorf("k = %q, want 50", v)
}
}
// Get reads from the replica side: a key present only on the write instance
// is invisible, a key present on the read instance is returned, missing = 0.
func TestGet_UsesReadSide(t *testing.T) {
s, w, r := testStore(t)
w.Set("only-on-write", "9")
r.Set("on-read", "5")
got, err := s.Get(context.Background(), []string{"on-read", "only-on-write", "missing"})
if err != nil {
t.Fatal(err)
}
if len(got) != 3 || got[0] != 5 || got[1] != 0 || got[2] != 0 {
t.Errorf("Get = %v, want [5 0 0]", got)
}
if got, err := s.Get(context.Background(), nil); err != nil || len(got) != 0 {
t.Errorf("Get(no keys) = %v, %v", got, err)
}
}
func TestGet_RoundRobinReaders(t *testing.T) {
w, r1, r2 := newMini(t), newMini(t), newMini(t)
r1.Set("k", "1")
r2.Set("k", "2")
s := NewWithClients(redis.NewClient(&redis.Options{Addr: w.Addr()}), []*redis.Client{
redis.NewClient(&redis.Options{Addr: r1.Addr()}),
redis.NewClient(&redis.Options{Addr: r2.Addr()}),
})
defer s.Close()
seen := map[int64]int{}
for i := 0; i < 4; i++ {
got, err := s.Get(context.Background(), []string{"k"})
if err != nil {
t.Fatal(err)
}
seen[got[0]]++
}
if seen[1] != 2 || seen[2] != 2 {
t.Errorf("reads not alternated between replicas: %v", seen)
}
}
func TestErrors_RedisDown(t *testing.T) {
s, w, r := testStore(t)
ctx := context.Background()
if err := s.Ping(ctx); err != nil {
t.Fatalf("Ping before close: %v", err)
}
w.Close()
if _, err := s.Incr(ctx, []string{"k"}); err == nil {
t.Error("Incr with write side down should fail")
}
if err := s.Ping(ctx); err == nil {
t.Error("Ping with write side down should fail")
}
r.Close()
if _, err := s.Get(ctx, []string{"k"}); err == nil {
t.Error("Get with read side down should fail")
}
}
func TestNew_Static(t *testing.T) {
w, r := newMini(t), newMini(t)
s, err := New(config.RedisConfig{WriteAddr: w.Addr(), ReadAddrs: []string{r.Addr()}, PoolSize: 2, TimeoutMS: 200})
if err != nil {
t.Fatalf("New: %v", err)
}
defer s.Close()
ctx := context.Background()
if err := s.Ping(ctx); err != nil {
t.Fatalf("Ping: %v", err)
}
if _, err := s.Incr(ctx, []string{"k"}); err != nil {
t.Fatal(err)
}
if v, _ := w.Get("k"); v != "1" {
t.Errorf("write went to %q, want the write_addr instance", v)
}
r.Set("k", "7")
if got, _ := s.Get(ctx, []string{"k"}); got[0] != 7 {
t.Errorf("Get = %v, want 7 from read_addrs", got)
}
}
func TestNew_StaticDefaultsReadToWrite(t *testing.T) {
w := newMini(t)
s, err := New(config.RedisConfig{WriteAddr: w.Addr(), PoolSize: 1, TimeoutMS: 200})
if err != nil {
t.Fatal(err)
}
defer s.Close()
w.Set("k", "3")
if got, err := s.Get(context.Background(), []string{"k"}); err != nil || got[0] != 3 {
t.Errorf("Get = %v, %v; want 3 from the write instance", got, err)
}
}
func TestNew_RejectsNonZeroDB(t *testing.T) {
if _, err := New(config.RedisConfig{WriteAddr: "127.0.0.1:6379", DB: 1, PoolSize: 1, TimeoutMS: 100}); err == nil {
t.Error("New with db=1 should fail: migadu_web reads DB 0")
}
}
// Sentinel mode constructs lazily; an unreachable sentinel surfaces on the
// first command, bounded by the configured timeout.
func TestNew_SentinelUnreachable(t *testing.T) {
s, err := New(config.RedisConfig{MasterName: "mymaster", Sentinels: []string{"127.0.0.1:1"}, PoolSize: 1, TimeoutMS: 200})
if err != nil {
t.Fatalf("New: %v", err)
}
defer s.Close()
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel()
start := time.Now()
if err := s.Ping(ctx); err == nil {
t.Error("Ping against an unreachable sentinel should fail")
}
if time.Since(start) > 2500*time.Millisecond {
t.Errorf("Ping took %v, should fail fast", time.Since(start))
}
}
func TestClose(t *testing.T) {
s, _, _ := testStore(t)
if err := s.Close(); err != nil {
t.Fatalf("Close: %v", err)
}
if err := s.Ping(context.Background()); err == nil {
t.Error("Ping after Close should fail")
}
}
internal/txlog/txlog.go
+300
View File
@@ -0,0 +1,300 @@
// Package txlog writes the counting transaction log: one block of lines per
// request, in the shape of the rspamd Lua plugin's log so existing grep
// habits keep working, with rspamd's task id replaced by the trace id:
//
// 2026-10-06 01:12:06.412 <ad8525> out; local domain: example.com account_id=7
// 2026-10-06 01:12:06.412 <ad8525> out; determined limit (level=[A] flow=[o] key=[cnt/A/x@example.com/d/2026-10-06/o]) => 100
// 2026-10-06 01:12:06.412 <ad8525> out; tracked address counters for x@example.com => count=12 limit=100 reached=false
//
// Each block is written with a single Write so concurrent requests never
// interleave. With format "json" the same events become one JSON object per
// line.
package txlog
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io"
"sort"
"strings"
"time"
"migadu_counters/internal/counting"
"migadu_counters/internal/domains"
"migadu_counters/internal/notify"
)
// Logger writes transaction blocks to w.
type Logger struct {
w io.Writer
json bool
now func() time.Time
}
// New creates a Logger. format is "text" (default) or "json"; now may be nil.
func New(w io.Writer, format string, now func() time.Time) *Logger {
if now == nil {
now = time.Now
}
return &Logger{w: w, json: format == "json", now: now}
}
// block accumulates the lines of one request.
type block struct {
l *Logger
ts time.Time
trace string
route string
buf bytes.Buffer
}
func (l *Logger) begin(trace, route string) *block {
return &block{l: l, ts: l.now().UTC(), trace: trace, route: route}
}
// line appends one event: msg in text mode, fields in json mode.
func (b *block) line(event, msg string, fields map[string]any) {
if b.l.json {
rec := map[string]any{
"ts": b.ts.Format("2006-01-02T15:04:05.000Z07:00"),
"trace_id": b.trace,
"route": b.route,
"event": event,
}
for k, v := range fields {
rec[k] = v
}
data, err := json.Marshal(rec)
if err != nil {
data = []byte(fmt.Sprintf(`{"event":"marshal_error","error":%q}`, err.Error()))
}
b.buf.Write(data)
b.buf.WriteByte('\n')
return
}
trace := b.trace
if trace == "" {
trace = "-"
}
fmt.Fprintf(&b.buf, "%s <%s> %s; %s\n", b.ts.Format("2006-01-02 15:04:05.000"), trace, b.route, msg)
}
func (b *block) flush() { _, _ = b.l.w.Write(b.buf.Bytes()) }
// Outcome logs an executed plan. dispatched[i] tells, for a Results[i] that
// reached its limit, whether a notification was handed to the notifier.
func (l *Logger) Outcome(out *counting.Outcome, dispatched []bool) {
b := l.begin(out.TraceID, out.Route)
b.request(out.Plan)
b.trail(out.Plan)
if out.Unknown != nil {
b.line("redis_unknown", fmt.Sprintf("redis outcome unknown: %v, re-read from master", out.Unknown),
map[string]any{"error": out.Unknown.Error(), "reread": true})
}
for _, r := range out.Results {
if r.Err != nil {
b.line("redis_key_error", fmt.Sprintf("redis error on key %s: %v (not counted)", r.Key, r.Err),
map[string]any{"level": r.Level, "id": r.ID, "flow": r.Flow, "key": r.Key, "error": r.Err.Error()})
continue
}
b.tracked(r, out.Unknown != nil)
}
for i, r := range out.Results {
if r.Reached {
b.reached(r, i < len(dispatched) && dispatched[i])
}
}
if out.Skipped != "" {
b.line("skipped", "skipped: "+out.Skipped, map[string]any{"reason": out.Skipped})
}
b.flush()
}
// Failed logs a plan whose increments could not be written to Redis.
func (l *Logger) Failed(plan *counting.Plan, err error) {
b := l.begin(plan.TraceID, plan.Route)
b.request(plan)
b.trail(plan)
b.line("redis_error", fmt.Sprintf("redis error: %v (nothing counted)", err), map[string]any{"error": err.Error()})
b.flush()
}
// Dropped logs a request that was rejected because the counting queue was full.
func (l *Logger) Dropped(route, traceID string) {
b := l.begin(traceID, route)
b.line("dropped", "dropped: queue full (not counted)", map[string]any{"reason": "queue full"})
b.flush()
}
// Error logs a request that failed before a plan existed (domain lookup error).
func (l *Logger) Error(route, traceID string, err error) {
b := l.begin(traceID, route)
b.line("error", fmt.Sprintf("error: %v (nothing counted)", err), map[string]any{"error": err.Error()})
b.flush()
}
// NotificationSent logs a delivered web app notification (notify hook).
func (l *Logger) NotificationSent(ev notify.Event, url string, attempt int) {
b := l.begin(ev.TraceID, ev.Route)
b.line("notification_sent", fmt.Sprintf("sent limit exceeded notification to url %q", url),
map[string]any{"url": url, "attempt": attempt, "level": ev.Level, "id": ev.ID, "flow": ev.Flow})
b.flush()
}
// NotificationFailed logs a notification given up on (notify hook).
func (l *Logger) NotificationFailed(ev notify.Event, url string, attempts int, err error) {
b := l.begin(ev.TraceID, ev.Route)
b.line("notification_failed", fmt.Sprintf("failed limit exceeded notification to url %q after %d attempts: %v", url, attempts, err),
map[string]any{"url": url, "attempts": attempts, "error": err.Error(), "level": ev.Level, "id": ev.ID, "flow": ev.Flow})
b.flush()
}
func (b *block) request(p *counting.Plan) {
if p.Route == counting.RouteOut {
msg := fmt.Sprintf("request auth_user=%s recipient=%s mail_from=%s", p.AuthUser, p.Recipient, p.MailFrom)
fields := map[string]any{"auth_user": p.AuthUser, "recipient": p.Recipient, "mail_from": p.MailFrom}
if p.LoginSuffix != "" {
msg += " login_suffix=" + p.LoginSuffix
fields["login_suffix"] = p.LoginSuffix
}
b.line("request", msg, fields)
} else {
b.line("request", fmt.Sprintf("request mail_from=%s recipient=%s junk=%t", p.MailFrom, p.Recipient, p.Junk),
map[string]any{"mail_from": p.MailFrom, "recipient": p.Recipient, "junk": p.Junk})
}
for _, u := range p.Unusable {
shown := u.Value
if shown == "" {
shown = "<empty>"
}
b.line("not_usable", fmt.Sprintf("%s address not usable: %s (%s, not counted)", u.Role, shown, u.Reason),
map[string]any{"role": u.Role, "value": u.Value, "reason": u.Reason})
}
}
// trail logs the domain lookups and the limits determined for each.
func (b *block) trail(p *counting.Plan) {
for _, lk := range p.Lookups {
if lk.Domain == nil {
if lk.Err == nil || errors.Is(lk.Err, domains.ErrNotFound) {
b.line("not_in_db", fmt.Sprintf("%s domain not in domains db: %s (not counted)", lk.Role, lk.DomainName),
map[string]any{"role": lk.Role, "domain": lk.DomainName})
} else if errors.Is(lk.Err, domains.ErrMalformed) {
detail := strings.TrimPrefix(lk.Err.Error(), domains.ErrMalformed.Error()+": ")
b.line("malformed_row", fmt.Sprintf("%s domain row malformed in domains db: %s: %s (not counted)", lk.Role, lk.DomainName, detail),
map[string]any{"role": lk.Role, "domain": lk.DomainName, "error": detail})
} else {
b.line("lookup_error", fmt.Sprintf("%s domain lookup failed: %s: %v", lk.Role, lk.DomainName, lk.Err),
map[string]any{"role": lk.Role, "domain": lk.DomainName, "error": lk.Err.Error()})
}
continue
}
d := lk.Domain
b.line("local_domain", fmt.Sprintf("local domain: %s account_id=%d", d.Name, d.OrgID),
map[string]any{"role": lk.Role, "domain": d.Name, "org_id": d.OrgID})
local := lk.Address
if i := strings.LastIndex(local, "@"); i >= 0 {
local = local[:i]
}
shown, other := summarize(d.Limits, local)
msg := fmt.Sprintf("domain: %s limits_map=%s", d.Name, renderLimits(shown, d.Limits))
if other > 0 {
msg += fmt.Sprintf(" (+%d other overrides)", other)
}
limitsJSON := make(map[string]int64, len(shown))
for _, k := range shown {
limitsJSON[k] = d.Limits[k]
}
b.line("limits", msg, map[string]any{"domain": d.Name, "limits": limitsJSON, "other_overrides": other})
flow := domains.FlowIn
if lk.Role == "sender" {
flow = domains.FlowOut
}
for _, inc := range p.Increments {
if inc.By != lk.Address || inc.Flow != flow {
continue
}
b.line("determined_limit",
fmt.Sprintf("determined limit (level=[%s] flow=[%s] key=[%s]) => %s", inc.Level, inc.Flow, inc.Key, limitText(inc)),
map[string]any{"level": inc.Level, "flow": inc.Flow, "key": inc.Key, "limit": limitJSON(inc), "limited": inc.Limited})
}
}
}
func (b *block) tracked(r counting.Result, reread bool) {
msg := fmt.Sprintf("tracked %s counters for %s => count=%d limit=%s reached=%t", levelWord(r.Level), r.ID, r.Count, limitText(r.Increment), r.Reached)
fields := map[string]any{"level": r.Level, "id": r.ID, "flow": r.Flow, "key": r.Key, "count": r.Count, "limit": limitJSON(r.Increment), "reached": r.Reached}
if reread {
msg += " (re-read)"
fields["reread"] = true
}
b.line("tracked", msg, fields)
}
func (b *block) reached(r counting.Result, dispatched bool) {
name := notify.URLName(r.Level)
verdict := "notifying " + name
if !dispatched {
verdict = "no " + name + " configured, not notifying"
}
b.line("limit_reached",
fmt.Sprintf("limit reached (level=[%s] id=[%s] flow=[%s] by=[%s]) => %s", r.Level, r.ID, r.Flow, r.By, verdict),
map[string]any{"level": r.Level, "id": r.ID, "flow": r.Flow, "by": r.By, "count": r.Count, "limit": r.Limit, "notifying": dispatched, "url_name": name})
}
// summarize picks the limit keys worth printing for one address: the base
// keys and the address's own overrides. Some domains carry thousands of
// per-mailbox overrides (boxmailhere.com: 9 798); dumping them all would make
// a 190 KB log line.
func summarize(l domains.Limits, local string) (shown []string, other int) {
keys := make([]string, 0, len(l))
for k := range l {
keys = append(keys, k)
}
sort.Strings(keys)
for _, k := range keys {
if !strings.Contains(k, "@/") || strings.HasPrefix(k, local+"@/") {
shown = append(shown, k)
} else {
other++
}
}
return shown, other
}
func renderLimits(keys []string, l domains.Limits) string {
parts := make([]string, len(keys))
for i, k := range keys {
parts[i] = fmt.Sprintf("%s=%d", k, l[k])
}
return "{" + strings.Join(parts, ", ") + "}"
}
func levelWord(level string) string {
switch level {
case domains.AddressLevel:
return "address"
case domains.DomainLevel:
return "domain"
case domains.OrganizationLevel:
return "account" // the Lua plugin's word for the organisation level
}
return level
}
func limitText(inc counting.Increment) string {
if !inc.Limited {
return "unlimited"
}
return fmt.Sprintf("%d", inc.Limit)
}
func limitJSON(inc counting.Increment) any {
if !inc.Limited {
return nil
}
return inc.Limit
}
internal/txlog/txlog_test.go
+413
View File
@@ -0,0 +1,413 @@
package txlog
import (
"bytes"
"encoding/json"
"errors"
"strings"
"sync"
"testing"
"time"
"migadu_counters/internal/counting"
"migadu_counters/internal/domains"
"migadu_counters/internal/notify"
)
var clock = time.Date(2026, 10, 6, 1, 12, 6, 412_000_000, time.UTC)
const ts = "2026-10-06 01:12:06.412"
type table map[string]*domains.Domain
func (t table) Lookup(d string) (*domains.Domain, error) {
if r, ok := t[d]; ok {
return r, nil
}
return nil, domains.ErrNotFound
}
func planner() *counting.Planner {
return &counting.Planner{
Lookup: table{
"wearemiraclesnow.com": {Name: "wearemiraclesnow.com", OrgID: 116681,
Limits: domains.ParseLimits("O/i=10000,O/o=2200,D/i=10000,D/o=2200,A/i=500,A/o=100")},
"nolimits.net": {Name: "nolimits.net", OrgID: 11, Limits: domains.Limits{}},
"many.example": {Name: "many.example", OrgID: 5,
Limits: domains.ParseLimits("A/i=1,A/o=2,bob@/i=3,bob@/o=4,carol@/i=5,dave@/o=6")},
},
Exempt: map[string]bool{"support@migadu.com": true},
Now: func() time.Time { return clock },
}
}
// outcome executes nothing: it attaches the given counts to the plan, with
// Reached derived exactly as Counter does.
func outcome(plan *counting.Plan, counts ...int64) *counting.Outcome {
out := &counting.Outcome{Plan: plan}
for i, inc := range plan.Increments {
out.Results = append(out.Results, counting.Result{
Increment: inc, Count: counts[i],
Reached: inc.Limited && counting.Crossed(counts[i], 1, inc.Limit),
})
}
return out
}
func newLogger(format string) (*Logger, *bytes.Buffer) {
var buf bytes.Buffer
return New(&buf, format, func() time.Time { return clock }), &buf
}
func TestOutcome_OutText(t *testing.T) {
plan, err := planner().PlanOut("ad8525", "noah@wearemiraclesnow.com", "noah@wearemiraclesnow.com", "y@example.org")
if err != nil {
t.Fatal(err)
}
l, buf := newLogger("text")
l.Outcome(outcome(plan, 12, 1830, 2200), []bool{false, false, true})
want := strings.Join([]string{
ts + " <ad8525> out; request auth_user=noah@wearemiraclesnow.com recipient=y@example.org mail_from=noah@wearemiraclesnow.com",
ts + " <ad8525> out; local domain: wearemiraclesnow.com account_id=116681",
ts + " <ad8525> out; domain: wearemiraclesnow.com limits_map={A/i=500, A/o=100, D/i=10000, D/o=2200, O/i=10000, O/o=2200}",
ts + " <ad8525> out; determined limit (level=[A] flow=[o] key=[cnt/A/noah@wearemiraclesnow.com/d/2026-10-06/o]) => 100",
ts + " <ad8525> out; determined limit (level=[D] flow=[o] key=[cnt/D/wearemiraclesnow.com/d/2026-10-06/o]) => 2200",
ts + " <ad8525> out; determined limit (level=[O] flow=[o] key=[cnt/O/116681/d/2026-10-06/o]) => 2200",
ts + " <ad8525> out; recipient domain not in domains db: example.org (not counted)",
ts + " <ad8525> out; tracked address counters for noah@wearemiraclesnow.com => count=12 limit=100 reached=false",
ts + " <ad8525> out; tracked domain counters for wearemiraclesnow.com => count=1830 limit=2200 reached=false",
ts + " <ad8525> out; tracked account counters for 116681 => count=2200 limit=2200 reached=true",
ts + " <ad8525> out; limit reached (level=[O] id=[116681] flow=[o] by=[noah@wearemiraclesnow.com]) => notifying org_url",
"",
}, "\n")
if got := buf.String(); got != want {
t.Errorf("block mismatch\n--- got ---\n%s--- want ---\n%s", got, want)
}
}
func TestOutcome_InTextUnlimited(t *testing.T) {
plan, err := planner().PlanIn("tr2", "s@gmail.com", "q@nolimits.net", false)
if err != nil {
t.Fatal(err)
}
l, buf := newLogger("text")
l.Outcome(outcome(plan, 1, 1, 1), nil)
want := strings.Join([]string{
ts + " <tr2> in; request mail_from=s@gmail.com recipient=q@nolimits.net junk=false",
ts + " <tr2> in; local domain: nolimits.net account_id=11",
ts + " <tr2> in; domain: nolimits.net limits_map={}",
ts + " <tr2> in; determined limit (level=[A] flow=[i] key=[cnt/A/q@nolimits.net/d/2026-10-06/i]) => unlimited",
ts + " <tr2> in; determined limit (level=[D] flow=[i] key=[cnt/D/nolimits.net/d/2026-10-06/i]) => unlimited",
ts + " <tr2> in; determined limit (level=[O] flow=[i] key=[cnt/O/11/d/2026-10-06/i]) => unlimited",
ts + " <tr2> in; tracked address counters for q@nolimits.net => count=1 limit=unlimited reached=false",
ts + " <tr2> in; tracked domain counters for nolimits.net => count=1 limit=unlimited reached=false",
ts + " <tr2> in; tracked account counters for 11 => count=1 limit=unlimited reached=false",
"",
}, "\n")
if got := buf.String(); got != want {
t.Errorf("block mismatch\n--- got ---\n%s--- want ---\n%s", got, want)
}
}
func TestOutcome_ReachedWithoutURL(t *testing.T) {
plan, _ := planner().PlanIn("tr", "s@gmail.com", "y@wearemiraclesnow.com", false)
l, buf := newLogger("text")
// D/i=10000 reached, A and O not; domain_url not configured -> dispatched false.
l.Outcome(outcome(plan, 3, 10000, 9), []bool{false, false, false})
want := ts + " <tr> in; limit reached (level=[D] id=[wearemiraclesnow.com] flow=[i] by=[y@wearemiraclesnow.com]) => no domain_url configured, not notifying\n"
if got := buf.String(); !strings.HasSuffix(got, want) {
t.Errorf("missing not-notifying line\n--- got ---\n%s--- want suffix ---\n%s", got, want)
}
if strings.Count(buf.String(), "limit reached") != 1 {
t.Errorf("want exactly one limit reached line:\n%s", buf.String())
}
}
func TestOutcome_Skips(t *testing.T) {
p := planner()
cases := []struct {
name string
plan func() (*counting.Plan, error)
want []string
}{
{"exempt", func() (*counting.Plan, error) {
return p.PlanIn("t", "support@migadu.com", "y@wearemiraclesnow.com", false)
},
[]string{
ts + " <t> in; request mail_from=support@migadu.com recipient=y@wearemiraclesnow.com junk=false",
ts + " <t> in; skipped: exempt",
}},
{"junk", func() (*counting.Plan, error) { return p.PlanIn("t", "s@gmail.com", "y@wearemiraclesnow.com", true) },
[]string{
ts + " <t> in; request mail_from=s@gmail.com recipient=y@wearemiraclesnow.com junk=true",
ts + " <t> in; skipped: junk",
}},
{"unknown", func() (*counting.Plan, error) { return p.PlanOut("t", "a@unknown.tld", "", "b@gmail.com") },
[]string{
ts + " <t> out; request auth_user=a@unknown.tld recipient=b@gmail.com mail_from=",
ts + " <t> out; sender domain not in domains db: unknown.tld (not counted)",
ts + " <t> out; recipient domain not in domains db: gmail.com (not counted)",
ts + " <t> out; skipped: unknown_domain",
}},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
plan, err := c.plan()
if err != nil {
t.Fatal(err)
}
l, buf := newLogger("text")
l.Outcome(outcome(plan), nil)
if got, want := buf.String(), strings.Join(c.want, "\n")+"\n"; got != want {
t.Errorf("--- got ---\n%s--- want ---\n%s", got, want)
}
})
}
}
func TestFailed_RedisError(t *testing.T) {
plan, _ := planner().PlanIn("tr", "s@gmail.com", "y@wearemiraclesnow.com", false)
l, buf := newLogger("text")
l.Failed(plan, errors.New("redis incr: dial tcp 127.0.0.1:6379: connect: connection refused"))
got := buf.String()
if !strings.HasPrefix(got, ts+" <tr> in; request mail_from=s@gmail.com") {
t.Errorf("missing request line:\n%s", got)
}
if !strings.Contains(got, "determined limit (level=[O]") {
t.Errorf("lookup trail missing before the error:\n%s", got)
}
if !strings.HasSuffix(got, ts+" <tr> in; redis error: redis incr: dial tcp 127.0.0.1:6379: connect: connection refused (nothing counted)\n") {
t.Errorf("missing error line:\n%s", got)
}
if strings.Contains(got, "tracked") {
t.Errorf("no tracked lines expected when nothing was counted:\n%s", got)
}
}
func TestNotificationLines(t *testing.T) {
l, buf := newLogger("text")
ev := notify.Event{TraceID: "ad8525", Route: "out", Level: "O", ID: "116681", Flow: "o", Address: "noah@wearemiraclesnow.com", Count: 2200, Limit: 2200}
u := "https://admin.migadu.com/api/limits?level=O&id=116681&flow=o&period=d&by=noah%40wearemiraclesnow.com&count=2200"
l.NotificationSent(ev, u, 1)
l.NotificationFailed(ev, u, 5, errors.New("unexpected status 500"))
want := ts + ` <ad8525> out; sent limit exceeded notification to url "` + u + `"` + "\n" +
ts + ` <ad8525> out; failed limit exceeded notification to url "` + u + `" after 5 attempts: unexpected status 500` + "\n"
if got := buf.String(); got != want {
t.Errorf("--- got ---\n%s--- want ---\n%s", got, want)
}
}
func TestEmptyTraceID(t *testing.T) {
plan, _ := planner().PlanIn("", "s@gmail.com", "y@wearemiraclesnow.com", true)
l, buf := newLogger("text")
l.Outcome(outcome(plan), nil)
if !strings.HasPrefix(buf.String(), ts+" <-> in; request") {
t.Errorf("empty trace id should print as <->:\n%s", buf.String())
}
}
// Domains with thousands of per-address overrides (boxmailhere.com has 9 798)
// must not dump them all: base keys + the address's own overrides + a count.
func TestLimitsMapSummarised(t *testing.T) {
plan, _ := planner().PlanOut("t", "bob@many.example", "", "ext@gmail.com")
l, buf := newLogger("text")
l.Outcome(outcome(plan, 1, 1, 1), nil)
want := ts + " <t> out; domain: many.example limits_map={A/i=1, A/o=2, bob@/i=3, bob@/o=4} (+2 other overrides)\n"
if !strings.Contains(buf.String(), want) {
t.Errorf("--- got ---\n%s--- want line ---\n%s", buf.String(), want)
}
}
func TestOutcome_JSON(t *testing.T) {
plan, _ := planner().PlanOut("ad8525", "noah@wearemiraclesnow.com", "noah@wearemiraclesnow.com", "y@example.org")
l, buf := newLogger("json")
l.Outcome(outcome(plan, 12, 1830, 2200), []bool{false, false, true})
lines := strings.Split(strings.TrimSpace(buf.String()), "\n")
var events []string
var records []map[string]any
for _, line := range lines {
var rec map[string]any
if err := json.Unmarshal([]byte(line), &rec); err != nil {
t.Fatalf("line is not JSON: %q: %v", line, err)
}
if rec["trace_id"] != "ad8525" || rec["route"] != "out" || rec["ts"] != "2026-10-06T01:12:06.412Z" {
t.Errorf("common fields wrong: %v", rec)
}
events = append(events, rec["event"].(string))
records = append(records, rec)
}
wantEvents := []string{"request", "local_domain", "limits", "determined_limit", "determined_limit", "determined_limit",
"not_in_db", "tracked", "tracked", "tracked", "limit_reached"}
if strings.Join(events, ",") != strings.Join(wantEvents, ",") {
t.Errorf("events = %v\nwant %v", events, wantEvents)
}
o := records[9]
if o["level"] != "O" || o["id"] != "116681" || o["count"].(float64) != 2200 || o["limit"].(float64) != 2200 || o["reached"] != true ||
o["key"] != "cnt/O/116681/d/2026-10-06/o" || o["flow"] != "o" {
t.Errorf("tracked O record = %v", o)
}
if r := records[10]; r["notifying"] != true || r["url_name"] != "org_url" || r["by"] != "noah@wearemiraclesnow.com" {
t.Errorf("limit_reached record = %v", r)
}
if r := records[1]; r["domain"] != "wearemiraclesnow.com" || r["org_id"].(float64) != 116681 {
t.Errorf("local_domain record = %v", r)
}
}
// countingWriter counts Write calls: each block must be one write so
// concurrent requests never interleave their lines.
type countingWriter struct {
mu sync.Mutex
writes int
buf bytes.Buffer
}
func (w *countingWriter) Write(p []byte) (int, error) {
w.mu.Lock()
defer w.mu.Unlock()
w.writes++
return w.buf.Write(p)
}
func TestBlockIsSingleWrite(t *testing.T) {
w := &countingWriter{}
l := New(w, "text", func() time.Time { return clock })
plan, _ := planner().PlanOut("t", "noah@wearemiraclesnow.com", "", "y@wearemiraclesnow.com")
l.Outcome(outcome(plan, 1, 1, 1, 1, 1, 1), nil)
l.Failed(plan, errors.New("x"))
l.NotificationSent(notify.Event{}, "u", 1)
if w.writes != 3 {
t.Errorf("writes = %d, want 3 (one per block/line)", w.writes)
}
}
func TestErrorLine(t *testing.T) {
l, buf := newLogger("text")
l.Error("out", "tr", errors.New("lookup sender domain x.com: cdb read error"))
want := ts + " <tr> out; error: lookup sender domain x.com: cdb read error (nothing counted)\n"
if got := buf.String(); got != want {
t.Errorf("got %q, want %q", got, want)
}
}
// --- review round 1 lines -----------------------------------------------------
func TestRequestLine_LoginSuffixAndUnusable(t *testing.T) {
plan, err := planner().Plan(counting.Job{Route: counting.RouteOut, TraceID: "m1", AuthUser: "x@wearemiraclesnow.com", LoginSuffix: "SUPPORTTOKEN",
MailFrom: "x@wearemiraclesnow.com", MailFromUsable: true,
RecipientUnusable: counting.Unusable{Role: "recipient", Value: "u@[192.0.2.1]", Reason: "invalid e-mail address"}})
if err != nil {
t.Fatal(err)
}
l, buf := newLogger("text")
l.Outcome(outcome(plan, 1, 1, 1), nil)
got := buf.String()
if !strings.Contains(got, ts+" <m1> out; request auth_user=x@wearemiraclesnow.com recipient= mail_from=x@wearemiraclesnow.com login_suffix=SUPPORTTOKEN\n") {
t.Errorf("request line should show the login suffix:\n%s", got)
}
if !strings.Contains(got, ts+" <m1> out; recipient address not usable: u@[192.0.2.1] (invalid e-mail address, not counted)\n") {
t.Errorf("missing unusable line:\n%s", got)
}
if !strings.Contains(got, "tracked address counters for x@wearemiraclesnow.com => count=1") {
t.Errorf("sender should still be tracked:\n%s", got)
}
plan, _ = planner().Plan(counting.Job{Route: counting.RouteOut, TraceID: "m2", Recipient: "y@wearemiraclesnow.com",
SenderUnusable: counting.Unusable{Role: "sender", Reason: "missing"}})
l, buf = newLogger("text")
l.Outcome(outcome(plan, 1, 1, 1), nil)
if got := buf.String(); !strings.Contains(got, " <m2> out; sender address not usable: <empty> (missing, not counted)\n") || strings.Contains(got, "login_suffix") {
t.Errorf("missing sender line / unexpected suffix:\n%s", got)
}
l, buf = newLogger("json")
l.Outcome(outcome(plan, 1, 1, 1), nil)
if got := buf.String(); !strings.Contains(got, `"event":"not_usable"`) || !strings.Contains(got, `"role":"sender"`) || !strings.Contains(got, `"reason":"missing"`) {
t.Errorf("json not_usable event missing:\n%s", got)
}
}
// A key whose own INCR failed gets an error line instead of a tracked line;
// the other keys are tracked normally.
func TestOutcome_PerKeyError(t *testing.T) {
plan, _ := planner().PlanOut("p1", "noah@wearemiraclesnow.com", "", "s@gmail.com")
out := outcome(plan, 100, 1830, 0)
out.Results[2].Err = errors.New("redis incr cnt/O/116681/d/2026-10-06/o: ERR value is not an integer or out of range")
l, buf := newLogger("text")
l.Outcome(out, []bool{true, false, false})
got := buf.String()
if !strings.Contains(got, ts+" <p1> out; tracked address counters for noah@wearemiraclesnow.com => count=100 limit=100 reached=true\n") {
t.Errorf("A should be tracked and reached:\n%s", got)
}
if !strings.Contains(got, ts+" <p1> out; redis error on key cnt/O/116681/d/2026-10-06/o: redis incr cnt/O/116681/d/2026-10-06/o: ERR value is not an integer or out of range (not counted)\n") {
t.Errorf("missing per-key error line:\n%s", got)
}
if strings.Contains(got, "tracked account counters") {
t.Errorf("failed key must not be reported as tracked:\n%s", got)
}
}
// A lost EXEC reply: the counts come from a master re-read and the block says so.
func TestOutcome_UnknownReRead(t *testing.T) {
plan, _ := planner().PlanOut("u1", "noah@wearemiraclesnow.com", "", "s@gmail.com")
out := outcome(plan, 100, 1830, 2200)
out.Unknown = errors.New("redis incr: read tcp 10.0.0.1:1->10.0.0.2:6379: i/o timeout")
l, buf := newLogger("text")
l.Outcome(out, []bool{true, false, true})
got := buf.String()
if !strings.Contains(got, ts+" <u1> out; redis outcome unknown: redis incr: read tcp 10.0.0.1:1->10.0.0.2:6379: i/o timeout, re-read from master\n") {
t.Errorf("missing unknown-outcome line:\n%s", got)
}
if !strings.Contains(got, "tracked address counters for noah@wearemiraclesnow.com => count=100 limit=100 reached=true (re-read)\n") {
t.Errorf("tracked lines should be marked re-read:\n%s", got)
}
if !strings.Contains(got, "limit reached (level=[A]") {
t.Errorf("crossing after re-read should still be reported:\n%s", got)
}
if strings.Contains(got, "nothing counted") {
t.Errorf("an unknown outcome is not 'nothing counted':\n%s", got)
}
l, buf = newLogger("json")
l.Outcome(out, nil)
if got := buf.String(); !strings.Contains(got, `"event":"redis_unknown"`) || !strings.Contains(got, `"reread":true`) {
t.Errorf("json unknown event missing:\n%s", got)
}
}
func TestDropped(t *testing.T) {
l, buf := newLogger("text")
l.Dropped(counting.RouteIn, "d1")
if got := buf.String(); got != ts+" <d1> in; dropped: queue full (not counted)\n" {
t.Errorf("got %q", got)
}
l, buf = newLogger("json")
l.Dropped(counting.RouteOut, "d2")
if got := buf.String(); !strings.Contains(got, `"event":"dropped"`) || !strings.Contains(got, `"route":"out"`) {
t.Errorf("got %q", got)
}
}
func TestTrail_MalformedRow(t *testing.T) {
bad := table{"good.com": {Name: "good.com", OrgID: 7, Limits: domains.ParseLimits("A/o=5")}}
p := &counting.Planner{Lookup: malformedTable{bad}, Now: func() time.Time { return clock }}
plan, err := p.PlanOut("mf", "a@good.com", "", "b@bad.com")
if err != nil {
t.Fatal(err)
}
l, buf := newLogger("text")
l.Outcome(outcome(plan, 1, 1, 1), nil)
if !strings.Contains(buf.String(), ts+" <mf> out; recipient domain row malformed in domains db: bad.com: "+`domain bad.com: invalid org id "x12"`+" (not counted)\n") {
t.Errorf("missing malformed line:\n%s", buf.String())
}
}
type malformedTable struct{ table }
func (m malformedTable) Lookup(d string) (*domains.Domain, error) {
if d == "bad.com" {
return domains.ParseRecord("bad.com", "x12|A/i=5")
}
return m.table.Lookup(d)
}
internal/watcher/watcher.go
+151
View File
@@ -0,0 +1,151 @@
// Package watcher reloads the domains CDB when the file changes on disk.
//
// It is the table mode of rcptd's watcher, trimmed: it watches the files
// themselves (kqueue on FreeBSD/macOS only reports in-place writes through a
// file watch) and their directories (an atomic `mv` over the file shows up as
// a Create in the directory), debounces bursts and calls one reload function.
package watcher
import (
"fmt"
"log/slog"
"os"
"path/filepath"
"sync"
"time"
"github.com/fsnotify/fsnotify"
)
// ReloadFunc is called, debounced, after a watched file changed. Its error is
// logged; the watcher keeps running either way.
type ReloadFunc func() error
// Watcher watches a set of files and calls reload on change.
type Watcher struct {
reload ReloadFunc
fsw *fsnotify.Watcher
files map[string]bool // absolute paths we care about
debounce time.Duration
done chan struct{}
stopOnce sync.Once
wg sync.WaitGroup
timerMu sync.Mutex
timer *time.Timer
runMu sync.Mutex // serialises reload calls
}
// New watches paths. Their directories must exist; the files themselves may
// appear later (the first deploy of the CDB triggers a reload).
func New(paths []string, reload ReloadFunc) (*Watcher, error) {
fsw, err := fsnotify.NewWatcher()
if err != nil {
return nil, err
}
w := &Watcher{
reload: reload,
fsw: fsw,
files: make(map[string]bool, len(paths)),
debounce: 500 * time.Millisecond,
done: make(chan struct{}),
}
dirs := map[string]bool{}
for _, p := range paths {
abs, err := filepath.Abs(p)
if err != nil {
fsw.Close()
return nil, fmt.Errorf("resolve path %s: %w", p, err)
}
w.files[abs] = true
if dir := filepath.Dir(abs); !dirs[dir] {
if err := fsw.Add(dir); err != nil {
fsw.Close()
return nil, fmt.Errorf("watch directory %s: %w", dir, err)
}
dirs[dir] = true
}
}
for abs := range w.files {
if _, err := os.Stat(abs); err == nil {
_ = fsw.Add(abs)
}
}
return w, nil
}
// Start begins watching in a background goroutine.
func (w *Watcher) Start() {
w.wg.Add(1)
go func() {
defer w.wg.Done()
for {
select {
case event, ok := <-w.fsw.Events:
if !ok {
return
}
if event.Op&(fsnotify.Write|fsnotify.Create) == 0 {
continue
}
abs, err := filepath.Abs(event.Name)
if err != nil || !w.files[abs] {
continue
}
// After a rename-over the old file watch is stale: re-add it so
// in-place writes to the new inode are seen too.
if event.Op&fsnotify.Create != 0 {
_ = w.fsw.Add(abs)
}
w.schedule()
case err, ok := <-w.fsw.Errors:
if !ok {
return
}
slog.Error("file watcher error", "error", err)
case <-w.done:
return
}
}
}()
}
// schedule (re)arms the debounce timer.
func (w *Watcher) schedule() {
w.timerMu.Lock()
defer w.timerMu.Unlock()
if w.timer != nil {
w.timer.Stop()
}
w.timer = time.AfterFunc(w.debounce, w.run)
}
func (w *Watcher) run() {
select {
case <-w.done:
return
default:
}
w.runMu.Lock()
defer w.runMu.Unlock()
if err := w.reload(); err != nil {
slog.Error("reload after file change failed", "error", err)
return
}
slog.Info("reloaded after file change")
}
// Stop ends watching and waits for the goroutine; safe to call twice.
func (w *Watcher) Stop() {
w.stopOnce.Do(func() {
w.timerMu.Lock()
if w.timer != nil {
w.timer.Stop()
}
w.timerMu.Unlock()
close(w.done)
w.fsw.Close()
w.wg.Wait()
})
}
internal/watcher/watcher_test.go
+189
View File
@@ -0,0 +1,189 @@
package watcher
import (
"errors"
"os"
"path/filepath"
"sync/atomic"
"testing"
"time"
)
func write(t *testing.T, path, content string) {
t.Helper()
if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
t.Fatal(err)
}
}
// replace mimics the agent's atomic deploy: write a temp file next to the
// target and rename it over the target.
func replace(t *testing.T, path, content string) {
t.Helper()
tmp := path + ".tmp"
write(t, tmp, content)
if err := os.Rename(tmp, path); err != nil {
t.Fatal(err)
}
}
type counter struct {
n atomic.Int32
done chan struct{}
err error
}
func newCounter() *counter { return &counter{done: make(chan struct{}, 16)} }
func (c *counter) reload() error {
c.n.Add(1)
select {
case c.done <- struct{}{}:
default:
}
return c.err
}
func (c *counter) wait(t *testing.T) {
t.Helper()
select {
case <-c.done:
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for reload")
}
}
// start creates and starts a watcher; debounce must be set before Start so
// the watcher goroutine never races with the test.
func start(t *testing.T, path string, c *counter, debounce time.Duration) *Watcher {
t.Helper()
w, err := New([]string{path}, c.reload)
if err != nil {
t.Fatal(err)
}
w.debounce = debounce
w.Start()
t.Cleanup(w.Stop)
time.Sleep(50 * time.Millisecond) // let the watches settle
return w
}
func TestWatcher_ReloadsOnInPlaceWrite(t *testing.T) {
path := filepath.Join(t.TempDir(), "rspamd_domains_db")
write(t, path, "v1")
c := newCounter()
start(t, path, c, 100*time.Millisecond)
write(t, path, "v2")
c.wait(t)
if c.n.Load() < 1 {
t.Errorf("reloads = %d, want >= 1", c.n.Load())
}
}
// The agent deploys the CDB with mv; after the first replacement the old
// kqueue file watch is stale, so a second replacement must still be seen.
func TestWatcher_ReloadsOnAtomicReplaceTwice(t *testing.T) {
path := filepath.Join(t.TempDir(), "rspamd_domains_db")
write(t, path, "v1")
c := newCounter()
start(t, path, c, 100*time.Millisecond)
replace(t, path, "v2")
c.wait(t)
first := c.n.Load()
time.Sleep(150 * time.Millisecond)
replace(t, path, "v3")
c.wait(t)
if c.n.Load() <= first {
t.Errorf("second replacement not detected: reloads %d -> %d", first, c.n.Load())
}
}
func TestWatcher_IgnoresOtherFilesInDir(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "rspamd_domains_db")
write(t, path, "v1")
c := newCounter()
start(t, path, c, 100*time.Millisecond)
write(t, filepath.Join(dir, "other_file"), "x")
time.Sleep(400 * time.Millisecond)
if c.n.Load() != 0 {
t.Errorf("reloads = %d for an unrelated file, want 0", c.n.Load())
}
}
func TestWatcher_DebouncesBurst(t *testing.T) {
path := filepath.Join(t.TempDir(), "rspamd_domains_db")
write(t, path, "v1")
c := newCounter()
start(t, path, c, 200*time.Millisecond)
for i := 0; i < 5; i++ {
write(t, path, "burst")
time.Sleep(20 * time.Millisecond)
}
c.wait(t)
time.Sleep(500 * time.Millisecond)
if n := c.n.Load(); n > 2 {
t.Errorf("reloads = %d, want 1-2 (debounced)", n)
}
}
func TestWatcher_ContinuesAfterReloadError(t *testing.T) {
path := filepath.Join(t.TempDir(), "rspamd_domains_db")
write(t, path, "v1")
c := newCounter()
c.err = errors.New("boom")
start(t, path, c, 100*time.Millisecond)
write(t, path, "v2")
c.wait(t)
time.Sleep(150 * time.Millisecond)
write(t, path, "v3")
c.wait(t)
if c.n.Load() < 2 {
t.Errorf("reloads = %d, want >= 2 (a failing reload must not stop the watcher)", c.n.Load())
}
}
func TestWatcher_StopPreventsReload(t *testing.T) {
path := filepath.Join(t.TempDir(), "rspamd_domains_db")
write(t, path, "v1")
c := newCounter()
w := start(t, path, c, 100*time.Millisecond)
w.Stop()
write(t, path, "v2")
time.Sleep(400 * time.Millisecond)
if c.n.Load() != 0 {
t.Errorf("reloads = %d after Stop, want 0", c.n.Load())
}
w.Stop() // idempotent
}
func TestNew_MissingDirectory(t *testing.T) {
if _, err := New([]string{"/nonexistent-dir-for-test/rspamd_domains_db"}, func() error { return nil }); err == nil {
t.Fatal("expected error for a missing directory")
}
}
func TestNew_MissingFileInExistingDirIsFine(t *testing.T) {
// The CDB may be deployed after the daemon starts; the directory watch
// picks up its creation.
path := filepath.Join(t.TempDir(), "rspamd_domains_db")
c := newCounter()
w, err := New([]string{path}, c.reload)
if err != nil {
t.Fatal(err)
}
w.debounce = 100 * time.Millisecond
w.Start()
defer w.Stop()
time.Sleep(50 * time.Millisecond)
write(t, path, "v1")
c.wait(t)
}