First commit Go counter service
This commit is contained in:
Dimitrije Dimitrijevic
@@ -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