First commit Go counter service
This commit is contained in:
@@ -0,0 +1,255 @@
|
||||
// counterctl shows today's (or a given day's) incoming/outgoing counters for
|
||||
// an address or a domain by asking the running migadu_counters daemon
|
||||
// (GET /counters), the same way rcptctl talks to rcptd.
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/BurntSushi/toml"
|
||||
|
||||
"migadu_counters/internal/api"
|
||||
"migadu_counters/internal/config"
|
||||
)
|
||||
|
||||
// Set by -ldflags at build time.
|
||||
var (
|
||||
version = "dev"
|
||||
commit = "none"
|
||||
date = "unknown"
|
||||
)
|
||||
|
||||
func main() {
|
||||
os.Exit(run(os.Args[1:], os.Stdout, os.Stderr, os.Getenv, time.Now))
|
||||
}
|
||||
|
||||
// cliConfig mirrors the daemon config fields the CLI needs; TOML ignores the
|
||||
// rest, so the daemon's own config file can be used as is.
|
||||
type cliConfig struct {
|
||||
Addr string `toml:"addr"`
|
||||
AuthToken string `toml:"auth_token"`
|
||||
}
|
||||
|
||||
func run(args []string, stdout, stderr io.Writer, getenv func(string) string, now func() time.Time) int {
|
||||
fs := flag.NewFlagSet("counterctl", flag.ContinueOnError)
|
||||
fs.SetOutput(stderr)
|
||||
address := fs.String("address", "", "mailbox address to show (A, D and O levels)")
|
||||
domain := fs.String("domain", "", "domain to show (D and O levels)")
|
||||
dateFlag := fs.String("date", "", "ISO 8601 calendar day YYYY-MM-DD, UTC (default: today)")
|
||||
jsonOut := fs.Bool("json", false, "print the raw API response")
|
||||
inOnly := fs.Bool("in", false, "show only the incoming counters")
|
||||
outOnly := fs.Bool("out", false, "show only the outgoing counters")
|
||||
configPath := fs.String("config", "", "config file (default: "+strings.Join(config.DefaultPaths, " or ")+")")
|
||||
urlFlag := fs.String("url", "", "service base URL (default: derived from addr in the config)")
|
||||
showVersion := fs.Bool("version", false, "print version and exit")
|
||||
fs.Usage = func() {
|
||||
fmt.Fprintf(stderr, "Usage: counterctl --address me@example.com | --domain example.com [--date YYYY-MM-DD] [--in|--out] [--json] [--config PATH] [--url URL]\n\n")
|
||||
fs.PrintDefaults()
|
||||
}
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return 1
|
||||
}
|
||||
if *showVersion {
|
||||
fmt.Fprintf(stdout, "counterctl %s (commit %s, built %s)\n", version, commit, date)
|
||||
return 0
|
||||
}
|
||||
if (*address == "") == (*domain == "") {
|
||||
fmt.Fprintln(stderr, "exactly one of --address or --domain is required")
|
||||
fs.Usage()
|
||||
return 1
|
||||
}
|
||||
|
||||
q := url.Values{}
|
||||
if *address != "" {
|
||||
addr, err := api.NormalizeEmail(*address)
|
||||
if err != nil {
|
||||
fmt.Fprintf(stderr, "invalid address %q: %v\n", *address, err)
|
||||
return 1
|
||||
}
|
||||
q.Set("address", addr)
|
||||
} else {
|
||||
q.Set("domain", strings.ToLower(strings.TrimSpace(*domain)))
|
||||
}
|
||||
if *dateFlag != "" {
|
||||
day, err := api.ParseDate(*dateFlag, now())
|
||||
if err != nil {
|
||||
fmt.Fprintln(stderr, err)
|
||||
return 1
|
||||
}
|
||||
q.Set("date", day)
|
||||
if t, _ := time.Parse("2006-01-02", day); now().UTC().Sub(t) > 31*24*time.Hour {
|
||||
fmt.Fprintf(stderr, "warning: %s is more than 31 days ago; counter keys expire after 31 days and will read as 0\n", day)
|
||||
}
|
||||
}
|
||||
|
||||
cfg, err := loadConfig(*configPath, *urlFlag != "")
|
||||
if err != nil {
|
||||
fmt.Fprintln(stderr, err)
|
||||
return 1
|
||||
}
|
||||
token := cfg.AuthToken
|
||||
if v := getenv(config.EnvAuthToken); v != "" {
|
||||
token = v
|
||||
}
|
||||
base := *urlFlag
|
||||
if base == "" {
|
||||
base = baseURL(cfg.Addr)
|
||||
if host := strings.TrimPrefix(strings.TrimPrefix(base, "http://"), "https://"); !strings.HasPrefix(host, "localhost") && !strings.HasPrefix(host, "127.0.0.1") && !strings.HasPrefix(host, "[::1]") {
|
||||
fmt.Fprintf(stderr, "warning: %s is not local; the auth token travels over plaintext HTTP\n", base)
|
||||
}
|
||||
}
|
||||
|
||||
req, err := http.NewRequest(http.MethodGet, base+"/counters?"+q.Encode(), nil)
|
||||
if err != nil {
|
||||
fmt.Fprintf(stderr, "bad URL: %v\n", err)
|
||||
return 1
|
||||
}
|
||||
if token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
}
|
||||
resp, err := (&http.Client{Timeout: 10 * time.Second}).Do(req)
|
||||
if err != nil {
|
||||
fmt.Fprintf(stderr, "request failed: %v\n", err)
|
||||
return 1
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
if err != nil {
|
||||
fmt.Fprintf(stderr, "reading response: %v\n", err)
|
||||
return 1
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
msg := strings.TrimSpace(string(body))
|
||||
var e api.ErrorResponse
|
||||
if json.Unmarshal(body, &e) == nil && e.Error != "" {
|
||||
msg = e.Error
|
||||
}
|
||||
fmt.Fprintf(stderr, "server answered %d: %s\n", resp.StatusCode, msg)
|
||||
return 1
|
||||
}
|
||||
if *jsonOut {
|
||||
stdout.Write(body)
|
||||
if len(body) == 0 || body[len(body)-1] != '\n' {
|
||||
io.WriteString(stdout, "\n")
|
||||
}
|
||||
return 0
|
||||
}
|
||||
var r api.CountersResponse
|
||||
if err := json.Unmarshal(body, &r); err != nil {
|
||||
fmt.Fprintf(stderr, "unexpected response: %v\n%s\n", err, body)
|
||||
return 1
|
||||
}
|
||||
showIn, showOut := flows(*inOnly, *outOnly)
|
||||
printTable(stdout, r, showIn, showOut)
|
||||
return 0
|
||||
}
|
||||
|
||||
// flows picks the table columns: --in or --out alone restrict the table to
|
||||
// that flow; neither or both show both flows.
|
||||
func flows(in, out bool) (showIn, showOut bool) {
|
||||
if in == out {
|
||||
return true, true
|
||||
}
|
||||
return in, out
|
||||
}
|
||||
|
||||
// loadConfig reads addr/auth_token from the daemon config. Without -config
|
||||
// the default paths are tried; when none exists the CLI can still run with
|
||||
// -url and the token from the environment.
|
||||
func loadConfig(path string, haveURL bool) (cliConfig, error) {
|
||||
var cfg cliConfig
|
||||
if path == "" {
|
||||
found, err := config.Discover("")
|
||||
if err != nil {
|
||||
if haveURL {
|
||||
return cfg, nil
|
||||
}
|
||||
return cfg, err
|
||||
}
|
||||
path = found
|
||||
}
|
||||
if _, err := toml.DecodeFile(path, &cfg); err != nil {
|
||||
return cfg, fmt.Errorf("read config %s: %v", path, err)
|
||||
}
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
// baseURL turns the daemon's listen address into a URL: ":5555" →
|
||||
// "http://localhost:5555"; a full URL is used as is.
|
||||
func baseURL(addr string) string {
|
||||
if addr == "" {
|
||||
addr = ":5555"
|
||||
}
|
||||
if strings.HasPrefix(addr, "http://") || strings.HasPrefix(addr, "https://") {
|
||||
return strings.TrimRight(addr, "/")
|
||||
}
|
||||
if strings.HasPrefix(addr, ":") {
|
||||
addr = "localhost" + addr
|
||||
}
|
||||
return "http://" + addr
|
||||
}
|
||||
|
||||
func printTable(w io.Writer, r api.CountersResponse, showIn, showOut bool) {
|
||||
dom, org := "-", "-"
|
||||
if r.Domain != "" {
|
||||
dom = r.Domain
|
||||
}
|
||||
if r.OrgID != nil {
|
||||
org = strconv.FormatInt(*r.OrgID, 10)
|
||||
}
|
||||
fmt.Fprintf(w, "%s date %s (UTC) domain %s org %s\n", r.Query, r.Date, dom, org)
|
||||
|
||||
tw := tabwriter.NewWriter(w, 0, 8, 2, ' ', 0)
|
||||
head := []string{"LEVEL", "ID"}
|
||||
if showIn {
|
||||
head = append(head, "IN", "LIMIT")
|
||||
}
|
||||
if showOut {
|
||||
head = append(head, "OUT", "LIMIT")
|
||||
}
|
||||
fmt.Fprintln(tw, strings.Join(head, "\t"))
|
||||
for _, l := range r.Levels {
|
||||
cols := []string{l.Level, l.ID}
|
||||
if showIn {
|
||||
cols = append(cols, count(l.In, l.InLimit), limit(l.InLimit))
|
||||
}
|
||||
if showOut {
|
||||
cols = append(cols, count(l.Out, l.OutLimit), limit(l.OutLimit))
|
||||
}
|
||||
fmt.Fprintln(tw, strings.Join(cols, "\t"))
|
||||
}
|
||||
tw.Flush()
|
||||
|
||||
if !r.Managed {
|
||||
fmt.Fprintln(w, "note: domain is not in rspamd_domains_db (not counted)")
|
||||
}
|
||||
if r.Exempt {
|
||||
fmt.Fprintln(w, "note: address is exempt from counting")
|
||||
}
|
||||
}
|
||||
|
||||
// count renders a value, marked with "!" once it reached its limit.
|
||||
func count(n int64, lim *int64) string {
|
||||
s := strconv.FormatInt(n, 10)
|
||||
if lim != nil && n >= *lim {
|
||||
s += "!"
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func limit(lim *int64) string {
|
||||
if lim == nil {
|
||||
return "-"
|
||||
}
|
||||
return strconv.FormatInt(*lim, 10)
|
||||
}
|
||||
@@ -0,0 +1,335 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"migadu_counters/internal/api"
|
||||
)
|
||||
|
||||
var now = time.Date(2026, 10, 13, 12, 0, 0, 0, time.UTC)
|
||||
|
||||
func i64(n int64) *int64 { return &n }
|
||||
|
||||
func sampleResponse() api.CountersResponse {
|
||||
return api.CountersResponse{
|
||||
Query: "me@example.com", Mode: "address", Date: "2026-10-13", Managed: true,
|
||||
Domain: "example.com", OrgID: i64(1234),
|
||||
Levels: []api.CountersLevel{
|
||||
{Level: "A", ID: "me@example.com", In: 12, InLimit: i64(500), Out: 200, OutLimit: i64(200)},
|
||||
{Level: "D", ID: "example.com", In: 40, InLimit: i64(1000), Out: 9, OutLimit: nil},
|
||||
{Level: "O", ID: "1234", In: 120, InLimit: i64(5000), Out: 30, OutLimit: i64(1000)},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// fakeAPI serves GET /counters and records what it was asked.
|
||||
type fakeAPI struct {
|
||||
mu sync.Mutex
|
||||
requests []*http.Request
|
||||
status int
|
||||
body any
|
||||
}
|
||||
|
||||
func (f *fakeAPI) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
f.mu.Lock()
|
||||
f.requests = append(f.requests, r)
|
||||
f.mu.Unlock()
|
||||
if f.status == 0 {
|
||||
f.status = 200
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(f.status)
|
||||
_ = json.NewEncoder(w).Encode(f.body)
|
||||
}
|
||||
|
||||
func (f *fakeAPI) last() *http.Request {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return f.requests[len(f.requests)-1]
|
||||
}
|
||||
|
||||
func writeConfig(t *testing.T, body string) string {
|
||||
t.Helper()
|
||||
p := filepath.Join(t.TempDir(), "config.toml")
|
||||
if err := os.WriteFile(p, []byte(body), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
func noEnv(string) string { return "" }
|
||||
|
||||
func runCLI(t *testing.T, env func(string) string, args ...string) (int, string, string) {
|
||||
t.Helper()
|
||||
var out, errb bytes.Buffer
|
||||
code := run(args, &out, &errb, env, func() time.Time { return now })
|
||||
return code, out.String(), errb.String()
|
||||
}
|
||||
|
||||
func TestRun_Table(t *testing.T) {
|
||||
f := &fakeAPI{body: sampleResponse()}
|
||||
s := httptest.NewServer(f)
|
||||
defer s.Close()
|
||||
cfg := writeConfig(t, "addr = \":8095\"\nauth_token = \"filetoken\"\n")
|
||||
|
||||
code, out, errOut := runCLI(t, noEnv, "--config", cfg, "--url", s.URL, "--address", "Me@Example.com")
|
||||
if code != 0 {
|
||||
t.Fatalf("exit %d, stderr %q", code, errOut)
|
||||
}
|
||||
req := f.last()
|
||||
if req.URL.Path != "/counters" || req.URL.Query().Get("address") != "me@example.com" || req.URL.Query().Has("date") || req.URL.Query().Has("domain") {
|
||||
t.Errorf("request = %s", req.URL)
|
||||
}
|
||||
if req.Header.Get("Authorization") != "Bearer filetoken" {
|
||||
t.Errorf("Authorization = %q", req.Header.Get("Authorization"))
|
||||
}
|
||||
|
||||
lines := strings.Split(strings.TrimSpace(out), "\n")
|
||||
if len(lines) != 5 {
|
||||
t.Fatalf("want header + column line + 3 rows, got %d lines:\n%s", len(lines), out)
|
||||
}
|
||||
for _, want := range []string{"me@example.com", "date 2026-10-13 (UTC)", "domain example.com", "org 1234"} {
|
||||
if !strings.Contains(lines[0], want) {
|
||||
t.Errorf("header %q lacks %q", lines[0], want)
|
||||
}
|
||||
}
|
||||
if cols := strings.Fields(lines[1]); strings.Join(cols, " ") != "LEVEL ID IN LIMIT OUT LIMIT" {
|
||||
t.Errorf("column line = %q", lines[1])
|
||||
}
|
||||
rows := [][]string{strings.Fields(lines[2]), strings.Fields(lines[3]), strings.Fields(lines[4])}
|
||||
want := [][]string{
|
||||
{"A", "me@example.com", "12", "500", "200!", "200"}, // out reached its limit
|
||||
{"D", "example.com", "40", "1000", "9", "-"}, // unlimited out
|
||||
{"O", "1234", "120", "5000", "30", "1000"},
|
||||
}
|
||||
for i := range want {
|
||||
if strings.Join(rows[i], " ") != strings.Join(want[i], " ") {
|
||||
t.Errorf("row %d = %v, want %v", i, rows[i], want[i])
|
||||
}
|
||||
}
|
||||
if errOut != "" {
|
||||
t.Errorf("unexpected stderr: %q", errOut)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_UnmanagedAndExemptNotes(t *testing.T) {
|
||||
f := &fakeAPI{body: api.CountersResponse{Query: "x@gmail.com", Mode: "address", Date: "2026-10-13",
|
||||
Levels: []api.CountersLevel{{Level: "A", ID: "x@gmail.com"}, {Level: "D", ID: "gmail.com"}}}}
|
||||
s := httptest.NewServer(f)
|
||||
defer s.Close()
|
||||
cfg := writeConfig(t, "auth_token = \"t\"\n")
|
||||
code, out, _ := runCLI(t, noEnv, "--config", cfg, "--url", s.URL, "--address", "x@gmail.com")
|
||||
if code != 0 || !strings.Contains(out, "not in rspamd_domains_db") || !strings.Contains(out, "org -") {
|
||||
t.Errorf("exit %d output:\n%s", code, out)
|
||||
}
|
||||
|
||||
f.body = api.CountersResponse{Query: "support@migadu.com", Mode: "address", Date: "2026-10-13", Managed: true, Exempt: true,
|
||||
Domain: "migadu.com", OrgID: i64(3743), Levels: []api.CountersLevel{}}
|
||||
code, out, _ = runCLI(t, noEnv, "--config", cfg, "--url", s.URL, "--address", "support@migadu.com")
|
||||
if code != 0 || !strings.Contains(out, "exempt from counting") {
|
||||
t.Errorf("exit %d output:\n%s", code, out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_JSON(t *testing.T) {
|
||||
f := &fakeAPI{body: sampleResponse()}
|
||||
s := httptest.NewServer(f)
|
||||
defer s.Close()
|
||||
cfg := writeConfig(t, "auth_token = \"t\"\n")
|
||||
code, out, _ := runCLI(t, noEnv, "--config", cfg, "--url", s.URL, "--address", "me@example.com", "--json")
|
||||
if code != 0 {
|
||||
t.Fatalf("exit %d", code)
|
||||
}
|
||||
var got api.CountersResponse
|
||||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||||
t.Fatalf("output is not the API JSON: %v\n%s", err, out)
|
||||
}
|
||||
if got.Query != "me@example.com" || len(got.Levels) != 3 || *got.Levels[0].InLimit != 500 {
|
||||
t.Errorf("decoded = %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_DomainMode(t *testing.T) {
|
||||
f := &fakeAPI{body: api.CountersResponse{Query: "example.com", Mode: "domain", Date: "2026-10-13", Managed: true,
|
||||
Domain: "example.com", OrgID: i64(1234),
|
||||
Levels: []api.CountersLevel{{Level: "D", ID: "example.com", In: 1, InLimit: i64(5), Out: 2, OutLimit: i64(6)}, {Level: "O", ID: "1234"}}}}
|
||||
s := httptest.NewServer(f)
|
||||
defer s.Close()
|
||||
cfg := writeConfig(t, "auth_token = \"t\"\n")
|
||||
code, out, _ := runCLI(t, noEnv, "--config", cfg, "--url", s.URL, "--domain", "Example.COM")
|
||||
if code != 0 {
|
||||
t.Fatalf("exit %d\n%s", code, out)
|
||||
}
|
||||
if q := f.last().URL.Query(); q.Get("domain") != "example.com" || q.Has("address") {
|
||||
t.Errorf("query = %v", q)
|
||||
}
|
||||
if !strings.Contains(out, "D example.com") && !strings.Contains(out, "D example.com") {
|
||||
t.Errorf("table lacks the D row:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_Date(t *testing.T) {
|
||||
f := &fakeAPI{body: sampleResponse()}
|
||||
s := httptest.NewServer(f)
|
||||
defer s.Close()
|
||||
cfg := writeConfig(t, "auth_token = \"t\"\n")
|
||||
|
||||
code, _, errOut := runCLI(t, noEnv, "--config", cfg, "--url", s.URL, "--address", "me@example.com", "--date", "2026-10-01")
|
||||
if code != 0 || f.last().URL.Query().Get("date") != "2026-10-01" {
|
||||
t.Errorf("exit %d stderr %q query %v", code, errOut, f.last().URL.Query())
|
||||
}
|
||||
if errOut != "" {
|
||||
t.Errorf("12 days back should not warn: %q", errOut)
|
||||
}
|
||||
// Older than the 31-day key TTL: still queried, but warned about.
|
||||
code, _, errOut = runCLI(t, noEnv, "--config", cfg, "--url", s.URL, "--address", "me@example.com", "--date", "2026-09-01")
|
||||
if code != 0 || !strings.Contains(errOut, "31 days") {
|
||||
t.Errorf("exit %d stderr %q, want a 31-day warning", code, errOut)
|
||||
}
|
||||
for _, bad := range []struct{ date, want string }{{"2099-01-01", "future"}, {"13.10.2026", "YYYY-MM-DD"}, {"2026-1-5", "YYYY-MM-DD"}} {
|
||||
before := len(f.requests)
|
||||
code, _, errOut = runCLI(t, noEnv, "--config", cfg, "--url", s.URL, "--address", "me@example.com", "--date", bad.date)
|
||||
if code != 1 || !strings.Contains(errOut, bad.want) {
|
||||
t.Errorf("--date %s: exit %d stderr %q, want exit 1 mentioning %q", bad.date, code, errOut, bad.want)
|
||||
}
|
||||
if len(f.requests) != before {
|
||||
t.Errorf("--date %s: request was sent despite the invalid date", bad.date)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_Usage(t *testing.T) {
|
||||
cfg := writeConfig(t, "auth_token = \"t\"\n")
|
||||
for _, args := range [][]string{
|
||||
{"--config", cfg},
|
||||
{"--config", cfg, "--address", "a@b.c", "--domain", "b.c"},
|
||||
{"--config", cfg, "--address", "not-an-address"},
|
||||
{"--config", cfg, "--bogus"},
|
||||
} {
|
||||
code, _, errOut := runCLI(t, noEnv, args...)
|
||||
if code != 1 || errOut == "" {
|
||||
t.Errorf("%v: exit %d stderr %q, want 1 with a message", args, code, errOut)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_ServerErrors(t *testing.T) {
|
||||
f := &fakeAPI{status: 401, body: api.ErrorResponse{Error: "unauthorized"}}
|
||||
s := httptest.NewServer(f)
|
||||
defer s.Close()
|
||||
cfg := writeConfig(t, "auth_token = \"wrong\"\n")
|
||||
code, _, errOut := runCLI(t, noEnv, "--config", cfg, "--url", s.URL, "--address", "me@example.com")
|
||||
if code != 1 || !strings.Contains(errOut, "401") || !strings.Contains(errOut, "unauthorized") {
|
||||
t.Errorf("exit %d stderr %q", code, errOut)
|
||||
}
|
||||
code, _, errOut = runCLI(t, noEnv, "--config", cfg, "--url", "http://127.0.0.1:1", "--address", "me@example.com")
|
||||
if code != 1 || errOut == "" {
|
||||
t.Errorf("unreachable: exit %d stderr %q", code, errOut)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_ConfigAndEnv(t *testing.T) {
|
||||
f := &fakeAPI{body: sampleResponse()}
|
||||
s := httptest.NewServer(f)
|
||||
defer s.Close()
|
||||
cfg := writeConfig(t, "addr = \":8095\"\nauth_token = \"filetoken\"\n")
|
||||
env := func(k string) string {
|
||||
if k == "MIGADU_COUNTERS_AUTH_TOKEN" {
|
||||
return "envtoken"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
if code, _, errOut := runCLI(t, env, "--config", cfg, "--url", s.URL, "--address", "me@example.com"); code != 0 {
|
||||
t.Fatalf("exit %d %q", code, errOut)
|
||||
}
|
||||
if got := f.last().Header.Get("Authorization"); got != "Bearer envtoken" {
|
||||
t.Errorf("env token not used: %q", got)
|
||||
}
|
||||
// Missing config file is an error naming the path.
|
||||
if code, _, errOut := runCLI(t, noEnv, "--config", "/nonexistent/x.toml", "--address", "me@example.com"); code != 1 || !strings.Contains(errOut, "/nonexistent/x.toml") {
|
||||
t.Errorf("exit %d stderr %q", code, errOut)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBaseURL(t *testing.T) {
|
||||
for in, want := range map[string]string{
|
||||
":5555": "http://localhost:5555",
|
||||
"10.0.0.5:5555": "http://10.0.0.5:5555",
|
||||
"localhost:1": "http://localhost:1",
|
||||
"http://x:1": "http://x:1",
|
||||
"https://x.example/": "https://x.example",
|
||||
"": "http://localhost:5555",
|
||||
} {
|
||||
if got := baseURL(in); got != want {
|
||||
t.Errorf("baseURL(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --in / --out restrict the table to one flow; both or neither give the full table.
|
||||
func TestRun_FlowFlags(t *testing.T) {
|
||||
f := &fakeAPI{body: sampleResponse()}
|
||||
s := httptest.NewServer(f)
|
||||
defer s.Close()
|
||||
cfg := writeConfig(t, "auth_token = \"t\"\n")
|
||||
base := []string{"--config", cfg, "--url", s.URL, "--address", "me@example.com"}
|
||||
|
||||
code, out, errOut := runCLI(t, noEnv, append(base, "--in")...)
|
||||
if code != 0 || errOut != "" {
|
||||
t.Fatalf("--in: exit %d stderr %q", code, errOut)
|
||||
}
|
||||
lines := strings.Split(strings.TrimSpace(out), "\n")
|
||||
if len(lines) != 5 || strings.Join(strings.Fields(lines[1]), " ") != "LEVEL ID IN LIMIT" {
|
||||
t.Errorf("--in table:\n%s", out)
|
||||
}
|
||||
wantIn := [][]string{{"A", "me@example.com", "12", "500"}, {"D", "example.com", "40", "1000"}, {"O", "1234", "120", "5000"}}
|
||||
for i, w := range wantIn {
|
||||
if got := strings.Join(strings.Fields(lines[2+i]), " "); got != strings.Join(w, " ") {
|
||||
t.Errorf("--in row %d = %q, want %v", i, got, w)
|
||||
}
|
||||
}
|
||||
|
||||
code, out, _ = runCLI(t, noEnv, append(base, "--out")...)
|
||||
if code != 0 {
|
||||
t.Fatalf("--out: exit %d", code)
|
||||
}
|
||||
lines = strings.Split(strings.TrimSpace(out), "\n")
|
||||
if len(lines) != 5 || strings.Join(strings.Fields(lines[1]), " ") != "LEVEL ID OUT LIMIT" {
|
||||
t.Errorf("--out table:\n%s", out)
|
||||
}
|
||||
wantOut := [][]string{{"A", "me@example.com", "200!", "200"}, {"D", "example.com", "9", "-"}, {"O", "1234", "30", "1000"}}
|
||||
for i, w := range wantOut {
|
||||
if got := strings.Join(strings.Fields(lines[2+i]), " "); got != strings.Join(w, " ") {
|
||||
t.Errorf("--out row %d = %q, want %v", i, got, w)
|
||||
}
|
||||
}
|
||||
|
||||
_, full, _ := runCLI(t, noEnv, base...)
|
||||
_, both, _ := runCLI(t, noEnv, append(base, "--in", "--out")...)
|
||||
if both != full {
|
||||
t.Errorf("--in --out should equal the full table:\n%s\nvs\n%s", both, full)
|
||||
}
|
||||
|
||||
// --json is the raw response whatever the flow flags say.
|
||||
code, out, _ = runCLI(t, noEnv, append(base, "--out", "--json")...)
|
||||
var got api.CountersResponse
|
||||
if code != 0 || json.Unmarshal([]byte(out), &got) != nil || len(got.Levels) != 3 || got.Levels[0].In != 12 {
|
||||
t.Errorf("--out --json: exit %d output %s", code, out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_Version(t *testing.T) {
|
||||
code, out, _ := runCLI(t, noEnv, "--version")
|
||||
if code != 0 || !strings.HasPrefix(out, "counterctl ") {
|
||||
t.Errorf("exit %d output %q", code, out)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,349 @@
|
||||
// migadu_counters counts mail per address / domain / organisation in Redis and
|
||||
// notifies the web app when a daily limit is reached. mailqueuer calls it
|
||||
// fire-and-forget once per ingested recipient (POST /in, POST /out).
|
||||
//
|
||||
// It is a foreground process: on FreeBSD it runs under rc.d + daemon(8)
|
||||
// (pidfile, log capture, restart). SIGHUP reopens the log files and reloads
|
||||
// the domains CDB; SIGTERM/SIGINT drain and exit.
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/prometheus/client_golang/prometheus/collectors"
|
||||
|
||||
"migadu_counters/internal/config"
|
||||
"migadu_counters/internal/counting"
|
||||
"migadu_counters/internal/domains"
|
||||
"migadu_counters/internal/metrics"
|
||||
"migadu_counters/internal/notify"
|
||||
"migadu_counters/internal/server"
|
||||
"migadu_counters/internal/store"
|
||||
"migadu_counters/internal/txlog"
|
||||
"migadu_counters/internal/watcher"
|
||||
)
|
||||
|
||||
// Set by -ldflags at build time.
|
||||
var (
|
||||
version = "dev"
|
||||
commit = "none"
|
||||
date = "unknown"
|
||||
)
|
||||
|
||||
func main() {
|
||||
os.Exit(run(os.Args[1:], os.Stdout, os.Stderr, nil))
|
||||
}
|
||||
|
||||
// run parses flags, starts the daemon and serves until SIGTERM/SIGINT. sigs
|
||||
// may be nil (then the real process signals are used); tests inject one.
|
||||
func run(args []string, stdout, stderr io.Writer, sigs <-chan os.Signal) int {
|
||||
fs := flag.NewFlagSet("migadu_counters", flag.ContinueOnError)
|
||||
fs.SetOutput(stderr)
|
||||
configPath := fs.String("config", "", "path to config file (TOML)")
|
||||
showVersion := fs.Bool("version", false, "print version and exit")
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return 1
|
||||
}
|
||||
if *showVersion {
|
||||
fmt.Fprintf(stdout, "migadu_counters %s (commit %s, built %s)\n", version, commit, date)
|
||||
return 0
|
||||
}
|
||||
|
||||
path, err := config.Discover(*configPath)
|
||||
if err != nil {
|
||||
fmt.Fprintln(stderr, err)
|
||||
return 1
|
||||
}
|
||||
cfg, err := config.Load(path)
|
||||
if err != nil {
|
||||
fmt.Fprintln(stderr, err)
|
||||
return 1
|
||||
}
|
||||
a, err := newApp(cfg)
|
||||
if err != nil {
|
||||
fmt.Fprintf(stderr, "startup failed: %v\n", err)
|
||||
return 1
|
||||
}
|
||||
a.logger.Info("migadu_counters starting", "version", version, "config", path)
|
||||
|
||||
apiLn, err := net.Listen("tcp", cfg.Addr)
|
||||
if err != nil {
|
||||
fmt.Fprintf(stderr, "listen %s: %v\n", cfg.Addr, err)
|
||||
return 1
|
||||
}
|
||||
adminLn, err := net.Listen("tcp", cfg.AdminAddr)
|
||||
if err != nil {
|
||||
apiLn.Close()
|
||||
fmt.Fprintf(stderr, "listen %s: %v\n", cfg.AdminAddr, err)
|
||||
return 1
|
||||
}
|
||||
a.Serve(apiLn, adminLn)
|
||||
|
||||
if sigs == nil {
|
||||
ch := make(chan os.Signal, 2)
|
||||
signal.Notify(ch, syscall.SIGHUP, syscall.SIGTERM, syscall.SIGINT)
|
||||
defer signal.Stop(ch)
|
||||
sigs = ch
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case sig := <-sigs:
|
||||
if sig == syscall.SIGHUP {
|
||||
a.logger.Info("SIGHUP received: reopening logs and reloading domains db")
|
||||
if err := a.Reload(); err != nil {
|
||||
a.logger.Warn("reload failed", "error", err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
a.logger.Info("shutting down", "signal", sig.String())
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
err := a.Shutdown(ctx)
|
||||
cancel()
|
||||
if err != nil {
|
||||
fmt.Fprintf(stderr, "shutdown: %v\n", err)
|
||||
}
|
||||
return 0
|
||||
case err := <-a.fatal:
|
||||
a.logger.Error("fatal server error", "error", err)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
_ = a.Shutdown(ctx)
|
||||
cancel()
|
||||
return 1
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// app is the assembled daemon.
|
||||
type app struct {
|
||||
cfg *config.Config
|
||||
logger *slog.Logger
|
||||
logWriter *config.LogWriter
|
||||
txWriter *config.LogWriter // nil when the counting log shares the app log or is off
|
||||
|
||||
db atomic.Pointer[domains.DB]
|
||||
st *store.Store
|
||||
notifier *notify.Notifier
|
||||
reg *prometheus.Registry
|
||||
m *metrics.Metrics
|
||||
srv *server.Server
|
||||
|
||||
apiSrv *http.Server
|
||||
adminSrv *http.Server
|
||||
w *watcher.Watcher
|
||||
fatal chan error
|
||||
reloadMu sync.Mutex
|
||||
}
|
||||
|
||||
// newApp opens everything the daemon needs. A missing CDB is fatal (nothing
|
||||
// could be counted); an unreachable Redis is only a warning (the service must
|
||||
// come up while Redis is briefly away).
|
||||
func newApp(cfg *config.Config) (*app, error) {
|
||||
logger, logWriter, err := config.NewLogger(cfg.Logging)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, w := range cfg.Warnings() {
|
||||
logger.Warn(w)
|
||||
}
|
||||
a := &app{cfg: cfg, logger: logger, logWriter: logWriter, fatal: make(chan error, 2)}
|
||||
|
||||
db, err := domains.OpenWithCache(cfg.DomainsCDB, cfg.MaxCachedDomains)
|
||||
if err != nil {
|
||||
logWriter.Close()
|
||||
return nil, err
|
||||
}
|
||||
a.db.Store(db)
|
||||
logger.Info("domains db loaded", "path", cfg.DomainsCDB, "domains", db.Count(), "max_cached_domains", cfg.MaxCachedDomains)
|
||||
|
||||
a.reg = prometheus.NewRegistry()
|
||||
a.reg.MustRegister(collectors.NewGoCollector(), collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}))
|
||||
a.m = metrics.New(a.reg)
|
||||
a.m.SetCDBDomains(db.Count())
|
||||
|
||||
a.st, err = store.New(cfg.Redis)
|
||||
if err != nil {
|
||||
db.Close()
|
||||
logWriter.Close()
|
||||
return nil, err
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
if err := a.st.Ping(ctx); err != nil {
|
||||
logger.Warn("redis not reachable at startup; commands will be retried on demand", "error", err)
|
||||
} else {
|
||||
mode := "static"
|
||||
if cfg.Redis.SentinelMode() {
|
||||
mode = "sentinel " + cfg.Redis.MasterName
|
||||
}
|
||||
logger.Info("redis reachable", "mode", mode)
|
||||
}
|
||||
cancel()
|
||||
|
||||
var tx *txlog.Logger
|
||||
if cfg.Logging.CountingEnabled() {
|
||||
w := logWriter
|
||||
if cfg.Logging.CountingOutput != cfg.Logging.Output {
|
||||
if a.txWriter, err = config.NewWriter(cfg.Logging.CountingOutput); err != nil {
|
||||
a.st.Close()
|
||||
db.Close()
|
||||
logWriter.Close()
|
||||
return nil, err
|
||||
}
|
||||
w = a.txWriter
|
||||
}
|
||||
tx = txlog.New(w, cfg.Logging.CountingFormat, nil)
|
||||
}
|
||||
|
||||
a.notifier = notify.New(notify.Config{
|
||||
OrgURL: cfg.Notify.OrgURL, AddressURL: cfg.Notify.AddressURL, DomainURL: cfg.Notify.DomainURL,
|
||||
User: cfg.Notify.User, Password: cfg.Notify.Password, Timeout: cfg.Notify.Timeout(),
|
||||
}, server.NotifyHooks(tx, a.m))
|
||||
|
||||
planner := &counting.Planner{
|
||||
Lookup: counting.LookupFunc(func(domain string) (*domains.Domain, error) {
|
||||
db := a.db.Load()
|
||||
if db == nil {
|
||||
return nil, errors.New("domains db not loaded")
|
||||
}
|
||||
return db.Lookup(domain)
|
||||
}),
|
||||
Exempt: cfg.ExemptSet(),
|
||||
CountJunk: cfg.CountJunk,
|
||||
}
|
||||
a.srv = server.New(server.Deps{
|
||||
AuthToken: cfg.AuthToken,
|
||||
AllowedCIDRs: cfg.AllowedPrefixes(),
|
||||
HandlerTimeout: cfg.HandlerTimeoutDuration(),
|
||||
QueueSize: cfg.QueueSize,
|
||||
Workers: cfg.Workers,
|
||||
Exempt: cfg.ExemptSet(),
|
||||
Counter: &counting.Counter{Planner: planner, Store: a.st},
|
||||
Store: a.st,
|
||||
Domains: func() *domains.DB { return a.db.Load() },
|
||||
Notifier: a.notifier,
|
||||
TxLog: tx,
|
||||
Metrics: a.m,
|
||||
Gatherer: a.reg,
|
||||
Logger: logger,
|
||||
})
|
||||
return a, nil
|
||||
}
|
||||
|
||||
// Serve starts the API and admin servers on the listeners and the CDB file
|
||||
// watcher. Listener failures are reported on a.fatal.
|
||||
func (a *app) Serve(apiLn, adminLn net.Listener) {
|
||||
newServer := func(h http.Handler) *http.Server {
|
||||
return &http.Server{
|
||||
Handler: h,
|
||||
ReadHeaderTimeout: time.Duration(a.cfg.ReadHeaderTimeout) * time.Second,
|
||||
WriteTimeout: time.Duration(a.cfg.WriteTimeout) * time.Second,
|
||||
IdleTimeout: time.Duration(a.cfg.IdleTimeout) * time.Second,
|
||||
MaxHeaderBytes: 1 << 20,
|
||||
}
|
||||
}
|
||||
a.apiSrv = newServer(a.srv.Handler())
|
||||
a.adminSrv = newServer(a.srv.AdminHandler())
|
||||
serve := func(name string, s *http.Server, ln net.Listener) {
|
||||
if err := s.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
a.fatal <- fmt.Errorf("%s server: %w", name, err)
|
||||
}
|
||||
}
|
||||
go serve("api", a.apiSrv, apiLn)
|
||||
go serve("admin", a.adminSrv, adminLn)
|
||||
|
||||
w, err := watcher.New([]string{a.cfg.DomainsCDB}, a.reloadDomains)
|
||||
if err != nil {
|
||||
a.logger.Error("domains db watcher not started; reload with SIGHUP", "error", err)
|
||||
} else {
|
||||
w.Start()
|
||||
a.w = w
|
||||
}
|
||||
a.logger.Info("listening", "addr", apiLn.Addr().String(), "admin_addr", adminLn.Addr().String())
|
||||
}
|
||||
|
||||
// Reload is the SIGHUP action: reopen both log files, then reload the CDB.
|
||||
func (a *app) Reload() error {
|
||||
if err := a.logWriter.Reopen(); err != nil {
|
||||
a.logger.Error("reopen log failed", "error", err)
|
||||
}
|
||||
if a.txWriter != nil {
|
||||
if err := a.txWriter.Reopen(); err != nil {
|
||||
a.logger.Error("reopen counting log failed", "error", err)
|
||||
}
|
||||
}
|
||||
return a.reloadDomains()
|
||||
}
|
||||
|
||||
// reloadDomains swaps in a freshly opened CDB. A file that cannot be opened,
|
||||
// or holds no domains at all, is rejected and the current table kept: a
|
||||
// half-deployed file must not silently stop all counting.
|
||||
func (a *app) reloadDomains() error {
|
||||
a.reloadMu.Lock()
|
||||
defer a.reloadMu.Unlock()
|
||||
db, err := domains.OpenWithCache(a.cfg.DomainsCDB, a.cfg.MaxCachedDomains)
|
||||
if err == nil && db.Count() == 0 {
|
||||
db.Close()
|
||||
err = fmt.Errorf("domains db %s holds no records", a.cfg.DomainsCDB)
|
||||
}
|
||||
if err != nil {
|
||||
a.m.RecordReload(false)
|
||||
a.logger.Warn("domains db reload failed, keeping current table", "error", err)
|
||||
return err
|
||||
}
|
||||
old := a.db.Swap(db)
|
||||
a.m.SetCDBDomains(db.Count())
|
||||
a.m.RecordReload(true)
|
||||
a.logger.Info("domains db reloaded", "path", a.cfg.DomainsCDB, "domains", db.Count())
|
||||
if old != nil {
|
||||
// Grace period for lookups still holding the old table.
|
||||
time.AfterFunc(30*time.Second, func() { old.Close() })
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Shutdown stops the listeners, drains the counting queue (every accepted
|
||||
// request is counted), waits for notifications, then closes Redis and the
|
||||
// log files — in that order, so nothing a worker needs disappears under it.
|
||||
func (a *app) Shutdown(ctx context.Context) error {
|
||||
if a.w != nil {
|
||||
a.w.Stop()
|
||||
}
|
||||
var errs []error
|
||||
if a.apiSrv != nil {
|
||||
errs = append(errs, a.apiSrv.Shutdown(ctx))
|
||||
}
|
||||
if a.adminSrv != nil {
|
||||
errs = append(errs, a.adminSrv.Shutdown(ctx))
|
||||
}
|
||||
if pending := a.srv.Pending(); pending > 0 {
|
||||
a.logger.Info("draining counting queue", "pending", pending)
|
||||
}
|
||||
if err := a.srv.Drain(ctx); err != nil {
|
||||
a.logger.Error("counting queue not fully drained", "error", err)
|
||||
errs = append(errs, err)
|
||||
}
|
||||
errs = append(errs, a.notifier.Shutdown(ctx))
|
||||
errs = append(errs, a.st.Close())
|
||||
if db := a.db.Swap(nil); db != nil {
|
||||
errs = append(errs, db.Close())
|
||||
}
|
||||
a.logger.Info("stopped")
|
||||
if a.txWriter != nil {
|
||||
errs = append(errs, a.txWriter.Close())
|
||||
}
|
||||
errs = append(errs, a.logWriter.Close())
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
@@ -0,0 +1,417 @@
|
||||
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())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user