lodge/backend/internal/server/store.go

package server

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"strconv"
	"time"

	"lodge/internal/protocol"

	_ "modernc.org/sqlite"
)

const (
	codeLifetime    = 10 * time.Minute
	sessionLifetime = 30 * 24 * time.Hour
	revisionLimit   = 500
	revisionRows    = `
		select r.seq, coalesce(m.name, iif(r.machine_id is null, 'web', 'revoked')), r.author, r.time, r.main_sha, r.main_branch, c.space, c.path, c.base, c.hash
		from revisions r
		join changes c on c.repo = r.repo and c.seq = r.seq
		left join machines m on m.id = r.machine_id`
	schema = `
create table if not exists machines (
	id integer primary key,
	name text not null,
	token_hash text not null unique,
	created_at integer not null
);
create table if not exists codes (
	code text primary key,
	expires_at integer not null
);
create table if not exists sessions (
	token_hash text primary key,
	expires_at integer not null
);
create table if not exists repos (
	key text primary key,
	name text not null,
	seq integer not null,
	time integer not null
);
create table if not exists blobs (
	hash text primary key,
	data blob not null
);
create table if not exists revisions (
	repo text not null,
	seq integer not null,
	machine_id integer,
	author text not null default '',
	time integer not null,
	main_sha text not null,
	main_branch text not null,
	primary key (repo, seq)
);
create table if not exists changes (
	repo text not null,
	seq integer not null,
	space text not null,
	path text not null,
	base text not null,
	hash text not null,
	primary key (repo, seq, space, path)
);
create index if not exists changes_by_path on changes (repo, space, path, seq);
`
)

var (
	ErrConflict = errors.New("state changed since base")
	ErrNotFound = errors.New("not found")
)

type Store struct {
	db *sql.DB
}

func OpenStore(path string) (*Store, error) {
	dsn := "file:" + path + "?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)&_pragma=synchronous(NORMAL)"
	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(1)
	if _, err := db.Exec(schema); err != nil {
		return nil, err
	}
	if err := addAuthor(db); err != nil {
		return nil, err
	}
	return &Store{db: db}, nil
}

func addAuthor(db *sql.DB) error {
	var found int
	if err := db.QueryRow(`select count(*) from pragma_table_info('revisions') where name = 'author'`).Scan(&found); err != nil || found > 0 {
		return err
	}
	_, err := db.Exec(`alter table revisions add column author text not null default ''`)
	return err
}

func (s *Store) Close() error {
	return s.db.Close()
}

func randomHex(bytes int) string {
	buf := make([]byte, bytes)
	if _, err := rand.Read(buf); err != nil {
		panic(err)
	}
	return hex.EncodeToString(buf)
}

func hashToken(token string) string {
	sum := sha256.Sum256([]byte(token))
	return hex.EncodeToString(sum[:])
}

func (s *Store) CreateCode(ctx context.Context) (string, time.Time, error) {
	code, expires := randomHex(12), time.Now().Add(codeLifetime)
	_, err := s.db.ExecContext(ctx, `insert into codes (code, expires_at) values (?, ?)`, code, expires.UnixMilli())
	return code, expires, err
}

func (s *Store) CodeValid(ctx context.Context, code string) bool {
	var found int
	err := s.db.QueryRowContext(ctx, `select 1 from codes where code = ? and expires_at > ?`, code, time.Now().UnixMilli()).Scan(&found)
	return err == nil
}

func (s *Store) Claim(ctx context.Context, code, name string) (string, error) {
	now := time.Now().UnixMilli()
	res, err := s.db.ExecContext(ctx, `delete from codes where code = ? and expires_at > ?`, code, now)
	if err != nil {
		return "", err
	}
	if deleted, _ := res.RowsAffected(); deleted == 0 {
		return "", ErrNotFound
	}
	token := randomHex(32)
	_, err = s.db.ExecContext(ctx, `insert into machines (name, token_hash, created_at) values (?, ?, ?)`, name, hashToken(token), now)
	return token, err
}

func (s *Store) MachineID(ctx context.Context, token string) (int64, error) {
	var id int64
	err := s.db.QueryRowContext(ctx, `select id from machines where token_hash = ?`, hashToken(token)).Scan(&id)
	return id, err
}

