Hammerhead/db.go

package hammerhead

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"git.veillax.com/Veillax/reef/model"
	"git.veillax.com/Veillax/reef/secretbox"
	_ "modernc.org/sqlite" // pure Go, so no C compiler is needed
)

// The configuration is held in memory (it is small and read constantly) and
// persisted to SQLite as one row per object, written incrementally in a
// transaction: only objects that changed since the last save are touched.
// Secrets are encrypted before they reach the file. Logs, traffic history and
// alerts are real tables with indexes, so they can be searched and pruned.

const schema = `
CREATE TABLE IF NOT EXISTS meta (k TEXT PRIMARY KEY, v TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS entities (
	kind TEXT NOT NULL, id TEXT NOT NULL, data TEXT NOT NULL,
	PRIMARY KEY (kind, id)
);
CREATE TABLE IF NOT EXISTS access_log (
	id INTEGER PRIMARY KEY AUTOINCREMENT, ts INTEGER NOT NULL, org_id TEXT NOT NULL DEFAULT '', mapping TEXT NOT NULL DEFAULT '',
	remote TEXT, method TEXT, host TEXT, path TEXT, status INTEGER, usr TEXT, result TEXT,
	bytes_in INTEGER, bytes_out INTEGER, ms INTEGER
);
CREATE INDEX IF NOT EXISTS access_org ON access_log (org_id, id);
CREATE INDEX IF NOT EXISTS access_map ON access_log (org_id, mapping, id);
CREATE INDEX IF NOT EXISTS access_ts ON access_log (ts);
CREATE TABLE IF NOT EXISTS audit_log (
	id INTEGER PRIMARY KEY AUTOINCREMENT, ts INTEGER NOT NULL, org_id TEXT NOT NULL DEFAULT '',
	actor TEXT, action TEXT, detail TEXT
);
CREATE INDEX IF NOT EXISTS audit_org ON audit_log (org_id, id);
CREATE TABLE IF NOT EXISTS traffic_hourly (
	mapping TEXT NOT NULL, hour INTEGER NOT NULL, conns INTEGER NOT NULL DEFAULT 0,
	bytes_in INTEGER NOT NULL DEFAULT 0, bytes_out INTEGER NOT NULL DEFAULT 0,
	PRIMARY KEY (mapping, hour)
);
CREATE TABLE IF NOT EXISTS alerts (
	key TEXT PRIMARY KEY, org_id TEXT NOT NULL DEFAULT '', kind TEXT, name TEXT, message TEXT,
	since INTEGER NOT NULL, fired INTEGER NOT NULL DEFAULT 0, notified INTEGER NOT NULL DEFAULT 0
);
CREATE TABLE IF NOT EXISTS alert_events (
	id INTEGER PRIMARY KEY AUTOINCREMENT, ts INTEGER NOT NULL, org_id TEXT NOT NULL DEFAULT '',
	key TEXT, status TEXT, name TEXT, message TEXT
);
CREATE INDEX IF NOT EXISTS alert_events_org ON alert_events (org_id, id);
` + dmarcSchema

type DB struct {
	sql *sql.DB
	box *secretbox.Box

	mu    sync.Mutex // serialises persist
	cache map[string]map[string]string
	ver   string
}

func openDB(path string, box *secretbox.Box) (*DB, error) {
	conn, err := sql.Open("sqlite", path+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)&_pragma=synchronous(NORMAL)")
	if err != nil {
		return nil, err
	}
	conn.SetMaxOpenConns(8)
	if _, err := conn.Exec(schema); err != nil {
		conn.Close()
		return nil, fmt.Errorf("prepare database: %w", err)
	}
	return &DB{sql: conn, box: box, cache: map[string]map[string]string{}}, nil
}

func (d *DB) Close() error { return d.sql.Close() }

func ms(t time.Time) int64 { return t.UnixMilli() }

// ---------- sealing of secret fields ----------

// sealed returns a copy of v with its secret fields encrypted.
func (d *DB) sealed(v any) any {
	b := d.box
	switch x := v.(type) {
	case *model.Edge:
		c := *x
		c.Token, c.AuthKey = b.Seal(c.Token), b.Seal(c.AuthKey)
		return &c
	case *model.Site:
		c := *x
		c.Token = b.Seal(c.Token)
		return &c
	case *User:
		c := *x
		c.TOTPSecret = b.Seal(c.TOTPSecret)
		return &c
	case *IdP:
		c := *x
		c.ClientSecret = b.Seal(c.ClientSecret)
		return &c
	case *model.MailEgress:
		c := *x
		c.Password = b.Seal(c.Password)
		return &c
	case *model.DKIMKey:
		c := *x
		c.PrivateKey = b.Seal(c.PrivateKey)
		return &c
	case *AlertChannel:
		c := *x
		c.Target, c.Token = b.Seal(c.Target), b.Seal(c.Token)
		return &c
	case Settings:
		x.SMTP.Pass = b.Seal(x.SMTP.Pass)
		return x
	}
	return v
}

