keeporg/internal/store/store.go

package store

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"sort"
	"sync"
	"time"

	"cc.tiredsysadmin.keeporg/internal/secret"
	"cc.tiredsysadmin.keeporg/internal/sqldb"
)

// Store is the server's metadata. It is held in memory as one State, which
// is what every access decision reads, and persisted in the database as one
// sealed row per object.
//
// Every mutation runs on a private copy of the state; the rows that changed
// are written in one transaction, and the copy is swapped in only once that
// transaction has committed. A reader never observes a half-applied update,
// and a failed one changes nothing.
//
// Several processes may share one database. Each write first checks the
// database's generation: if another process has committed since this one
// last loaded, the state is reloaded and the mutation applied again on top.
// A background check picks up other processes' writes within a couple of
// seconds, so a session revoked through one process stops working in all.
type Store struct {
	db  *sqldb.DB
	key *secret.Keyring

	mu   sync.RWMutex
	s    *State
	gen  int64
	rows map[string][32]byte // row id → digest of the plaintext last stored

	stop chan struct{}
	done chan struct{}
}

// watchInterval is how often the store checks for writes by other processes.
const watchInterval = 2 * time.Second

// Open loads the state from db. key seals and names the rows; it must be the
// key the rows were written with, which sqldb.CheckKey establishes first.
func Open(ctx context.Context, db *sqldb.DB, key *secret.Keyring) (*Store, error) {
	st := &Store{db: db, key: key, stop: make(chan struct{}), done: make(chan struct{})}
	if err := st.reload(ctx); err != nil {
		return nil, err
	}
	go st.watch()
	return st, nil
}

// Close stops watching for other processes' writes. The database stays open;
// it belongs to whoever opened it.
func (st *Store) Close() {
	select {
	case <-st.stop:
	default:
		close(st.stop)
		<-st.done
	}
}

// Read runs fn against a consistent view of the state. fn must not retain or
// mutate anything it is given.
func (st *Store) Read(fn func(*State)) {
	st.mu.RLock()
	defer st.mu.RUnlock()
	fn(st.s)
}

// Empty reports whether nothing has ever been stored.
func (st *Store) Empty() bool {
	st.mu.RLock()
	defer st.mu.RUnlock()
	return len(st.rows) == 0
}

// errStale means another process committed since this one loaded.
var errStale = errors.New("store: stale generation")

// Update applies fn to a copy of the state and commits it. If fn or the write
// fails, the stored state is left untouched. fn may run more than once when
// another process writes at the same moment, so it must only change the
// state it is given and the variables it reports results through.
func (st *Store) Update(fn func(*State) error) error {
	return st.UpdateWith(fn, nil)
}

// UpdateWith is Update with more writes in the same transaction: inTx runs
// after the generation is locked and before the commit, so what it writes
// lands together with the state or not at all, and no other process can
// write in between. Like fn, it may run more than once.
func (st *Store) UpdateWith(fn func(*State) error, inTx func(context.Context, *sql.Tx) error) error {
	st.mu.Lock()
	defer st.mu.Unlock()
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	for attempt := 0; attempt < 5; attempt++ {
		next, err := clone(st.s)
		if err != nil {
			return err
		}
		if err := fn(next); err != nil {
			return err
		}
		rows, err := st.flatten(next)
		if err != nil {
			return err
		}
		digests := make(map[string][32]byte, len(rows))
		var upserts []string
		for id, plain := range rows {
			d := sha256.Sum256(plain)
			digests[id] = d
			if old, ok := st.rows[id]; !ok || old != d {
				upserts = append(upserts, id)
			}
		}
		var deletes []string
		for id := range st.rows {
			if _, ok := rows[id]; !ok {
				deletes = append(deletes, id)
			}
		}
		if len(upserts) == 0 && len(deletes) == 0 && inTx == nil {
			st.s = next
			return nil
		}
		err = st.db.Tx(ctx, func(tx *sql.Tx) error {
			gen, err := st.db.LockGeneration(ctx, tx)
			if err != nil {
				return err
			}
			if gen != st.gen {
				return errStale
			}
			if inTx != nil {
				if err := inTx(ctx, tx); err != nil {
					return err
				}
			}
			for _, id := range upserts {
				if _, err := tx.ExecContext(ctx, st.db.Q(`INSERT INTO keeporg_records (id, body)
					VALUES (?, ?) ON CONFLICT (id) DO UPDATE SET body = excluded.body`),
					id, st.key.Seal([]byte(id), rows[id])); err != nil {
					return err
				}
			}
			for _, id := range deletes {
				if _, err := tx.ExecContext(ctx, st.db.Q(`DELETE FROM keeporg_records WHERE id = ?`), id); err != nil {
					return err
				}
			}
			return st.db.SetGeneration(ctx, tx, gen+1)
		})
		if errors.Is(err, errStale) {
			if err := st.reloadLocked(ctx); err != nil {
				return err
			}
			continue
		}
		if err != nil {
			return fmt.Errorf("store: %w", err)
		}
		st.s, st.rows, st.gen = next, digests, st.gen+1
		return nil
	}
	return errors.New("store: the database kept changing under this update; try again")
}