func (s *Store) CreateSession(ctx context.Context) (string, time.Time, error) {
	token, expires := randomHex(32), time.Now().Add(sessionLifetime)
	_, err := s.db.ExecContext(ctx, `insert into sessions (token_hash, expires_at) values (?, ?)`, hashToken(token), expires.UnixMilli())
	return token, expires, err
}

func (s *Store) SessionValid(ctx context.Context, token string) bool {
	var found int
	err := s.db.QueryRowContext(ctx, `select 1 from sessions where token_hash = ? and expires_at > ?`, hashToken(token), time.Now().UnixMilli()).Scan(&found)
	return err == nil
}

func (s *Store) DeleteSession(ctx context.Context, token string) error {
	_, err := s.db.ExecContext(ctx, `delete from sessions where token_hash = ? or expires_at <= ?`, hashToken(token), time.Now().UnixMilli())
	return err
}

func (s *Store) Machines(ctx context.Context) ([]protocol.Machine, error) {
	rows, err := s.db.QueryContext(ctx, `
		select m.id, m.name, m.created_at, coalesce(max(r.time), 0), count(r.seq)
		from machines m left join revisions r on r.machine_id = m.id
		group by m.id order by m.created_at`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	machines := []protocol.Machine{}
	for rows.Next() {
		var machine protocol.Machine
		if err := rows.Scan(&machine.ID, &machine.Name, &machine.CreatedAt, &machine.LastPush, &machine.Revisions); err != nil {
			return nil, err
		}
		machines = append(machines, machine)
	}
	return machines, rows.Err()
}

func (s *Store) DeleteMachine(ctx context.Context, id int64) error {
	res, err := s.db.ExecContext(ctx, `delete from machines where id = ?`, id)
	if err != nil {
		return err
	}
	if deleted, _ := res.RowsAffected(); deleted == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *Store) DeleteRepo(ctx context.Context, repo string) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if _, err := tx.ExecContext(ctx, `delete from changes where repo = ?`, repo); err != nil {
		return err
	}
	if _, err := tx.ExecContext(ctx, `delete from revisions where repo = ?`, repo); err != nil {
		return err
	}
	res, err := tx.ExecContext(ctx, `delete from repos where key = ?`, repo)
	if err != nil {
		return err
	}
	if deleted, _ := res.RowsAffected(); deleted == 0 {
		return ErrNotFound
	}
	if _, err := tx.ExecContext(ctx, `delete from blobs where hash not in (select hash from changes)`); err != nil {
		return err
	}
	return tx.Commit()
}

func (s *Store) PutBlob(ctx context.Context, hash string, data []byte) error {
	_, err := s.db.ExecContext(ctx, `insert or ignore into blobs (hash, data) values (?, ?)`, hash, data)
	return err
}

func (s *Store) Blob(ctx context.Context, hash string) ([]byte, error) {
	var data []byte
	err := s.db.QueryRowContext(ctx, `select data from blobs where hash = ?`, hash).Scan(&data)
	if errors.Is(err, sql.ErrNoRows) {
		return nil, ErrNotFound
	}
	return data, err
}

