counters/internal/server/server.go

// Package server serves the API (/in, /out, /counters) behind the HTTP
// validation layer, and the admin endpoints (/health, /metrics).
package server

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"net/netip"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"

	"migadu_counters/internal/api"
	"migadu_counters/internal/counting"
	"migadu_counters/internal/domains"
	"migadu_counters/internal/metrics"
	"migadu_counters/internal/notify"
	"migadu_counters/internal/store"
	"migadu_counters/internal/txlog"
)

// Deps wires the server to the rest of the daemon.
type Deps struct {
	AuthToken      string
	AllowedCIDRs   []netip.Prefix
	HandlerTimeout time.Duration // bounds the synchronous part (validation) and each queued job
	QueueSize      int           // counting queue capacity (default 10 000)
	Workers        int           // counting workers (default 8)
	Exempt         map[string]bool

	Counter  *counting.Counter
	Store    *store.Store
	Domains  func() *domains.DB // current CDB (swapped on reload); nil = not loaded
	Notifier *notify.Notifier   // nil = never notify
	TxLog    *txlog.Logger      // nil = transaction log disabled
	Metrics  *metrics.Metrics
	Gatherer prometheus.Gatherer // served on /metrics; nil = no endpoint
	Logger   *slog.Logger
	Now      func() time.Time
}

// Server holds the handlers and the counting queue.
type Server struct {
	d     Deps
	queue *counting.Queue

	panicHook func() // tests: injected into the synchronous path
}

// New applies defaults and starts the counting workers.
func New(d Deps) *Server {
	if d.Now == nil {
		d.Now = time.Now
	}
	if d.HandlerTimeout <= 0 {
		d.HandlerTimeout = 4 * time.Second
	}
	d.Logger = d.logger()
	s := &Server{d: d}
	obs := counting.Observer{
		Panic: func(job counting.Job, p any, stack []byte) {
			s.d.Logger.Error("panic while counting", "panic", fmt.Sprint(p), "route", job.Route, "trace_id", job.TraceID, "stack", string(stack))
		},
	}
	if d.Metrics != nil {
		obs.Depth = d.Metrics.SetQueueDepth
		obs.Wait = func(w time.Duration) { d.Metrics.ObserveQueueWait(w.Seconds()) }
	}
	s.queue = counting.NewQueue(counting.QueueConfig{Size: d.QueueSize, Workers: d.Workers, Timeout: d.HandlerTimeout}, s.process, obs)
	return s
}

// Drain stops accepting jobs and waits for the queued ones, bounded by ctx.
// Call it after the HTTP server stopped and before the notifier shuts down.
func (s *Server) Drain(ctx context.Context) error { return s.queue.Drain(ctx) }

// Pending is the number of accepted jobs not yet counted.
func (s *Server) Pending() int { return s.queue.Pending() }

// Depth is the number of jobs waiting for a worker.
func (s *Server) Depth() int { return s.queue.Depth() }

// NotifyHooks wires notification outcomes into the transaction log and the
// metrics; used by the daemon and by tests.
func NotifyHooks(tx *txlog.Logger, m *metrics.Metrics) notify.Hooks {
	return notify.Hooks{
		Sent: func(ev notify.Event, url string, attempt int) {
			if m != nil {
				m.RecordNotification(ev.Level, "sent")
			}
			if tx != nil {
				tx.NotificationSent(ev, url, attempt)
			}
		},
		Failed: func(ev notify.Event, url string, attempts int, err error) {
			if m != nil {
				m.RecordNotification(ev.Level, "failed")
			}
			if tx != nil {
				tx.NotificationFailed(ev, url, attempts, err)
			}
		},
	}
}

// Handler returns the API handler: recover → metrics → timeout → routes.
func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.Handle("/out", s.route("/out", http.MethodPost, s.handleOut))
	mux.Handle("/in", s.route("/in", http.MethodPost, s.handleIn))
	mux.Handle("/counters", s.route("/counters", http.MethodGet, s.handleCounters))
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		writeError(w, http.StatusNotFound, "not found")
	})
	timed := http.TimeoutHandler(mux, s.d.HandlerTimeout, `{"error":"request timeout"}`)
	return s.recoverer(s.observe(timed))
}

// AdminHandler returns /health and /metrics (no auth; bind admin_addr locally).
func (s *Server) AdminHandler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/health", s.handleHealth)
	if s.d.Gatherer != nil {
		mux.Handle("/metrics", promhttp.HandlerFor(s.d.Gatherer, promhttp.HandlerOpts{}))
	}
	return mux
}

