package server import ( "bytes" "context" "encoding/json" "fmt" "io" "log/slog" "net" "net/http" "net/http/httptest" "net/netip" "strconv" "strings" "sync" "sync/atomic" "testing" "time" "github.com/alicebob/miniredis/v2" "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/testutil" "github.com/redis/go-redis/v9" "migadu_counters/internal/api" "migadu_counters/internal/config" "migadu_counters/internal/counting" "migadu_counters/internal/domains" "migadu_counters/internal/metrics" "migadu_counters/internal/notify" "migadu_counters/internal/store" "migadu_counters/internal/txlog" ) // Real CDB: yagolivas.com = org 1049, O/i=200,O/o=20,D/i=200,D/o=20,A/i=200,A/o=20. const ( testCDB = "../../test_files/rspamd_domains_db.cdb" token = "secret-token" day = "2026-10-13" ) var clock = time.Date(2026, 10, 13, 12, 0, 0, 0, time.UTC) type syncBuffer struct { mu sync.Mutex buf bytes.Buffer } func (b *syncBuffer) Write(p []byte) (int, error) { b.mu.Lock() defer b.mu.Unlock() return b.buf.Write(p) } func (b *syncBuffer) String() string { b.mu.Lock() defer b.mu.Unlock() return b.buf.String() } // capture records notification requests. type capture struct { mu sync.Mutex urls []string } func (c *capture) ServeHTTP(w http.ResponseWriter, r *http.Request) { c.mu.Lock() c.urls = append(c.urls, r.URL.String()) c.mu.Unlock() w.WriteHeader(http.StatusOK) } func (c *capture) count() int { c.mu.Lock() defer c.mu.Unlock() return len(c.urls) } type options struct { token string cidrs []netip.Prefix handlerTimeout time.Duration writeAddr string // override the write Redis (e.g. a black hole) domainURL bool exempt map[string]bool queueSize int workers int } type harness struct { t *testing.T api *httptest.Server admin *httptest.Server write *miniredis.Miniredis read *miniredis.Miniredis st *store.Store db atomic.Pointer[domains.DB] lookup atomic.Pointer[counting.LookupFunc] notifier *notify.Notifier addrURL *capture orgURL *capture domURL *capture log *syncBuffer tx *syncBuffer reg *prometheus.Registry m *metrics.Metrics token string srv *Server } func newHarness(t *testing.T, o options) *harness { t.Helper() if o.token == "" { o.token = token } if o.handlerTimeout == 0 { o.handlerTimeout = 4 * time.Second } if o.exempt == nil { o.exempt = map[string]bool{"support@migadu.com": true, "contact@migadu.com": true} } h := &harness{t: t, token: o.token, log: &syncBuffer{}, tx: &syncBuffer{}} h.write, h.read = miniredis.RunT(t), miniredis.RunT(t) writeAddr := h.write.Addr() if o.writeAddr != "" { writeAddr = o.writeAddr } var err error h.st, err = store.New(config.RedisConfig{WriteAddr: writeAddr, ReadAddrs: []string{h.read.Addr()}, PoolSize: 4, TimeoutMS: 2000}) if err != nil { t.Fatal(err) } t.Cleanup(func() { h.st.Close() }) db, err := domains.Open(testCDB) if err != nil { t.Fatal(err) } t.Cleanup(func() { db.Close() }) h.db.Store(db) defaultLookup := counting.LookupFunc(func(d string) (*domains.Domain, error) { return h.db.Load().Lookup(d) }) h.lookup.Store(&defaultLookup) h.addrURL, h.orgURL, h.domURL = &capture{}, &capture{}, &capture{} sa, so, sd := httptest.NewServer(h.addrURL), httptest.NewServer(h.orgURL), httptest.NewServer(h.domURL) t.Cleanup(sa.Close) t.Cleanup(so.Close) t.Cleanup(sd.Close) q := "/api/limits?level=$LEVEL&id=$ID&flow=$FLOW&period=$PERIOD&by=$ADDRESS&count=$COUNT" ncfg := notify.Config{AddressURL: sa.URL + q, OrgURL: so.URL + q, User: "u", Password: "p", Timeout: time.Second, Attempts: 2, Backoff: time.Millisecond} if o.domainURL { ncfg.DomainURL = sd.URL + q } h.reg = prometheus.NewRegistry() h.m = metrics.New(h.reg) tx := txlog.New(h.tx, "text", func() time.Time { return clock }) h.notifier = notify.New(ncfg, NotifyHooks(tx, h.m)) t.Cleanup(func() { h.notifier.Shutdown(context.Background()) }) logger := slog.New(slog.NewTextHandler(h.log, &slog.HandlerOptions{Level: slog.LevelDebug})) planner := &counting.Planner{ Lookup: counting.LookupFunc(func(d string) (*domains.Domain, error) { return (*h.lookup.Load())(d) }), Exempt: o.exempt, Now: func() time.Time { return clock }, } s := New(Deps{ AuthToken: o.token, AllowedCIDRs: o.cidrs, HandlerTimeout: o.handlerTimeout, QueueSize: o.queueSize, Workers: o.workers, Exempt: o.exempt, Counter: &counting.Counter{Planner: planner, Store: h.st, ReReadTimeout: 200 * time.Millisecond}, Store: h.st, Domains: func() *domains.DB { return h.db.Load() }, Notifier: h.notifier, TxLog: tx, Metrics: h.m, Gatherer: h.reg, Logger: logger, Now: func() time.Time { return clock }, }) h.srv = s h.api = httptest.NewServer(s.Handler()) h.admin = httptest.NewServer(s.AdminHandler()) t.Cleanup(h.api.Close) t.Cleanup(h.admin.Close) t.Cleanup(func() { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _ = s.Drain(ctx) }) return h } // settle waits until every accepted job has been processed. func (h *harness) settle() { h.t.Helper() waitFor(h.t, func() bool { return h.srv.Pending() == 0 }, "queue did not settle") } func waitFor(t *testing.T, cond func() bool, msg string) { t.Helper() deadline := time.Now().Add(5 * time.Second) for !cond() { if time.Now().After(deadline) { t.Fatal(msg) } time.Sleep(2 * time.Millisecond) } } // accept posts a body that must be accepted (202) and waits for it to be counted. func (h *harness) accept(path, body string) { h.t.Helper() st, b := h.post(path, body) if st != 202 || strings.TrimSpace(b) != `{"accepted":true}` { h.t.Fatalf("%s %s: %d %s, want 202 accepted", path, body, st, b) } h.settle() } func (h *harness) redisInt(key string) int64 { h.t.Helper() v, err := h.write.Get(key) if err != nil { return 0 } n, err := strconv.ParseInt(v, 10, 64) if err != nil { h.t.Fatalf("key %s holds %q", key, v) } return n } // rawPostAndClose is a fire-and-forget client: it writes the request and // closes the socket without reading the reply. func rawPostAndClose(t *testing.T, addr, path, body string) { t.Helper() c, err := net.Dial("tcp", addr) if err != nil { t.Fatal(err) } req := fmt.Sprintf("POST %s HTTP/1.1\r\nHost: x\r\nAuthorization: Bearer %s\r\nContent-Type: application/json\r\nContent-Length: %d\r\n\r\n%s", path, token, len(body), body) if _, err := c.Write([]byte(req)); err != nil { t.Fatal(err) } c.Close() } // do sends a request with the default headers unless overridden; a header // value of "-" removes it. func (h *harness) do(method, path, body string, hdr map[string]string) (int, http.Header, string) { h.t.Helper() req, err := http.NewRequest(method, h.api.URL+path, strings.NewReader(body)) if err != nil { h.t.Fatal(err) } req.Header.Set("Authorization", "Bearer "+h.token) if body != "" { req.Header.Set("Content-Type", "application/json") } for k, v := range hdr { if v == "-" { req.Header.Del(k) } else { req.Header.Set(k, v) } } resp, err := http.DefaultClient.Do(req) if err != nil { h.t.Fatal(err) } defer resp.Body.Close() b, _ := io.ReadAll(resp.Body) return resp.StatusCode, resp.Header, string(b) } func (h *harness) post(path, body string) (int, string) { st, _, b := h.do(http.MethodPost, path, body, nil) return st, b } func (h *harness) get(path string) (int, string) { st, _, b := h.do(http.MethodGet, path, "", nil) return st, b } func (h *harness) metric(c prometheus.Collector) float64 { return testutil.ToFloat64(c) } func TestOut_CountsSenderAndExternalRecipient(t *testing.T) { h := newHarness(t, options{}) h.accept("/out", `{"trace_id":"tr1","auth_user":"a@yagolivas.com","mail_from":"a@yagolivas.com","recipient":"ext@gmail.com"}`) for _, k := range []string{store.Key("A", "a@yagolivas.com", day, "o"), store.Key("D", "yagolivas.com", day, "o"), store.Key("O", "1049", day, "o")} { if h.redisInt(k) != 1 { t.Errorf("redis %s = %d, want 1", k, h.redisInt(k)) } } if keys := h.write.Keys(); len(keys) != 3 { t.Errorf("redis keys = %v, want the 3 outgoing keys only", keys) } if got := h.tx.String(); !strings.Contains(got, " out; tracked address counters for a@yagolivas.com => count=1 limit=20 reached=false") || !strings.Contains(got, "recipient domain not in domains db: gmail.com (not counted)") { t.Errorf("transaction log:\n%s", got) } if h.metric(h.m.Requests.WithLabelValues("/out", "202")) != 1 || h.metric(h.m.Increments.WithLabelValues("A", "o")) != 1 { t.Error("metrics not recorded") } } func TestOut_LocalRecipientCountsIncomingToo(t *testing.T) { h := newHarness(t, options{}) h.accept("/out", `{"trace_id":"t","auth_user":"a@yagolivas.com","recipient":"b@yagolivas.com"}`) if keys := h.write.Keys(); len(keys) != 6 { t.Errorf("redis keys = %v, want 6", keys) } if h.redisInt(store.Key("A", "b@yagolivas.com", day, "i")) != 1 || h.redisInt(store.Key("O", "1049", day, "i")) != 1 { t.Errorf("recipient incoming not counted: %v", h.write.Keys()) } } func TestOut_ReachedNotifiesConfiguredLevelsOnce(t *testing.T) { h := newHarness(t, options{}) for i := 1; i <= 21; i++ { h.accept("/out", `{"trace_id":"t","auth_user":"a@yagolivas.com","recipient":"ext@gmail.com"}`) } if err := h.notifier.Shutdown(context.Background()); err != nil { t.Fatal(err) } if h.addrURL.count() != 1 || h.orgURL.count() != 1 || h.domURL.count() != 0 { t.Errorf("notifications addr=%d org=%d dom=%d, want 1/1/0", h.addrURL.count(), h.orgURL.count(), h.domURL.count()) } if u := h.addrURL.urls[0]; !strings.Contains(u, "level=A") || !strings.Contains(u, "id=a%40yagolivas.com") || !strings.Contains(u, "count=20") || !strings.Contains(u, "by=a%40yagolivas.com") { t.Errorf("address notification url = %s", u) } tx := h.tx.String() for _, want := range []string{ "tracked address counters for a@yagolivas.com => count=20 limit=20 reached=true", "tracked address counters for a@yagolivas.com => count=21 limit=20 reached=false", "limit reached (level=[A] id=[a@yagolivas.com] flow=[o] by=[a@yagolivas.com]) => notifying address_url", "limit reached (level=[O] id=[1049] flow=[o] by=[a@yagolivas.com]) => notifying org_url", "limit reached (level=[D] id=[yagolivas.com] flow=[o] by=[a@yagolivas.com]) => no domain_url configured, not notifying", "sent limit exceeded notification to url", } { if !strings.Contains(tx, want) { t.Errorf("transaction log lacks %q:\n%s", want, tx) } } if h.metric(h.m.LimitReached.WithLabelValues("A", "o")) != 1 || h.metric(h.m.Notifications.WithLabelValues("A", "sent")) != 1 || h.metric(h.m.Notifications.WithLabelValues("D", "disabled")) != 1 { t.Error("limit/notification metrics not recorded") } } func TestIn_CountsRecipientOnly(t *testing.T) { h := newHarness(t, options{}) h.accept("/in", `{"trace_id":"t","mail_from":"s@gmail.com","recipient":"b@yagolivas.com","junk":false}`) if keys := h.write.Keys(); len(keys) != 3 { t.Errorf("redis keys = %v", keys) } if h.redisInt(store.Key("O", "1049", day, "i")) != 1 { t.Error("O/i not counted") } if !strings.Contains(h.tx.String(), " in; tracked account counters for 1049 => count=1 limit=200 reached=false") { t.Errorf("transaction log:\n%s", h.tx.String()) } } func TestSkips(t *testing.T) { h := newHarness(t, options{}) cases := []struct{ path, body, skipped string }{ {"/in", `{"trace_id":"s1","mail_from":"s@gmail.com","recipient":"b@yagolivas.com","junk":true}`, "junk"}, {"/in", `{"trace_id":"s2","mail_from":"s@gmail.com","recipient":"nobody@gmail.com"}`, "unknown_domain"}, {"/in", `{"trace_id":"s3","mail_from":"s@gmail.com","recipient":"support@migadu.com"}`, "exempt"}, {"/out", `{"trace_id":"s4","auth_user":"contact@migadu.com","recipient":"b@yagolivas.com"}`, "exempt"}, {"/out", `{"trace_id":"s5","auth_user":"a@unknown.tld","recipient":"b@gmail.com"}`, "unknown_domain"}, } for i, c := range cases { h.accept(c.path, c.body) if !strings.Contains(h.tx.String(), fmt.Sprintf(" %s; skipped: %s", i+1, strings.TrimPrefix(c.path, "/"), c.skipped)) { t.Errorf("%s %s: transaction log lacks the skip line:\n%s", c.path, c.body, h.tx.String()) } } if keys := h.write.Keys(); len(keys) != 0 { t.Errorf("skipped requests wrote keys: %v", keys) } } // --- review round 1: counting is decoupled from the request (#1) ------------- // A fire-and-forget client that closes the socket before the reply must not // cancel the count (reviewer A: 20 of 20 lost before the fix). func TestClientDisconnectStillCounts(t *testing.T) { h := newHarness(t, options{}) real := *h.lookup.Load() slow := counting.LookupFunc(func(d string) (*domains.Domain, error) { time.Sleep(20 * time.Millisecond) return real(d) }) h.lookup.Store(&slow) addr := strings.TrimPrefix(h.api.URL, "http://") const n = 20 for i := 0; i < n; i++ { rawPostAndClose(t, addr, "/in", `{"trace_id":"dc","mail_from":"s@gmail.com","recipient":"a@yagolivas.com"}`) } key := store.Key("A", "a@yagolivas.com", day, "i") waitFor(t, func() bool { return h.redisInt(key) == n }, "disconnected clients' requests were not all counted") h.lookup.Store(&real) for i := 0; i < 300; i++ { rawPostAndClose(t, addr, "/in", `{"trace_id":"dc","mail_from":"s@gmail.com","recipient":"b@yagolivas.com"}`) } key = store.Key("A", "b@yagolivas.com", day, "i") waitFor(t, func() bool { return h.redisInt(key) == 300 }, fmt.Sprintf("immediate disconnects: counter = %d, want 300", h.redisInt(key))) } // The reply arrives before the Redis work: even a stalled Redis does not // delay the 202. func TestAcceptedBeforeCounting(t *testing.T) { ln, err := net.Listen("tcp", "127.0.0.1:0") // black hole: accepts, never answers if err != nil { t.Fatal(err) } defer ln.Close() go func() { for { c, err := ln.Accept() if err != nil { return } defer c.Close() } }() h := newHarness(t, options{writeAddr: ln.Addr().String(), handlerTimeout: 300 * time.Millisecond}) start := time.Now() st, body := h.post("/in", `{"trace_id":"bh","recipient":"b@yagolivas.com"}`) if st != 202 || time.Since(start) > 200*time.Millisecond { t.Errorf("status %d %s after %v, want an immediate 202", st, body, time.Since(start)) } h.settle() if tx := h.tx.String(); !strings.Contains(tx, " in; redis error:") || !strings.Contains(tx, "nothing counted") { t.Errorf("worker should log the Redis failure:\n%s", tx) } if h.metric(h.m.RedisErrors.WithLabelValues("incr")) < 1 { t.Error("redis error metric not recorded") } } // A full queue answers 503 and the drop is visible in metrics and the log. func TestQueueFull(t *testing.T) { release := make(chan struct{}) h := newHarness(t, options{queueSize: 1, workers: 1}) real := *h.lookup.Load() blocking := counting.LookupFunc(func(d string) (*domains.Domain, error) { <-release return real(d) }) h.lookup.Store(&blocking) // First job occupies the worker, second fills the queue. for i := 0; i < 2; i++ { if st, body := h.post("/in", `{"trace_id":"q","recipient":"b@yagolivas.com"}`); st != 202 { t.Fatalf("request %d: %d %s", i, st, body) } } waitFor(t, func() bool { return h.srv.Depth() == 1 }, "worker did not take the first job") st, body := h.post("/in", `{"trace_id":"full","recipient":"b@yagolivas.com"}`) if st != 503 || strings.TrimSpace(body) != `{"error":"queue full"}` { t.Errorf("queue full = %d %s", st, body) } if h.metric(h.m.Dropped.WithLabelValues("/in")) != 1 || h.metric(h.m.Requests.WithLabelValues("/in", "503")) != 1 { t.Error("dropped / 503 metrics not recorded") } if !strings.Contains(h.tx.String(), " in; dropped: queue full (not counted)") { t.Errorf("transaction log:\n%s", h.tx.String()) } close(release) h.settle() if h.redisInt(store.Key("A", "b@yagolivas.com", day, "i")) != 2 { t.Errorf("accepted jobs should still be counted: %v", h.write.Keys()) } } // Drain finishes every accepted job before returning. func TestDrainCountsEverything(t *testing.T) { h := newHarness(t, options{workers: 1}) real := *h.lookup.Load() slow := counting.LookupFunc(func(d string) (*domains.Domain, error) { time.Sleep(10 * time.Millisecond) return real(d) }) h.lookup.Store(&slow) for i := 0; i < 10; i++ { if st, _ := h.post("/in", `{"recipient":"b@yagolivas.com"}`); st != 202 { t.Fatal(st) } } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if err := h.srv.Drain(ctx); err != nil { t.Fatal(err) } if h.redisInt(store.Key("A", "b@yagolivas.com", day, "i")) != 10 { t.Errorf("after Drain counter = %d, want 10", h.redisInt(store.Key("A", "b@yagolivas.com", day, "i"))) } if st, body := h.post("/in", `{"recipient":"b@yagolivas.com"}`); st != 503 { t.Errorf("after Drain = %d %s, want 503 (not accepting)", st, body) } } // --- review round 1: only counted fields are validated (#2), token logins (#3) // Reviewer B: every odd-but-accepted-by-mizu mail_from used to 400 the // request and lose the recipient's count. func TestIn_OddMailFromStillCountsRecipient(t *testing.T) { h := newHarness(t, options{}) senders := []string{ "john smith@example.com", "a@b@example.com", "info@bücher.de", "postmaster@[192.0.2.1]", "bounce@mail_relay.example.com", strings.Repeat("v", 250) + "@example.com", } for _, mf := range senders { h.accept("/in", `{"trace_id":"odd","mail_from":"`+mf+`","recipient":"b@yagolivas.com"}`) } if got := h.redisInt(store.Key("A", "b@yagolivas.com", day, "i")); got != int64(len(senders)) { t.Errorf("recipient counted %d times, want %d", got, len(senders)) } if !strings.Contains(h.tx.String(), "request mail_from=john smith@example.com recipient=b@yagolivas.com") || !strings.Contains(h.tx.String(), "request mail_from=info@xn--bcher-kva.de") { t.Errorf("transaction log should show the raw / punycode sender:\n%s", h.tx.String()) } } // Reviewer B: an unusual external recipient used to drop the sender's count. func TestOut_OddRecipientStillCountsSender(t *testing.T) { h := newHarness(t, options{}) for _, rc := range []string{"x@[192.0.2.1]", "x@host_name.example.com", "", "nobody"} { h.accept("/out", `{"trace_id":"odd","auth_user":"a@yagolivas.com","mail_from":"a@yagolivas.com","recipient":"`+rc+`"}`) } if got := h.redisInt(store.Key("A", "a@yagolivas.com", day, "o")); got != 4 { t.Errorf("sender outgoing = %d, want 4", got) } tx := h.tx.String() for _, want := range []string{ "recipient address not usable: x@[192.0.2.1] (invalid e-mail address, not counted)", "recipient address not usable: x@host_name.example.com (invalid e-mail address, not counted)", "recipient address not usable: (missing, not counted)", "recipient address not usable: nobody (invalid e-mail address, not counted)", } { if !strings.Contains(tx, want) { t.Errorf("transaction log lacks %q:\n%s", want, tx) } } // IDN recipients are counted under their punycode key when hosted. h.accept("/out", `{"trace_id":"idn","auth_user":"a@yagolivas.com","recipient":"x@Bücher.de"}`) if !strings.Contains(h.tx.String(), " out; recipient domain not in domains db: xn--bcher-kva.de (not counted)") { t.Errorf("IDN recipient should be looked up as punycode:\n%s", h.tx.String()) } } // Reviewer A: master/support logins arrive as user@dom@SUFFIX. func TestOut_LoginSuffixCountsBaseAddress(t *testing.T) { h := newHarness(t, options{}) h.accept("/out", `{"trace_id":"m1","auth_user":"x@yagolivas.com@SUPPORTTOKEN","mail_from":"x@yagolivas.com","recipient":"y@yagolivas.com"}`) if h.redisInt(store.Key("A", "x@yagolivas.com", day, "o")) != 1 || h.redisInt(store.Key("A", "y@yagolivas.com", day, "i")) != 1 { t.Errorf("base address not counted: %v", h.write.Keys()) } if !strings.Contains(h.tx.String(), " out; request auth_user=x@yagolivas.com recipient=y@yagolivas.com mail_from=x@yagolivas.com login_suffix=SUPPORTTOKEN") { t.Errorf("transaction log:\n%s", h.tx.String()) } } // --- review round 1: honest handling of ambiguous Redis results (#4) --------- // Reviewer A: one key holding garbage used to fail the whole request and lose // the crossing on the other keys. func TestPartialExecFailureStillNotifies(t *testing.T) { h := newHarness(t, options{}) h.write.Set(store.Key("O", "1049", day, "o"), "garbage") h.write.Set(store.Key("A", "x@yagolivas.com", day, "o"), "19") // A/o=20: this request crosses h.accept("/out", `{"trace_id":"p1","auth_user":"x@yagolivas.com","recipient":"z@gmail.com"}`) if h.redisInt(store.Key("A", "x@yagolivas.com", day, "o")) != 20 || h.redisInt(store.Key("D", "yagolivas.com", day, "o")) != 1 { t.Errorf("good keys not counted: %v", h.write.Keys()) } if err := h.notifier.Shutdown(context.Background()); err != nil { t.Fatal(err) } if h.addrURL.count() != 1 { t.Errorf("address notifications = %d, want 1", h.addrURL.count()) } tx := h.tx.String() if !strings.Contains(tx, " out; tracked address counters for x@yagolivas.com => count=20 limit=20 reached=true") || !strings.Contains(tx, " out; redis error on key cnt/O/1049/d/"+day+"/o: ") || strings.Contains(tx, "nothing counted") { t.Errorf("transaction log:\n%s", tx) } if h.metric(h.m.RedisErrors.WithLabelValues("incr")) != 1 { t.Errorf("redis_errors_total{incr} = %v, want 1 (the failed key)", h.metric(h.m.RedisErrors.WithLabelValues("incr"))) } } // delayProxy forwards to backend and, once it has seen EXEC, holds back the // next reply longer than the client's read timeout. func delayProxy(t *testing.T, backend string, 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(1) 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() } // Reviewer B: a lost EXEC reply used to be logged "nothing counted" although // the counters moved, so the crossing was never notified that day. func TestLostReplyStillNotifies(t *testing.T) { backend := miniredis.RunT(t) proxy := delayProxy(t, backend.Addr(), 3*time.Second) h := newHarness(t, options{writeAddr: proxy}) key := store.Key("A", "a@yagolivas.com", day, "o") // A/o=20 backend.Set(key, "19") h.accept("/out", `{"trace_id":"lost","auth_user":"a@yagolivas.com","recipient":"x@gmail.com"}`) if v, _ := backend.Get(key); v != "20" { t.Fatalf("redis value = %s, the increment should have been applied", v) } if err := h.notifier.Shutdown(context.Background()); err != nil { t.Fatal(err) } if h.addrURL.count() != 1 { t.Errorf("address notifications = %d, want 1 after the master re-read", h.addrURL.count()) } tx := h.tx.String() if !strings.Contains(tx, " out; redis outcome unknown: ") || !strings.Contains(tx, "re-read from master") || !strings.Contains(tx, "tracked address counters for a@yagolivas.com => count=20 limit=20 reached=true (re-read)") { t.Errorf("transaction log:\n%s", tx) } if h.metric(h.m.RedisErrors.WithLabelValues("incr_unknown")) != 1 { t.Error("redis_errors_total{incr_unknown} not recorded") } } // --- HTTP validation layer, one test per step ------------------------------- func TestValidation_Method(t *testing.T) { h := newHarness(t, options{}) st, hdr, body := h.do(http.MethodGet, "/out", "", nil) if st != 405 || hdr.Get("Allow") != "POST" || !strings.Contains(body, `"error"`) { t.Errorf("GET /out = %d %v %s", st, hdr, body) } if st, _, _ := h.do(http.MethodPost, "/counters?address=a@yagolivas.com", `{}`, nil); st != 405 { t.Errorf("POST /counters = %d", st) } if st, _, _ := h.do(http.MethodPut, "/in", `{}`, nil); st != 405 { t.Errorf("PUT /in = %d", st) } } func TestValidation_BodyTooLarge(t *testing.T) { h := newHarness(t, options{}) big := `{"recipient":"b@yagolivas.com","pad":"` + strings.Repeat("x", 65*1024) + `"}` if st, body := h.post("/in", big); st != 413 { t.Errorf("65 KB body = %d %s", st, body) } } func TestValidation_BearerToken(t *testing.T) { h := newHarness(t, options{}) body := `{"mail_from":"s@gmail.com","recipient":"b@yagolivas.com"}` for name, hdr := range map[string]map[string]string{ "missing": {"Authorization": "-"}, "wrong": {"Authorization": "Bearer nope"}, "basic": {"Authorization": "Basic dTpw"}, "empty": {"Authorization": "Bearer "}, } { st, h2, b := h.do(http.MethodPost, "/in", body, hdr) if st != 401 || h2.Get("WWW-Authenticate") != "Bearer" { t.Errorf("%s: %d %v %s", name, st, h2, b) } } if st, _, _ := h.do(http.MethodGet, "/counters?address=b@yagolivas.com", "", map[string]string{"Authorization": "-"}); st != 401 { t.Errorf("GET /counters without token = %d", st) } if keys := h.write.Keys(); len(keys) != 0 { t.Errorf("rejected requests wrote keys: %v", keys) } if h.metric(h.m.Requests.WithLabelValues("/in", "401")) != 4 { t.Errorf("401s not counted: %v", h.metric(h.m.Requests.WithLabelValues("/in", "401"))) } } func TestValidation_SourceIP(t *testing.T) { denied := newHarness(t, options{cidrs: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}}) if st, body := denied.post("/in", `{"recipient":"b@yagolivas.com"}`); st != 403 { t.Errorf("from 127.0.0.1 with allowlist 10/8 = %d %s", st, body) } allowed := newHarness(t, options{cidrs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8"), netip.MustParsePrefix("::1/128")}}) if st, body := allowed.post("/in", `{"recipient":"b@yagolivas.com"}`); st != 202 { t.Errorf("from loopback with loopback allowlist = %d %s", st, body) } } func TestValidation_ContentType(t *testing.T) { h := newHarness(t, options{}) body := `{"recipient":"b@yagolivas.com"}` if st, _, _ := h.do(http.MethodPost, "/in", body, map[string]string{"Content-Type": "text/plain"}); st != 415 { t.Errorf("text/plain = %d", st) } if st, _, _ := h.do(http.MethodPost, "/in", body, map[string]string{"Content-Type": "-"}); st != 415 { t.Errorf("missing content type = %d", st) } if st, _, b := h.do(http.MethodPost, "/in", body, map[string]string{"Content-Type": "application/json; charset=utf-8"}); st != 202 { t.Errorf("json with charset = %d %s", st, b) } } func TestValidation_JSONAndFields(t *testing.T) { h := newHarness(t, options{}) cases := []struct { name, path, body, wantErr string }{ {"syntax", "/in", `{"recipient": `, "invalid json"}, {"two objects", "/in", `{"recipient":"b@yagolivas.com"}{}`, "exactly one JSON object"}, {"array", "/in", `[]`, "invalid json"}, {"missing recipient", "/in", `{"mail_from":"s@gmail.com"}`, "recipient: missing"}, {"bad recipient", "/in", `{"recipient":"nobody"}`, "recipient: invalid e-mail address"}, {"both sides missing", "/out", `{"mail_from":"a@yagolivas.com"}`, "auth_user: missing; recipient: missing"}, {"both sides bad", "/out", `{"auth_user":"nobody","recipient":"x@[192.0.2.1]"}`, "auth_user: invalid e-mail address; recipient: invalid e-mail address"}, {"wrong type", "/in", `{"recipient":"b@yagolivas.com","junk":"yes"}`, "invalid json"}, } for _, c := range cases { st, body := h.post(c.path, c.body) if st != 400 || !strings.Contains(body, c.wantErr) { t.Errorf("%s: %d %s (want 400 containing %q)", c.name, st, body, c.wantErr) } } // Unknown fields are accepted so mailqueuer can add fields later; a bad // mail_from and a missing auth_user never reject on their own (#2). for _, body := range []string{ `{"recipient":"b@yagolivas.com","future_field":1}`, `{"auth_user":"a@yagolivas.com","recipient":"b@yagolivas.com","mail_from":"x"}`, `{"recipient":"b@yagolivas.com"}`, } { if st, b := h.post("/out", body); st != 202 { t.Errorf("%s = %d %s, want 202", body, st, b) } } h.settle() if h.redisInt(store.Key("A", "b@yagolivas.com", day, "i")) != 3 || h.redisInt(store.Key("A", "a@yagolivas.com", day, "o")) != 1 { t.Errorf("only the accepted requests may write keys: %v", h.write.Keys()) } if !strings.Contains(h.log.String(), "recipient: missing") { t.Errorf("rejections should be logged:\n%s", h.log.String()) } } // The synchronous part is still bounded by handler_timeout: a validation // layer that stalls (here: a slow body) answers 503 instead of hanging. func TestValidation_HandlerTimeout(t *testing.T) { h := newHarness(t, options{handlerTimeout: 100 * time.Millisecond}) pr, pw := io.Pipe() go func() { pw.Write([]byte(`{"recipient":`)) time.Sleep(400 * time.Millisecond) pw.Write([]byte(`"b@yagolivas.com"}`)) pw.Close() }() req, _ := http.NewRequest(http.MethodPost, h.api.URL+"/in", pr) req.Header.Set("Authorization", "Bearer "+token) req.Header.Set("Content-Type", "application/json") start := time.Now() resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatal(err) } b, _ := io.ReadAll(resp.Body) resp.Body.Close() if resp.StatusCode != 503 || !strings.Contains(string(b), "timeout") || time.Since(start) > time.Second { t.Errorf("status %d %s after %v", resp.StatusCode, b, time.Since(start)) } } // A panic while counting (here: in the domain lookup) is recovered in the // worker: the request was already accepted, the panic is logged and the // worker keeps serving. func TestWorkerRecoversFromPanic(t *testing.T) { h := newHarness(t, options{workers: 1}) boom := counting.LookupFunc(func(string) (*domains.Domain, error) { panic("boom") }) prev := h.lookup.Swap(&boom) h.accept("/in", `{"trace_id":"pan","recipient":"b@yagolivas.com"}`) h.lookup.Store(prev) h.accept("/in", `{"recipient":"b@yagolivas.com"}`) if h.redisInt(store.Key("A", "b@yagolivas.com", day, "i")) != 1 { t.Error("worker must keep serving after a panic") } if !strings.Contains(h.log.String(), "panic") || !strings.Contains(h.log.String(), "pan") { t.Errorf("panic should be logged with the trace id:\n%s", h.log.String()) } } // A panic in the synchronous part is still answered 500 and the daemon lives. func TestValidation_RecoverFromPanic(t *testing.T) { h := newHarness(t, options{}) h.srv.panicHook = func() { panic("sync boom") } st, body := h.post("/in", `{"recipient":"b@yagolivas.com"}`) if st != 500 || !strings.Contains(body, "internal error") { t.Errorf("panic = %d %s", st, body) } h.srv.panicHook = nil h.accept("/in", `{"recipient":"b@yagolivas.com"}`) if !strings.Contains(h.log.String(), "panic") { t.Error("panic should be logged") } } func TestUnknownPath(t *testing.T) { h := newHarness(t, options{}) if st, body := h.get("/nope"); st != 404 || !strings.Contains(body, `"error"`) { t.Errorf("/nope = %d %s", st, body) } } func TestRedisDown_LoggedNotCounted(t *testing.T) { h := newHarness(t, options{}) h.write.Close() h.accept("/in", `{"trace_id":"tr","recipient":"b@yagolivas.com"}`) if !strings.Contains(h.tx.String(), " in; redis error:") || !strings.Contains(h.tx.String(), "(nothing counted)") { t.Errorf("transaction log should record the redis error:\n%s", h.tx.String()) } if h.metric(h.m.RedisErrors.WithLabelValues("incr")) != 1 { t.Error("redis error metric not recorded") } if !strings.Contains(h.log.String(), "counting failed") { t.Errorf("app log:\n%s", h.log.String()) } } func TestLookupError_Logged(t *testing.T) { h := newHarness(t, options{}) broken := counting.LookupFunc(func(string) (*domains.Domain, error) { return nil, io.ErrUnexpectedEOF }) h.lookup.Store(&broken) h.accept("/in", `{"trace_id":"tr","recipient":"b@yagolivas.com"}`) if !strings.Contains(h.tx.String(), " in; error: lookup recipient domain yagolivas.com") { t.Errorf("transaction log:\n%s", h.tx.String()) } } // --- /counters -------------------------------------------------------------- func decodeCounters(t *testing.T, body string) api.CountersResponse { t.Helper() var r api.CountersResponse if err := json.Unmarshal([]byte(body), &r); err != nil { t.Fatalf("bad JSON %q: %v", body, err) } return r } func TestCounters_AddressReadsReplica(t *testing.T) { h := newHarness(t, options{}) // Counters live on the replica (reads) — the write instance holds other values. h.read.Set(store.Key("A", "a@yagolivas.com", day, "i"), "12") h.read.Set(store.Key("A", "a@yagolivas.com", day, "o"), "3") h.read.Set(store.Key("D", "yagolivas.com", day, "i"), "40") h.read.Set(store.Key("O", "1049", day, "o"), "30") h.write.Set(store.Key("A", "a@yagolivas.com", day, "i"), "999") st, body := h.get("/counters?address=A@Yagolivas.com") if st != 200 { t.Fatalf("status %d: %s", st, body) } r := decodeCounters(t, body) if r.Query != "a@yagolivas.com" || r.Mode != "address" || r.Date != day || !r.Managed || r.Exempt || r.Domain != "yagolivas.com" || r.OrgID == nil || *r.OrgID != 1049 { t.Errorf("header = %+v", r) } if len(r.Levels) != 3 { t.Fatalf("levels = %+v", r.Levels) } a, d, o := r.Levels[0], r.Levels[1], r.Levels[2] if a.Level != "A" || a.ID != "a@yagolivas.com" || a.In != 12 || *a.InLimit != 200 || a.Out != 3 || *a.OutLimit != 20 { t.Errorf("A = %+v", a) } if d.Level != "D" || d.ID != "yagolivas.com" || d.In != 40 || d.Out != 0 || *d.InLimit != 200 { t.Errorf("D = %+v", d) } if o.Level != "O" || o.ID != "1049" || o.Out != 30 || *o.OutLimit != 20 || *o.InLimit != 200 { t.Errorf("O = %+v", o) } } func TestCounters_DomainMode(t *testing.T) { h := newHarness(t, options{}) st, body := h.get("/counters?domain=yagolivas.com") if st != 200 { t.Fatalf("status %d: %s", st, body) } r := decodeCounters(t, body) if r.Mode != "domain" || len(r.Levels) != 2 || r.Levels[0].Level != "D" || r.Levels[1].Level != "O" || !r.Managed { t.Errorf("response = %+v", r) } } func TestCounters_DateParam(t *testing.T) { h := newHarness(t, options{}) h.read.Set(store.Key("A", "a@yagolivas.com", "2026-10-01", "i"), "7") st, body := h.get("/counters?address=a@yagolivas.com&date=2026-10-01") if st != 200 { t.Fatalf("status %d: %s", st, body) } if r := decodeCounters(t, body); r.Date != "2026-10-01" || r.Levels[0].In != 7 { t.Errorf("response = %+v", r) } for _, q := range []string{"address=a@yagolivas.com&date=2026-10-14", "address=a@yagolivas.com&date=13.10.2026", "address=a@yagolivas.com&domain=x.com", "", "address=nobody", "domain=yagolivas.com&date=2027-01-01"} { if st, body := h.get("/counters?" + q); st != 400 { t.Errorf("%q = %d %s, want 400", q, st, body) } } } func TestCounters_UnmanagedAndExempt(t *testing.T) { h := newHarness(t, options{}) st, body := h.get("/counters?address=x@gmail.com") if st != 200 { t.Fatalf("status %d: %s", st, body) } r := decodeCounters(t, body) if r.Managed || r.OrgID != nil || len(r.Levels) != 2 || r.Levels[0].InLimit != nil || r.Levels[1].Level != "D" { t.Errorf("unmanaged response = %+v", r) } st, body = h.get("/counters?address=support@migadu.com") if st != 200 { t.Fatalf("status %d: %s", st, body) } if r := decodeCounters(t, body); !r.Exempt || !r.Managed { t.Errorf("exempt response = %+v", r) } } func TestCounters_RedisDown(t *testing.T) { h := newHarness(t, options{}) h.read.Close() if st, body := h.get("/counters?address=a@yagolivas.com"); st != 503 { t.Errorf("status %d %s", st, body) } } // --- admin ------------------------------------------------------------------- func (h *harness) adminGet(path string) (int, string) { resp, err := http.Get(h.admin.URL + path) if err != nil { h.t.Fatal(err) } defer resp.Body.Close() b, _ := io.ReadAll(resp.Body) return resp.StatusCode, string(b) } func TestHealth(t *testing.T) { h := newHarness(t, options{}) if st, body := h.adminGet("/health"); st != 200 || !strings.Contains(body, "ok") { t.Errorf("healthy = %d %s", st, body) } h.write.Close() if st, body := h.adminGet("/health"); st != 503 || !strings.Contains(body, "redis") { t.Errorf("redis down = %d %s", st, body) } } func TestHealth_NoDomains(t *testing.T) { h := newHarness(t, options{}) h.db.Store(nil) if st, body := h.adminGet("/health"); st != 503 || !strings.Contains(body, "domains") { t.Errorf("no cdb = %d %s", st, body) } } func TestMetricsEndpoint(t *testing.T) { h := newHarness(t, options{}) h.accept("/in", `{"recipient":"b@yagolivas.com"}`) st, body := h.adminGet("/metrics") if st != 200 || !strings.Contains(body, `migadu_counters_requests_total{route="/in",status="202"} 1`) || !strings.Contains(body, "migadu_counters_queue_depth 0") { t.Errorf("metrics = %d\n%s", st, body) } } func TestNotifyHooks(t *testing.T) { reg := prometheus.NewRegistry() m := metrics.New(reg) buf := &syncBuffer{} hooks := NotifyHooks(txlog.New(buf, "text", func() time.Time { return clock }), m) ev := notify.Event{TraceID: "t", Route: "out", Level: "O"} hooks.Sent(ev, "http://x", 1) hooks.Failed(ev, "http://x", 5, io.EOF) if testutil.ToFloat64(m.Notifications.WithLabelValues("O", "sent")) != 1 || testutil.ToFloat64(m.Notifications.WithLabelValues("O", "failed")) != 1 { t.Error("hooks did not record metrics") } if s := buf.String(); !strings.Contains(s, "sent limit exceeded") || !strings.Contains(s, "failed limit exceeded") { t.Errorf("hooks did not log:\n%s", s) } } // Redis-backed /counters must use a client that is not the write client. var _ = redis.Nil // --- review #7 / #8 ----------------------------------------------------------- // Reviewer B: plus-addressed mail used to create cnt/A/b+1@… counters that the // web app never reads; now it counts against the mailbox. func TestIn_PlusTagCountsMailbox(t *testing.T) { h := newHarness(t, options{}) for _, r := range []string{"b+1@yagolivas.com", "b+2@yagolivas.com", "B+3@yagolivas.com"} { h.accept("/in", `{"trace_id":"plus","recipient":"`+r+`"}`) } if h.redisInt(store.Key("A", "b@yagolivas.com", day, "i")) != 3 { t.Errorf("mailbox counter = %d, want 3; keys %v", h.redisInt(store.Key("A", "b@yagolivas.com", day, "i")), h.write.Keys()) } for _, k := range h.write.Keys() { if strings.Contains(k, "+") { t.Errorf("tagged key written: %s", k) } } if !strings.Contains(h.tx.String(), "request mail_from= recipient=b+1@yagolivas.com") || !strings.Contains(h.tx.String(), "tracked address counters for b@yagolivas.com => count=1") { t.Errorf("log should show the tagged request and the mailbox counter:\n%s", h.tx.String()) } // /counters resolves the tag the same way (it reads the replica side, // which holds no data in this harness: only the identity is checked). h.read.Set(store.Key("A", "b@yagolivas.com", day, "i"), "3") st, body := h.get("/counters?address=b%2B9@yagolivas.com") if st != 200 { t.Fatalf("%d %s", st, body) } if r := decodeCounters(t, body); r.Query != "b@yagolivas.com" || r.Levels[0].ID != "b@yagolivas.com" || r.Levels[0].In != 3 { t.Errorf("/counters for a tagged address = %+v", r) } } // Reviewer B: a malformed row (non-numeric org id) used to 503 the request and // lose the sender's count. func TestMalformedRowCountsOtherSide(t *testing.T) { h := newHarness(t, options{}) db, err := domains.Open("../../review/_proof_tests/b/bad.cdb") if err != nil { t.Fatal(err) } defer db.Close() lf := counting.LookupFunc(db.Lookup) h.lookup.Store(&lf) h.db.Store(db) h.accept("/out", `{"trace_id":"bad1","auth_user":"a@good.com","recipient":"b@bad.com"}`) h.accept("/in", `{"trace_id":"bad2","recipient":"b@bad.com"}`) if h.redisInt(store.Key("A", "a@good.com", day, "o")) != 1 || len(h.write.Keys()) != 3 { t.Errorf("sender not counted / unexpected keys: %v", h.write.Keys()) } tx := h.tx.String() if !strings.Contains(tx, " out; recipient domain row malformed in domains db: bad.com:") || !strings.Contains(tx, " in; skipped: unknown_domain") { t.Errorf("transaction log:\n%s", tx) } if !strings.Contains(h.log.String(), "malformed") || !strings.Contains(h.log.String(), "bad.com") { t.Errorf("app log should warn about the row:\n%s", h.log.String()) } // /counters on such a domain: unmanaged, not an error. st, body := h.get("/counters?address=b@bad.com") if st != 200 || decodeCounters(t, body).Managed { t.Errorf("/counters = %d %s", st, body) } }