func (s *Store) Repos(ctx context.Context) ([]protocol.Repo, error) {
	rows, err := s.db.QueryContext(ctx, `select key, name, seq, time from repos order by time desc`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	repos := []protocol.Repo{}
	for rows.Next() {
		var repo protocol.Repo
		if err := rows.Scan(&repo.Key, &repo.Name, &repo.Seq, &repo.Time); err != nil {
			return nil, err
		}
		repos = append(repos, repo)
	}
	return repos, rows.Err()
}

func (s *Store) resolveSeq(ctx context.Context, repo, at string) (int64, error) {
	var seq sql.NullInt64
	var err error
	if n, convErr := strconv.ParseInt(at, 10, 64); at != "" && convErr == nil {
		return n, nil
	}
	if at == "" {
		err = s.db.QueryRowContext(ctx, `select max(seq) from revisions where repo = ?`, repo).Scan(&seq)
	} else {
		err = s.db.QueryRowContext(ctx, `select max(seq) from revisions where repo = ? and main_sha like ? || '%'`, repo, at).Scan(&seq)
		if err == nil && !seq.Valid {
			return 0, ErrNotFound
		}
	}
	return seq.Int64, err
}

func (s *Store) State(ctx context.Context, repo, at string) (protocol.State, error) {
	state := protocol.State{Spaces: map[string]protocol.Manifest{}}
	seq, err := s.resolveSeq(ctx, repo, at)
	if err != nil {
		return state, err
	}
	state.Seq = seq
	rows, err := s.db.QueryContext(ctx, `
		select c.space, c.path, c.hash, r.time
		from changes c join revisions r on r.repo = c.repo and r.seq = c.seq
		where c.repo = ?1 and c.hash != '' and c.seq = (
			select max(seq) from changes
			where repo = c.repo and space = c.space and path = c.path and seq <= ?2
		)`, repo, seq)
	if err != nil {
		return state, err
	}
	defer rows.Close()
	for rows.Next() {
		var space, path string
		var file protocol.File
		if err := rows.Scan(&space, &path, &file.Hash, &file.Time); err != nil {
			return state, err
		}
		if state.Spaces[space] == nil {
			state.Spaces[space] = protocol.Manifest{}
		}
		state.Spaces[space][path] = file
	}
	return state, rows.Err()
}

func (s *Store) Revisions(ctx context.Context, repo, path string) ([]protocol.Revision, error) {
	return s.revisions(ctx, revisionRows+`
		where r.repo = ?1 and (?2 = '' or c.path = ?2)
		order by r.seq desc, c.path
		limit ?3`, repo, path, revisionLimit)
}

func (s *Store) RevisionsAfter(ctx context.Context, repo, path string, after int64) ([]protocol.Revision, error) {
	return s.revisions(ctx, revisionRows+`
		where r.repo = ?1 and (?2 = '' or c.path = ?2) and r.seq in (
			select distinct seq from changes
			where repo = ?1 and (?2 = '' or path = ?2) and seq > ?4
			order by seq limit ?3
		)
		order by r.seq, c.path`, repo, path, revisionLimit, after)
}

func (s *Store) revisions(ctx context.Context, query string, args ...any) ([]protocol.Revision, error) {
	rows, err := s.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	revisions := []protocol.Revision{}
	for rows.Next() {
		var rev protocol.Revision
		var change protocol.Change
		if err := rows.Scan(&rev.Seq, &rev.Machine, &rev.Author, &rev.Time, &rev.MainSHA, &rev.MainBranch, &change.Space, &change.Path, &change.Base, &change.Hash); err != nil {
			return nil, err
		}
		if last := len(revisions) - 1; last >= 0 && revisions[last].Seq == rev.Seq {
			revisions[last].Changes = append(revisions[last].Changes, change)
			continue
		}
		rev.Changes = []protocol.Change{change}
		revisions = append(revisions, rev)
	}
	return revisions, rows.Err()
}

func (s *Store) Push(ctx context.Context, repo string, machineID int64, push protocol.Push) (int64, error) {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()

	for _, change := range push.Changes {
		var current string
		err := tx.QueryRowContext(ctx, `select hash from changes where repo = ? and space = ? and path = ? order by seq desc limit 1`, repo, change.Space, change.Path).Scan(&current)
		if err != nil && !errors.Is(err, sql.ErrNoRows) {
			return 0, err
		}
		if current != change.Base {
			return 0, ErrConflict
		}
		if change.Hash == "" {
			continue
		}
		var found int
		if err := tx.QueryRowContext(ctx, `select 1 from blobs where hash = ?`, change.Hash).Scan(&found); err != nil {
			return 0, fmt.Errorf("blob %s was not uploaded: %w", change.Hash, ErrNotFound)
		}
	}

	now := time.Now().UnixMilli()
	var seq int64
	err = tx.QueryRowContext(ctx, `
		insert into repos (key, name, seq, time) values (?, ?, 1, ?)
		on conflict (key) do update set seq = seq + 1, name = excluded.name, time = excluded.time
		returning seq`, repo, push.Name, now).Scan(&seq)
	if err != nil {
		return 0, err
	}
	machine := sql.NullInt64{Int64: machineID, Valid: machineID != 0}
	if _, err := tx.ExecContext(ctx, `insert into revisions (repo, seq, machine_id, author, time, main_sha, main_branch) values (?, ?, ?, ?, ?, ?, ?)`, repo, seq, machine, push.Author, now, push.MainSHA, push.MainBranch); err != nil {
		return 0, err
	}
	for _, change := range push.Changes {
		if _, err := tx.ExecContext(ctx, `insert into changes (repo, seq, space, path, base, hash) values (?, ?, ?, ?, ?, ?)`, repo, seq, change.Space, change.Path, change.Base, change.Hash); err != nil {
			return 0, err
		}
	}
	return seq, tx.Commit()
}