First commit Go counter service

This commit is contained in:
Dimitrije Dimitrijevic
2026-10-15 02:12:18 +02:00
commit c43ff17627
59 changed files with 11205 additions and 0 deletions
+255
View File
@@ -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)
}
+335
View File
@@ -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)
}
}