func (d *DB) open1(s *string) error {
	v, err := d.box.Open(*s)
	*s = v
	return err
}

// unseal decrypts the secret fields of a loaded object in place.
func (d *DB) unseal(v any) error {
	switch x := v.(type) {
	case *model.Edge:
		return errors.Join(d.open1(&x.Token), d.open1(&x.AuthKey))
	case *model.Site:
		return d.open1(&x.Token)
	case *User:
		return d.open1(&x.TOTPSecret)
	case *IdP:
		return d.open1(&x.ClientSecret)
	case *model.MailEgress:
		return d.open1(&x.Password)
	case *model.DKIMKey:
		return d.open1(&x.PrivateKey)
	case *AlertChannel:
		return errors.Join(d.open1(&x.Target), d.open1(&x.Token))
	case *Settings:
		return d.open1(&x.SMTP.Pass)
	}
	return nil
}

// ---------- configuration ----------

type ent struct {
	kind, id string
	v        any
}

func collect(st *state) []ent {
	var out []ent
	add := func(kind, id string, v any) { out = append(out, ent{kind, id, v}) }
	add("settings", "main", st.Settings)
	for _, x := range st.Users {
		add("user", x.ID, x)
	}
	for _, x := range st.Orgs {
		add("org", x.ID, x)
	}
	for _, x := range st.Roles {
		add("role", x.ID, x)
	}
	for _, x := range st.Members {
		add("member", x.UserID+"|"+x.OrgID, x)
	}
	for _, x := range st.Invites {
		add("invite", x.ID, x)
	}
	for _, x := range st.APIKeys {
		add("apikey", x.ID, x)
	}
	for _, x := range st.IdPs {
		add("idp", x.ID, x)
	}
	for _, x := range st.Edges {
		add("edge", x.ID, x)
	}
	for _, x := range st.Sites {
		add("site", x.ID, x)
	}
	for _, x := range st.Domains {
		add("domain", x.ID, x)
	}
	for _, x := range st.Mappings {
		add("mapping", x.ID, x)
	}
	for _, x := range st.MailEgress {
		add("mailegress", x.ID, x)
	}
	for _, x := range st.DKIM {
		add("dkim", x.ID, x)
	}
	for _, x := range st.Channels {
		add("channel", x.ID, x)
	}
	for _, x := range st.Shares {
		add("share", x.ID, x)
	}
	for k, x := range st.Secrets {
		add("secret", k, x)
	}
	for k, x := range st.Totals {
		add("total", k, x)
	}
	for k, x := range st.Sessions {
		add("session", k, x)
	}
	return out
}

// persist writes whatever changed since the last call, in one transaction.
func (d *DB) persist(st *state) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	ents := collect(st)
	type change struct {
		e     ent
		plain string
	}
	var changes []change
	seen := map[string]map[string]bool{}
	for _, e := range ents {
		b, err := json.Marshal(e.v)
		if err != nil {
			return err
		}
		if seen[e.kind] == nil {
			seen[e.kind] = map[string]bool{}
		}
		seen[e.kind][e.id] = true
		if d.cache[e.kind][e.id] != string(b) {
			changes = append(changes, change{e, string(b)})
		}
	}
	type gone struct{ kind, id string }
	var removed []gone
	for kind, ids := range d.cache {
		for id := range ids {
			if !seen[kind][id] {
				removed = append(removed, gone{kind, id})
			}
		}
	}
	ver := fmt.Sprint(st.Version)
	if len(changes) == 0 && len(removed) == 0 && ver == d.ver {
		return nil
	}
	tx, err := d.sql.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	for _, c := range changes {
		sealed, err := json.Marshal(d.sealed(c.e.v))
		if err != nil {
			return err
		}
		if _, err := tx.Exec(`INSERT INTO entities (kind, id, data) VALUES (?, ?, ?) ON CONFLICT (kind, id) DO UPDATE SET data = excluded.data`,
			c.e.kind, c.e.id, string(sealed)); err != nil {
			return err
		}
	}
	for _, g := range removed {
		if _, err := tx.Exec(`DELETE FROM entities WHERE kind = ? AND id = ?`, g.kind, g.id); err != nil {
			return err
		}
	}
	if ver != d.ver {
		if _, err := tx.Exec(`INSERT INTO meta (k, v) VALUES ('version', ?) ON CONFLICT (k) DO UPDATE SET v = excluded.v`, ver); err != nil {
			return err
		}
	}
	if err := tx.Commit(); err != nil {
		return err
	}
	for _, c := range changes {
		if d.cache[c.e.kind] == nil {
			d.cache[c.e.kind] = map[string]string{}
		}
		d.cache[c.e.kind][c.e.id] = c.plain
	}
	for _, g := range removed {
		delete(d.cache[g.kind], g.id)
	}
	d.ver = ver
	return nil
}