func (s *Server) handleOut(w http.ResponseWriter, r *http.Request) {
	var req api.OutRequest
	if e := decodeJSON(r, &req); e != nil {
		s.reject(w, r, "/out", e)
		return
	}
	v, err := api.ValidateOut(req)
	if err != nil {
		s.reject(w, r, "/out", &httpError{http.StatusBadRequest, err.Error()})
		return
	}
	job := counting.Job{
		Route: counting.RouteOut, TraceID: v.TraceID, Received: s.d.Now(),
		AuthUser: v.AuthUser.Value, LoginSuffix: v.LoginSuffix,
		MailFrom: v.MailFrom, MailFromUsable: v.MailFromUsable,
		Recipient: v.Recipient.Value,
	}
	if !v.AuthUser.Usable() {
		job.SenderUnusable = counting.Unusable{Role: "sender", Value: v.AuthUser.Raw, Reason: v.AuthUser.Reason}
	}
	if !v.Recipient.Usable() {
		job.RecipientUnusable = counting.Unusable{Role: "recipient", Value: v.Recipient.Raw, Reason: v.Recipient.Reason}
	}
	s.enqueue(w, r, "/out", job)
}

func (s *Server) handleIn(w http.ResponseWriter, r *http.Request) {
	var req api.InRequest
	if e := decodeJSON(r, &req); e != nil {
		s.reject(w, r, "/in", e)
		return
	}
	v, err := api.ValidateIn(req)
	if err != nil {
		s.reject(w, r, "/in", &httpError{http.StatusBadRequest, err.Error()})
		return
	}
	s.enqueue(w, r, "/in", counting.Job{
		Route: counting.RouteIn, TraceID: v.TraceID, Received: s.d.Now(),
		MailFrom: v.MailFrom, MailFromUsable: v.MailFromUsable, Recipient: v.Recipient, Junk: v.Junk,
	})
}

// enqueue hands the validated job to the workers and answers 202; a full
// queue answers 503 and the drop is counted and logged.
func (s *Server) enqueue(w http.ResponseWriter, r *http.Request, route string, job counting.Job) {
	if s.panicHook != nil {
		s.panicHook()
	}
	if !s.queue.Enqueue(job) {
		if s.d.Metrics != nil {
			s.d.Metrics.RecordDropped(route)
		}
		if s.d.TxLog != nil {
			s.d.TxLog.Dropped(job.Route, job.TraceID)
		}
		s.d.Logger.Error("counting queue full, request dropped", "route", route, "trace_id", job.TraceID, "remote", r.RemoteAddr)
		writeError(w, http.StatusServiceUnavailable, "queue full")
		return
	}
	writeJSON(w, http.StatusAccepted, api.CountResponse{Accepted: true})
}

// process is the worker body: plan, execute, dispatch notifications for
// crossed limits and write the transaction log. ctx is the queue's own
// bounded context, never a request's.
func (s *Server) process(ctx context.Context, job counting.Job) {
	plan, err := s.d.Counter.Planner.Plan(job)
	if err != nil {
		s.d.Logger.Error("domain lookup failed", "route", job.Route, "trace_id", job.TraceID, "error", err)
		if s.d.TxLog != nil {
			s.d.TxLog.Error(job.Route, job.TraceID, err)
		}
		return
	}
	for _, lk := range plan.Lookups {
		if errors.Is(lk.Err, domains.ErrMalformed) {
			s.d.Logger.Warn("malformed row in domains db, domain treated as not hosted", "domain", lk.DomainName, "route", job.Route, "trace_id", job.TraceID, "error", lk.Err)
		}
	}
	out, err := s.d.Counter.Execute(ctx, plan)
	if err != nil {
		if s.d.Metrics != nil {
			s.d.Metrics.RecordRedisError("incr")
		}
		if s.d.TxLog != nil {
			s.d.TxLog.Failed(plan, err)
		}
		s.d.Logger.Error("counting failed", "route", plan.Route, "trace_id", plan.TraceID, "error", err)
		return
	}
	if out.Unknown != nil {
		if s.d.Metrics != nil {
			s.d.Metrics.RecordRedisError("incr_unknown")
		}
		s.d.Logger.Warn("redis outcome unknown, counts re-read from master", "route", plan.Route, "trace_id", plan.TraceID, "error", out.Unknown)
	}

	dispatched := make([]bool, len(out.Results))
	for i, res := range out.Results {
		if res.Err != nil {
			if s.d.Metrics != nil {
				s.d.Metrics.RecordRedisError("incr")
			}
			s.d.Logger.Error("redis error on key", "route", plan.Route, "trace_id", plan.TraceID, "key", res.Key, "error", res.Err)
			continue
		}
		if s.d.Metrics != nil {
			s.d.Metrics.RecordIncrement(res.Level, res.Flow)
		}
		if res.Reached {
			if s.d.Metrics != nil {
				s.d.Metrics.RecordLimitReached(res.Level, res.Flow)
			}
			ev := notify.Event{TraceID: plan.TraceID, Route: plan.Route, Level: res.Level, ID: res.ID,
				Flow: res.Flow, Address: res.By, Count: res.Count, Limit: res.Limit}
			dispatched[i] = s.d.Notifier != nil && s.d.Notifier.Notify(ev)
			if !dispatched[i] && s.d.Metrics != nil {
				s.d.Metrics.RecordNotification(res.Level, "disabled")
			}
		}
	}
	if s.d.TxLog != nil {
		s.d.TxLog.Outcome(out, dispatched)
	}
}

