Files

418 lines
12 KiB
Go

package main
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"sync"
"syscall"
"testing"
"time"
"github.com/alicebob/miniredis/v2"
"github.com/colinmarc/cdb"
"migadu_counters/internal/api"
"migadu_counters/internal/config"
)
const realCDB = "../../test_files/rspamd_domains_db.cdb"
// env is one configured daemon with its own temp dir, miniredis and notify capture.
type env struct {
t *testing.T
dir string
cdbPath string
cfgPath string
logPath string
txPath string
redis *miniredis.Miniredis
notified *capture
}
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(200)
}
func copyFile(t *testing.T, src, dst string) {
t.Helper()
in, err := os.Open(src)
if err != nil {
t.Fatal(err)
}
defer in.Close()
out, err := os.Create(dst)
if err != nil {
t.Fatal(err)
}
if _, err := io.Copy(out, in); err != nil {
t.Fatal(err)
}
if err := out.Close(); err != nil {
t.Fatal(err)
}
}
// writeSmallCDB atomically replaces path with a CDB holding one domain.
func writeSmallCDB(t *testing.T, path string) {
t.Helper()
tmp := path + ".tmp"
w, err := cdb.Create(tmp)
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)
}
if err := os.Rename(tmp, path); err != nil {
t.Fatal(err)
}
}
func newEnv(t *testing.T, extra string) *env {
t.Helper()
e := &env{t: t, dir: t.TempDir(), redis: miniredis.RunT(t), notified: &capture{}}
e.cdbPath = filepath.Join(e.dir, "rspamd_domains_db")
copyFile(t, realCDB, e.cdbPath)
ns := httptest.NewServer(e.notified)
t.Cleanup(ns.Close)
e.logPath = filepath.Join(e.dir, "app.log")
e.txPath = filepath.Join(e.dir, "counting.log")
e.cfgPath = filepath.Join(e.dir, "config.toml")
q := "/api/limits?level=$LEVEL&id=$ID&flow=$FLOW&period=$PERIOD&by=$ADDRESS&count=$COUNT"
body := fmt.Sprintf(`
addr = "127.0.0.1:0"
admin_addr = "127.0.0.1:0"
auth_token = "t"
domains_cdb = %q
%s
[logging]
output = %q
counting_output = %q
[redis]
write_addr = %q
pool_size = 4
timeout_ms = 500
[notify]
org_url = %q
address_url = %q
user = "u"
password = "p"
timeout_seconds = 2
`, e.cdbPath, extra, e.logPath, e.txPath, e.redis.Addr(), ns.URL+q, ns.URL+q)
if err := os.WriteFile(e.cfgPath, []byte(body), 0o600); err != nil {
t.Fatal(err)
}
return e
}
func (e *env) logs() string {
b, _ := os.ReadFile(e.logPath)
return string(b)
}
func listen(t *testing.T) net.Listener {
t.Helper()
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
return ln
}
func httpDo(t *testing.T, method, url, body string) (int, string) {
t.Helper()
req, _ := http.NewRequest(method, url, strings.NewReader(body))
req.Header.Set("Authorization", "Bearer t")
if body != "" {
req.Header.Set("Content-Type", "application/json")
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatalf("%s %s: %v", method, url, err)
}
defer resp.Body.Close()
b, _ := io.ReadAll(resp.Body)
return resp.StatusCode, string(b)
}
func counters(t *testing.T, base, address string) api.CountersResponse {
t.Helper()
st, body := httpDo(t, http.MethodGet, base+"/counters?address="+address, "")
if st != 200 {
t.Fatalf("/counters %s: %d %s", address, st, body)
}
var r api.CountersResponse
if err := json.Unmarshal([]byte(body), &r); err != nil {
t.Fatal(err)
}
return r
}
func TestApp_EndToEnd(t *testing.T) {
e := newEnv(t, "")
cfg, err := config.Load(e.cfgPath)
if err != nil {
t.Fatal(err)
}
a, err := newApp(cfg)
if err != nil {
t.Fatalf("newApp: %v", err)
}
apiLn, adminLn := listen(t), listen(t)
a.Serve(apiLn, adminLn)
base := "http://" + apiLn.Addr().String()
admin := "http://" + adminLn.Addr().String()
// Count through the real stack: accepted immediately, counted by a worker.
st, body := httpDo(t, http.MethodPost, base+"/in", `{"trace_id":"e2e","mail_from":"s@gmail.com","recipient":"b@yagolivas.com"}`)
if st != 202 || strings.TrimSpace(body) != `{"accepted":true}` {
t.Fatalf("/in = %d %s", st, body)
}
key := "cnt/A/b@yagolivas.com/d/" + time.Now().UTC().Format("2006-01-02") + "/i"
for deadline := time.Now().Add(5 * time.Second); !e.redis.Exists(key); {
if time.Now().After(deadline) {
t.Fatalf("counter key missing; keys: %v", e.redis.Keys())
}
time.Sleep(10 * time.Millisecond)
}
if r := counters(t, base, "b@yagolivas.com"); r.Levels[0].In != 1 || *r.Levels[0].InLimit != 200 {
t.Errorf("/counters after /in = %+v", r.Levels[0])
}
if st, body := httpDo(t, http.MethodGet, admin+"/health", ""); st != 200 {
t.Errorf("/health = %d %s", st, body)
}
if st, body := httpDo(t, http.MethodGet, admin+"/metrics", ""); st != 200 || !strings.Contains(body, "migadu_counters_cdb_domains 93549") {
t.Errorf("/metrics = %d (cdb_domains gauge missing)", st)
}
for deadline := time.Now().Add(5 * time.Second); ; {
tx, _ := os.ReadFile(e.txPath)
if strings.Contains(string(tx), "<e2e> in; tracked address counters for b@yagolivas.com => count=1 limit=200 reached=false") {
break
}
if time.Now().After(deadline) {
t.Fatalf("counting log:\n%s", tx)
}
time.Sleep(10 * time.Millisecond)
}
if !strings.Contains(e.logs(), "listening") {
t.Errorf("app log should announce listening:\n%s", e.logs())
}
// SIGHUP path: the CDB was replaced by a different table.
writeSmallCDB(t, e.cdbPath)
if err := a.Reload(); err != nil {
t.Fatalf("Reload: %v", err)
}
if r := counters(t, base, "x@newdomain.test"); !r.Managed || r.OrgID == nil || *r.OrgID != 5 {
t.Errorf("after reload newdomain.test = %+v", r)
}
if r := counters(t, base, "a@yagolivas.com"); r.Managed {
t.Errorf("after reload yagolivas.com should be unmanaged: %+v", r)
}
if st, body := httpDo(t, http.MethodGet, admin+"/metrics", ""); !strings.Contains(body, "migadu_counters_cdb_domains 1") || !strings.Contains(body, `migadu_counters_cdb_reloads_total{result="success"} 1`) {
t.Errorf("/metrics after reload = %d:\n%s", st, body)
}
// A corrupt replacement (deployed atomically, as the agent's mv does) must
// be rejected and the old table kept. Never overwrite the CDB in place:
// that truncates the inode the daemon holds open.
if err := os.WriteFile(e.cdbPath+".bad", []byte("garbage"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.Rename(e.cdbPath+".bad", e.cdbPath); err != nil {
t.Fatal(err)
}
if err := a.Reload(); err == nil {
t.Error("Reload with a corrupt CDB should fail")
}
if r := counters(t, base, "x@newdomain.test"); !r.Managed {
t.Errorf("old table should survive a failed reload: %+v", r)
}
// Watcher path: an atomic replacement reloads without a signal.
tmp := e.cdbPath + ".new"
copyFile(t, realCDB, tmp)
if err := os.Rename(tmp, e.cdbPath); err != nil {
t.Fatal(err)
}
deadline := time.Now().Add(8 * time.Second)
for {
if r := counters(t, base, "a@yagolivas.com"); r.Managed {
break
}
if time.Now().After(deadline) {
t.Fatalf("watcher did not reload the replaced CDB; log:\n%s", e.logs())
}
time.Sleep(100 * time.Millisecond)
}
// Log rotation: after renaming the app log, Reload reopens it.
rotated := e.logPath + ".0"
if err := os.Rename(e.logPath, rotated); err != nil {
t.Fatal(err)
}
if err := a.Reload(); err != nil {
t.Fatal(err)
}
httpDo(t, http.MethodPost, base+"/in", `{}`) // a rejection -> warn line
if !strings.Contains(e.logs(), "request rejected") {
t.Errorf("new log file should receive lines after rotation:\n%s", e.logs())
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
if err := a.Shutdown(ctx); err != nil {
t.Fatalf("Shutdown: %v", err)
}
if _, err := http.Get(base + "/counters"); err == nil {
t.Error("API still answering after Shutdown")
}
if !strings.Contains(e.logs(), "stopped") {
t.Errorf("log should record the stop:\n%s", e.logs())
}
}
func TestApp_ReachedNotifies(t *testing.T) {
e := newEnv(t, "")
cfg, _ := config.Load(e.cfgPath)
a, err := newApp(cfg)
if err != nil {
t.Fatal(err)
}
apiLn, adminLn := listen(t), listen(t)
a.Serve(apiLn, adminLn)
base := "http://" + apiLn.Addr().String()
for i := 0; i < 20; i++ { // A/o=20 for yagolivas.com
if st, body := httpDo(t, http.MethodPost, base+"/out", `{"auth_user":"a@yagolivas.com","recipient":"x@gmail.com"}`); st != 202 {
t.Fatalf("/out = %d %s", st, body)
}
}
// Shutdown drains the queue, then the notifier: every notification lands.
if err := a.Shutdown(context.Background()); err != nil {
t.Fatal(err)
}
if v, _ := e.redis.Get("cnt/A/a@yagolivas.com/d/" + time.Now().UTC().Format("2006-01-02") + "/o"); v != "20" {
t.Errorf("after Shutdown A/o = %q, want 20 (queue drained)", v)
}
e.notified.mu.Lock()
n := len(e.notified.urls)
e.notified.mu.Unlock()
if n != 2 { // address_url + org_url (domain_url not configured)
t.Errorf("notifications = %d, want 2: %v", n, e.notified.urls)
}
if tx, _ := os.ReadFile(e.txPath); strings.Count(string(tx), "sent limit exceeded notification") != 2 {
t.Errorf("counting log should have two sent lines:\n%s", tx)
}
}
func TestNewApp_MissingCDBIsFatal(t *testing.T) {
e := newEnv(t, "")
os.Remove(e.cdbPath)
cfg, _ := config.Load(e.cfgPath)
if _, err := newApp(cfg); err == nil || !strings.Contains(err.Error(), "rspamd_domains_db") {
t.Errorf("newApp without CDB: %v", err)
}
}
func TestNewApp_RedisDownIsNotFatal(t *testing.T) {
e := newEnv(t, "")
e.redis.Close()
cfg, _ := config.Load(e.cfgPath)
a, err := newApp(cfg)
if err != nil {
t.Fatalf("Redis being down must not prevent startup: %v", err)
}
apiLn, adminLn := listen(t), listen(t)
a.Serve(apiLn, adminLn)
defer a.Shutdown(context.Background())
if st, body := httpDo(t, http.MethodGet, "http://"+adminLn.Addr().String()+"/health", ""); st != 503 || !strings.Contains(body, "redis") {
t.Errorf("/health with redis down = %d %s", st, body)
}
if !strings.Contains(e.logs(), "redis") {
t.Errorf("startup should warn about redis:\n%s", e.logs())
}
}
func TestRun_FlagsAndExitCodes(t *testing.T) {
var out, errb bytes.Buffer
if code := run([]string{"-version"}, &out, &errb, nil); code != 0 || !strings.HasPrefix(out.String(), "migadu_counters ") {
t.Errorf("-version: %d %q", code, out.String())
}
out.Reset()
errb.Reset()
if code := run([]string{"-config", "/nonexistent/config.toml"}, &out, &errb, nil); code != 1 || !strings.Contains(errb.String(), "/nonexistent/config.toml") {
t.Errorf("missing config: %d %q", code, errb.String())
}
errb.Reset()
bad := filepath.Join(t.TempDir(), "bad.toml")
os.WriteFile(bad, []byte("domains_cdb = \"x\"\nauth_token = \"t\"\n[redis]\ndb = 1\nwrite_addr = \"127.0.0.1:1\"\n"), 0o600)
if code := run([]string{"-config", bad}, &out, &errb, nil); code != 1 || !strings.Contains(errb.String(), "redis.db") {
t.Errorf("invalid config: %d %q", code, errb.String())
}
// Review #6: an empty / mistyped token never starts the daemon.
for _, body := range []string{"domains_cdb = \"x\"\n[redis]\nwrite_addr = \"127.0.0.1:1\"\n", "domains_cdb = \"x\"\nauth-token = \"t\"\n[redis]\nwrite_addr = \"127.0.0.1:1\"\n"} {
errb.Reset()
os.WriteFile(bad, []byte(body), 0o600)
if code := run([]string{"-config", bad}, &out, &errb, nil); code != 1 || !strings.Contains(errb.String(), "auth") {
t.Errorf("config %q: exit %d stderr %q, want 1 mentioning the token", body, code, errb.String())
}
}
}
// run() serves until SIGTERM; SIGHUP in between reloads.
func TestRun_SignalLoop(t *testing.T) {
e := newEnv(t, "")
sigs := make(chan os.Signal, 2)
done := make(chan int, 1)
var out, errb bytes.Buffer
go func() { done <- run([]string{"-config", e.cfgPath}, &out, &errb, sigs) }()
deadline := time.Now().Add(5 * time.Second)
for !strings.Contains(e.logs(), "listening") {
if time.Now().After(deadline) {
t.Fatalf("daemon did not start; stderr %q log:\n%s", errb.String(), e.logs())
}
time.Sleep(50 * time.Millisecond)
}
sigs <- syscall.SIGHUP
deadline = time.Now().Add(5 * time.Second)
for !strings.Contains(e.logs(), "reloaded") {
if time.Now().After(deadline) {
t.Fatalf("SIGHUP did not reload; log:\n%s", e.logs())
}
time.Sleep(50 * time.Millisecond)
}
sigs <- syscall.SIGTERM
select {
case code := <-done:
if code != 0 {
t.Errorf("exit code %d, stderr %q", code, errb.String())
}
case <-time.After(15 * time.Second):
t.Fatal("run did not return after SIGTERM")
}
if !strings.Contains(e.logs(), "stopped") {
t.Errorf("log:\n%s", e.logs())
}
}