// load reads every stored object into st. It reports whether the database had any.
func (d *DB) load(st *state) (bool, error) {
	d.mu.Lock()
	defer d.mu.Unlock()
	rows, err := d.sql.Query(`SELECT kind, id, data FROM entities`)
	if err != nil {
		return false, err
	}
	defer rows.Close()
	any1 := false
	for rows.Next() {
		var kind, id, data string
		if err := rows.Scan(&kind, &id, &data); err != nil {
			return false, err
		}
		any1 = true
		var target any
		switch kind {
		case "settings":
			target = &st.Settings
		case "user":
			target = new(User)
		case "org":
			target = new(Org)
		case "role":
			target = new(Role)
		case "member":
			target = new(Member)
		case "invite":
			target = new(Invite)
		case "apikey":
			target = new(APIKey)
		case "idp":
			target = new(IdP)
		case "edge":
			target = new(model.Edge)
		case "site":
			target = new(model.Site)
		case "domain":
			target = new(model.Domain)
		case "mapping":
			target = new(model.Mapping)
		case "mailegress":
			target = new(model.MailEgress)
		case "dkim":
			target = new(model.DKIMKey)
		case "channel":
			target = new(AlertChannel)
		case "share":
			target = new(ShareLink)
		case "secret":
			target = new(Secret)
		case "total":
			target = new(Total)
		case "session":
			target = new(Session)
		default:
			continue
		}
		if err := json.Unmarshal([]byte(data), target); err != nil {
			return false, fmt.Errorf("read %s %s: %w", kind, id, err)
		}
		if err := d.unseal(target); err != nil {
			return false, fmt.Errorf("%s %s: %w", kind, id, err)
		}
		switch x := target.(type) {
		case *User:
			st.Users = append(st.Users, x)
		case *Org:
			st.Orgs = append(st.Orgs, x)
		case *Role:
			st.Roles = append(st.Roles, x)
		case *Member:
			st.Members = append(st.Members, x)
		case *Invite:
			st.Invites = append(st.Invites, x)
		case *APIKey:
			st.APIKeys = append(st.APIKeys, x)
		case *IdP:
			st.IdPs = append(st.IdPs, x)
		case *model.Edge:
			st.Edges = append(st.Edges, x)
		case *model.Site:
			st.Sites = append(st.Sites, x)
		case *model.Domain:
			st.Domains = append(st.Domains, x)
		case *model.Mapping:
			st.Mappings = append(st.Mappings, x)
		case *model.MailEgress:
			st.MailEgress = append(st.MailEgress, x)
		case *model.DKIMKey:
			st.DKIM = append(st.DKIM, x)
		case *AlertChannel:
			st.Channels = append(st.Channels, x)
		case *ShareLink:
			st.Shares = append(st.Shares, x)
		case *Secret:
			st.Secrets[id] = x
		case *Total:
			st.Totals[id] = x
		case *Session:
			st.Sessions[id] = x
		}
	}
	if err := rows.Err(); err != nil {
		return false, err
	}
	var v string
	if d.sql.QueryRow(`SELECT v FROM meta WHERE k = 'version'`).Scan(&v) == nil {
		fmt.Sscan(v, &st.Version)
	}
	// remember exactly what is stored, so the next persist only writes differences
	d.cache = map[string]map[string]string{}
	for _, e := range collect(st) {
		b, _ := json.Marshal(e.v)
		if d.cache[e.kind] == nil {
			d.cache[e.kind] = map[string]string{}
		}
		d.cache[e.kind][e.id] = string(b)
	}
	d.ver = fmt.Sprint(st.Version)
	return any1, nil
}

// ---------- access log ----------

type accessFilter struct {
	Org, Mapping, Q, Result string
	StatusClass             int // 2, 3, 4 or 5; 0 = any
	Before                  int64
	Limit                   int
}