// GC drops expired sessions, device logins and invites.
func (st *Store) GC(now time.Time) error {
	return st.Update(func(s *State) error {
		for k, v := range s.Sessions {
			if now.After(v.ExpiresAt) {
				delete(s.Sessions, k)
			}
		}
		for k, v := range s.Devices {
			if now.After(v.ExpiresAt) {
				delete(s.Devices, k)
			}
		}
		for k, v := range s.Invites {
			if !v.Usable(now) {
				delete(s.Invites, k)
			}
		}
		return nil
	})
}

// watch reloads the state when another process has written to the database.
func (st *Store) watch() {
	defer close(st.done)
	t := time.NewTicker(watchInterval)
	defer t.Stop()
	for {
		select {
		case <-st.stop:
			return
		case <-t.C:
		}
		ctx, cancel := context.WithTimeout(context.Background(), watchInterval)
		gen, err := st.db.Generation(ctx, st.db)
		st.mu.RLock()
		stale := err == nil && gen != st.gen
		st.mu.RUnlock()
		if stale {
			// A failed reload is retried on the next tick; meanwhile the
			// last good state keeps serving.
			_ = st.reload(ctx)
		}
		cancel()
	}
}

func (st *Store) reload(ctx context.Context) error {
	st.mu.Lock()
	defer st.mu.Unlock()
	return st.reloadLocked(ctx)
}

// reloadLocked reads every row. The generation and the rows are read in one
// transaction, so they describe the same moment.
func (st *Store) reloadLocked(ctx context.Context) error {
	s := newState()
	digests := map[string][32]byte{}
	var gen int64
	err := st.db.Tx(ctx, func(tx *sql.Tx) error {
		var err error
		if gen, err = st.db.Generation(ctx, tx); err != nil {
			return err
		}
		rows, err := tx.QueryContext(ctx, `SELECT id, body FROM keeporg_records`)
		if err != nil {
			return err
		}
		defer rows.Close()
		for rows.Next() {
			var (
				id   string
				body []byte
			)
			if err := rows.Scan(&id, &body); err != nil {
				return err
			}
			plain, err := st.key.Open([]byte(id), body)
			if err != nil {
				return fmt.Errorf("record %s: %w", id, err)
			}
			if err := s.addRecord(plain); err != nil {
				return fmt.Errorf("record %s: %w", id, err)
			}
			digests[id] = sha256.Sum256(plain)
		}
		return rows.Err()
	})
	if err != nil {
		return fmt.Errorf("store: load: %w", err)
	}
	s.sortLists()
	st.s, st.rows, st.gen = s, digests, gen
	return nil
}

// record is the plaintext of one row: what kind of object, and the object.
type record struct {
	Kind string          `json:"k"`
	Data json.RawMessage `json:"d"`
}

const (
	kindUser       = "user"
	kindGroup      = "group"
	kindDatabase   = "database"
	kindInvite     = "invite"
	kindSession    = "session"
	kindDevice     = "device"
	kindMembership = "membership"
	kindGrant      = "grant"
)