// handleCounters answers GET /counters?address=|domain=[&date=] from the
// read side.
func (s *Server) handleCounters(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	address, domain := strings.TrimSpace(q.Get("address")), strings.ToLower(strings.TrimSpace(q.Get("domain")))
	if (address == "") == (domain == "") {
		s.reject(w, r, "/counters", &httpError{http.StatusBadRequest, "exactly one of address or domain is required"})
		return
	}
	date, err := api.ParseDate(q.Get("date"), s.d.Now())
	if err != nil {
		s.reject(w, r, "/counters", &httpError{http.StatusBadRequest, "date: " + err.Error()})
		return
	}

	resp := api.CountersResponse{Date: date, Levels: []api.CountersLevel{}}
	var local string
	if address != "" {
		addr, err := api.NormalizeEmail(address)
		if err != nil {
			s.reject(w, r, "/counters", &httpError{http.StatusBadRequest, "address: invalid e-mail address"})
			return
		}
		addr = domains.Mailbox(addr)
		resp.Query, resp.Mode, resp.Exempt = addr, "address", s.d.Exempt[addr]
		at := strings.LastIndex(addr, "@")
		local, domain = addr[:at], addr[at+1:]
	} else {
		resp.Query, resp.Mode = domain, "domain"
	}

	db := s.d.Domains()
	if db == nil {
		writeError(w, http.StatusServiceUnavailable, "domains db not loaded")
		return
	}
	d, err := db.Lookup(domain)
	if errors.Is(err, domains.ErrMalformed) {
		s.d.Logger.Warn("malformed row in domains db, domain treated as not hosted", "domain", domain, "route", "/counters", "error", err)
	} else if err != nil && !errors.Is(err, domains.ErrNotFound) {
		writeError(w, http.StatusServiceUnavailable, "domain lookup failed: "+err.Error())
		return
	}
	managed := err == nil
	if managed {
		org := d.OrgID
		resp.Managed, resp.Domain, resp.OrgID = true, d.Name, &org
	}

	type row struct{ level, id string }
	var rows []row
	if resp.Mode == "address" {
		rows = append(rows, row{domains.AddressLevel, resp.Query})
	}
	rows = append(rows, row{domains.DomainLevel, domain})
	if managed {
		rows = append(rows, row{domains.OrganizationLevel, strconv.FormatInt(d.OrgID, 10)})
	}
	keys := make([]string, 0, 2*len(rows))
	for _, rw := range rows {
		keys = append(keys, store.Key(rw.level, rw.id, date, domains.FlowIn), store.Key(rw.level, rw.id, date, domains.FlowOut))
	}
	values, err := s.d.Store.Get(r.Context(), keys)
	if err != nil {
		if s.d.Metrics != nil {
			s.d.Metrics.RecordRedisError("mget")
		}
		writeError(w, http.StatusServiceUnavailable, "redis unavailable: "+err.Error())
		return
	}
	limit := func(level, flow string) *int64 {
		if !managed {
			return nil
		}
		if n, ok := d.Limits.Limit(level, flow, local); ok {
			return &n
		}
		return nil
	}
	for i, rw := range rows {
		resp.Levels = append(resp.Levels, api.CountersLevel{
			Level: rw.level, ID: rw.id,
			In: values[2*i], InLimit: limit(rw.level, domains.FlowIn),
			Out: values[2*i+1], OutLimit: limit(rw.level, domains.FlowOut),
		})
	}
	writeJSON(w, http.StatusOK, resp)
}

func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
	db := s.d.Domains()
	if db == nil {
		writeError(w, http.StatusServiceUnavailable, "domains db not loaded")
		return
	}
	if err := db.Ping(); err != nil {
		writeError(w, http.StatusServiceUnavailable, "domains db not readable: "+err.Error())
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 200*time.Millisecond)
	defer cancel()
	if err := s.d.Store.Ping(ctx); err != nil {
		if s.d.Metrics != nil {
			s.d.Metrics.RecordRedisError("ping")
		}
		writeError(w, http.StatusServiceUnavailable, fmt.Sprintf("redis: %v", err))
		return
	}
	writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}