func likeEscape(s string) string {
	return strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`).Replace(s)
}

func (d *DB) addAccess(es []model.LogEntry) error {
	if len(es) == 0 {
		return nil
	}
	tx, err := d.sql.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	st, err := tx.Prepare(`INSERT INTO access_log (ts, org_id, mapping, remote, method, host, path, status, usr, result, bytes_in, bytes_out, ms)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`)
	if err != nil {
		return err
	}
	defer st.Close()
	for _, e := range es {
		if e.Time.IsZero() {
			e.Time = time.Now()
		}
		if _, err := st.Exec(ms(e.Time), e.OrgID, e.Mapping, e.Remote, e.Method, e.Host, e.Path, e.Status, e.User, e.Result, e.BytesIn, e.BytesOut, e.Millis); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func (d *DB) queryAccess(f accessFilter) []model.LogEntry {
	q := `SELECT id, ts, mapping, remote, method, host, path, status, usr, result, bytes_in, bytes_out, ms FROM access_log WHERE org_id = ?`
	args := []any{f.Org}
	if f.Mapping != "" {
		q += ` AND mapping = ?`
		args = append(args, f.Mapping)
	}
	if f.Q != "" {
		l := "%" + likeEscape(f.Q) + "%"
		q += ` AND (host LIKE ? ESCAPE '\' OR path LIKE ? ESCAPE '\' OR remote LIKE ? ESCAPE '\' OR usr LIKE ? ESCAPE '\')`
		args = append(args, l, l, l, l)
	}
	if f.StatusClass >= 2 && f.StatusClass <= 5 {
		q += ` AND status >= ? AND status < ?`
		args = append(args, f.StatusClass*100, f.StatusClass*100+100)
	}
	if f.Result != "" {
		q += ` AND result LIKE ? ESCAPE '\'`
		args = append(args, likeEscape(f.Result)+"%")
	}
	if f.Before > 0 {
		q += ` AND id < ?`
		args = append(args, f.Before)
	}
	if f.Limit <= 0 || f.Limit > 1000 {
		f.Limit = 200
	}
	q += ` ORDER BY id DESC LIMIT ?`
	args = append(args, f.Limit)
	rows, err := d.sql.Query(q, args...)
	out := []model.LogEntry{}
	if err != nil {
		return out
	}
	defer rows.Close()
	for rows.Next() {
		var e model.LogEntry
		var ts int64
		var remote, method, host, path, usr, result sql.NullString
		var status sql.NullInt64
		if rows.Scan(&e.ID, &ts, &e.Mapping, &remote, &method, &host, &path, &status, &usr, &result, &e.BytesIn, &e.BytesOut, &e.Millis) != nil {
			continue
		}
		e.Time, e.OrgID = time.UnixMilli(ts), f.Org
		e.Remote, e.Method, e.Host, e.Path, e.User, e.Result = remote.String, method.String, host.String, path.String, usr.String, result.String
		e.Status = int(status.Int64)
		out = append(out, e)
	}
	return out
}

// ---------- audit log ----------

func (d *DB) addAudit(e AuditEntry) {
	d.sql.Exec(`INSERT INTO audit_log (ts, org_id, actor, action, detail) VALUES (?, ?, ?, ?, ?)`, ms(e.Time), e.OrgID, e.Actor, e.Action, e.Detail)
}

func (d *DB) queryAudit(org string, limit int, before int64) []AuditEntry {
	if limit <= 0 || limit > 1000 {
		limit = 200
	}
	q, args := `SELECT id, ts, actor, action, detail FROM audit_log WHERE org_id = ?`, []any{org}
	if before > 0 {
		q += ` AND id < ?`
		args = append(args, before)
	}
	rows, err := d.sql.Query(q+` ORDER BY id DESC LIMIT ?`, append(args, limit)...)
	out := []AuditEntry{}
	if err != nil {
		return out
	}
	defer rows.Close()
	for rows.Next() {
		var e AuditEntry
		var ts int64
		var actor, action, detail sql.NullString
		if rows.Scan(&e.ID, &ts, &actor, &action, &detail) != nil {
			continue
		}
		e.Time, e.OrgID, e.Actor, e.Action, e.Detail = time.UnixMilli(ts), org, actor.String, action.String, detail.String
		out = append(out, e)
	}
	return out
}

// purge deletes logs older than the retention periods.
func (d *DB) purge(accessDays, auditDays int) {
	d.purgeDMARC(time.Now().AddDate(0, 0, -dmarcKeepDays))
	if accessDays > 0 {
		d.sql.Exec(`DELETE FROM access_log WHERE ts < ?`, ms(time.Now().AddDate(0, 0, -accessDays)))
	}
	if auditDays > 0 {
		d.sql.Exec(`DELETE FROM audit_log WHERE ts < ?`, ms(time.Now().AddDate(0, 0, -auditDays)))
	}
	d.sql.Exec(`DELETE FROM traffic_hourly WHERE hour < ?`, time.Now().Add(-31*24*time.Hour).Unix()/3600)
	d.sql.Exec(`DELETE FROM alert_events WHERE ts < ?`, ms(time.Now().AddDate(0, 0, -90)))
}