// flatten turns a state into rows. A row's id is a keyed digest of its kind
// and natural key, so ids are stable across writes, and a session token or
// an invite token never appears in the database except inside a sealed row.
func (st *Store) flatten(s *State) (map[string][]byte, error) {
	out := map[string][]byte{}
	put := func(kind, key string, v any) error {
		data, err := json.Marshal(v)
		if err != nil {
			return err
		}
		plain, err := json.Marshal(record{Kind: kind, Data: data})
		if err != nil {
			return err
		}
		out[st.key.MAC(kind+"\x00"+key)] = plain
		return nil
	}
	var err error
	for k, v := range s.Users {
		err = errors.Join(err, put(kindUser, k, v))
	}
	for k, v := range s.Groups {
		err = errors.Join(err, put(kindGroup, k, v))
	}
	for k, v := range s.Databases {
		err = errors.Join(err, put(kindDatabase, k, v))
	}
	for k, v := range s.Invites {
		err = errors.Join(err, put(kindInvite, k, v))
	}
	for k, v := range s.Sessions {
		err = errors.Join(err, put(kindSession, k, v))
	}
	for k, v := range s.Devices {
		err = errors.Join(err, put(kindDevice, k, v))
	}
	for _, m := range s.Memberships {
		err = errors.Join(err, put(kindMembership, m.UserID+"\x00"+m.GroupID, m))
	}
	for _, g := range s.Grants {
		err = errors.Join(err, put(kindGrant, g.GroupID+"\x00"+g.DatabaseID, g))
	}
	return out, err
}

// addRecord puts one row's object back where it belongs in the state.
func (s *State) addRecord(plain []byte) error {
	var r record
	if err := json.Unmarshal(plain, &r); err != nil {
		return err
	}
	switch r.Kind {
	case kindUser:
		v := &User{}
		if err := json.Unmarshal(r.Data, v); err != nil {
			return err
		}
		s.Users[v.ID] = v
	case kindGroup:
		v := &Group{}
		if err := json.Unmarshal(r.Data, v); err != nil {
			return err
		}
		s.Groups[v.ID] = v
	case kindDatabase:
		v := &Database{}
		if err := json.Unmarshal(r.Data, v); err != nil {
			return err
		}
		s.Databases[v.ID] = v
	case kindInvite:
		v := &Invite{}
		if err := json.Unmarshal(r.Data, v); err != nil {
			return err
		}
		s.Invites[v.Token] = v
	case kindSession:
		v := &Session{}
		if err := json.Unmarshal(r.Data, v); err != nil {
			return err
		}
		s.Sessions[v.Token] = v
	case kindDevice:
		v := &DeviceAuth{}
		if err := json.Unmarshal(r.Data, v); err != nil {
			return err
		}
		s.Devices[v.DeviceCode] = v
	case kindMembership:
		var v Membership
		if err := json.Unmarshal(r.Data, &v); err != nil {
			return err
		}
		s.Memberships = append(s.Memberships, v)
	case kindGrant:
		var v Grant
		if err := json.Unmarshal(r.Data, &v); err != nil {
			return err
		}
		s.Grants = append(s.Grants, v)
	default:
		return fmt.Errorf("unknown record kind %q", r.Kind)
	}
	return nil
}

// sortLists gives the two lists a stable order, since rows come back in
// whatever order the database likes.
func (s *State) sortLists() {
	sort.Slice(s.Memberships, func(i, j int) bool {
		a, b := s.Memberships[i], s.Memberships[j]
		return a.UserID+"\x00"+a.GroupID < b.UserID+"\x00"+b.GroupID
	})
	sort.Slice(s.Grants, func(i, j int) bool {
		a, b := s.Grants[i], s.Grants[j]
		return a.GroupID+"\x00"+a.DatabaseID < b.GroupID+"\x00"+b.DatabaseID
	})
}

func clone(s *State) (*State, error) {
	b, err := json.Marshal(s)
	if err != nil {
		return nil, err
	}
	out := newState()
	if err := json.Unmarshal(b, out); err != nil {
		return nil, err
	}
	out.init()
	return out, nil
}

// ReadLegacy reads the state.json document that servers before 0.11 kept,
// for importing it into a database.
func ReadLegacy(path string) (*State, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	s := newState()
	if err := json.Unmarshal(b, s); err != nil {
		return nil, fmt.Errorf("store: parse %s: %w", path, err)
	}
	s.init()
	return s, nil
}

// Replace overwrites the whole state with s, for an import.
func (st *Store) Replace(s *State) error {
	return st.Update(func(next *State) error {
		*next = *s
		next.init()
		return nil
	})
}