350 lines
10 KiB
Go
350 lines
10 KiB
Go
// 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...)
|
|
}
|