// ---------- traffic history ----------

func (d *DB) addTraffic(mapping string, st model.Stat) {
	d.sql.Exec(`INSERT INTO traffic_hourly (mapping, hour, conns, bytes_in, bytes_out) VALUES (?, ?, ?, ?, ?)
		ON CONFLICT (mapping, hour) DO UPDATE SET conns = conns + excluded.conns, bytes_in = bytes_in + excluded.bytes_in, bytes_out = bytes_out + excluded.bytes_out`,
		mapping, time.Now().Unix()/3600, st.Conns, st.BytesIn, st.BytesOut)
}

// trafficHistory returns bytes (in + out) per hour for each mapping, oldest first.
func (d *DB) trafficHistory(ids []string, hours int) map[string][]int64 {
	out := map[string][]int64{}
	if hours <= 0 || hours > 24*30 {
		hours = 24
	}
	now := time.Now().Unix() / 3600
	for _, id := range ids {
		out[id] = make([]int64, hours)
	}
	rows, err := d.sql.Query(`SELECT mapping, hour, bytes_in + bytes_out FROM traffic_hourly WHERE hour > ?`, now-int64(hours))
	if err != nil {
		return out
	}
	defer rows.Close()
	for rows.Next() {
		var id string
		var h, n int64
		if rows.Scan(&id, &h, &n) != nil {
			continue
		}
		if series, ok := out[id]; ok {
			if i := int(h - (now - int64(hours) + 1)); i >= 0 && i < hours {
				series[i] += n
			}
		}
	}
	return out
}

// ---------- alerts ----------

func (d *DB) saveAlert(a *alertState) {
	fired := 0
	if a.Fired {
		fired = 1
	}
	d.sql.Exec(`INSERT INTO alerts (key, org_id, kind, name, message, since, fired, notified) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
		ON CONFLICT (key) DO UPDATE SET name = excluded.name, message = excluded.message, fired = excluded.fired, notified = excluded.notified`,
		a.Key, a.OrgID, a.Kind, a.Name, a.Message, ms(a.Since), fired, ms(a.Notified))
}

func (d *DB) deleteAlert(key string) { d.sql.Exec(`DELETE FROM alerts WHERE key = ?`, key) }

func (d *DB) loadAlertRows() []alertState {
	rows, err := d.sql.Query(`SELECT key, org_id, kind, name, message, since, fired, notified FROM alerts`)
	if err != nil {
		return nil
	}
	defer rows.Close()
	var out []alertState
	for rows.Next() {
		var a alertState
		var since, notified, fired int64
		var kind, name, msg sql.NullString
		if rows.Scan(&a.Key, &a.OrgID, &kind, &name, &msg, &since, &fired, &notified) != nil {
			continue
		}
		a.Kind, a.Name, a.Message = kind.String, name.String, msg.String
		a.Since, a.Fired, a.Notified = time.UnixMilli(since), fired == 1, time.UnixMilli(notified)
		out = append(out, a)
	}
	return out
}

type alertEvent struct {
	ID      int64     `json:"id"`
	Time    time.Time `json:"time"`
	Status  string    `json:"status"`
	Name    string    `json:"name"`
	Message string    `json:"message"`
}

func (d *DB) addAlertEvent(t time.Time, org, key, status, name, msg string) {
	d.sql.Exec(`INSERT INTO alert_events (ts, org_id, key, status, name, message) VALUES (?, ?, ?, ?, ?, ?)`, ms(t), org, key, status, name, msg)
}

func (d *DB) queryAlertEvents(org string, limit int) []alertEvent {
	out := []alertEvent{}
	rows, err := d.sql.Query(`SELECT id, ts, status, name, message FROM alert_events WHERE org_id = ? ORDER BY id DESC LIMIT ?`, org, limit)
	if err != nil {
		return out
	}
	defer rows.Close()
	for rows.Next() {
		var e alertEvent
		var ts int64
		var status, name, msg sql.NullString
		if rows.Scan(&e.ID, &ts, &status, &name, &msg) != nil {
			continue
		}
		e.Time, e.Status, e.Name, e.Message = time.UnixMilli(ts), status.String, name.String, msg.String
		out = append(out, e)
	}
	return out
}