Warp/server/internal/storage/pocketbase/admin.go

package pocketbase

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"

	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// Admin is the PocketBase implementation of storage.AdminStore.
type Admin struct{ *base }

// registrationTokenColumns reads uses_allowed as NULL for a token without a
// limit, as the SQL drivers store it.
const registrationTokenColumns = `token, CASE WHEN has_uses_allowed THEN uses_allowed END,
	completed, expires_at_ms, created_at_ms`

// CreateRegistrationToken implements storage.AdminStore.
func (s *Admin) CreateRegistrationToken(ctx context.Context, rt storage.RegistrationToken) error {
	err := s.write(ctx, func(t *tx) error {
		taken, err := t.exists("admin_registration_tokens", `token = ?`, rt.Token)
		if err != nil {
			return err
		}
		if taken {
			return storage.ErrRegistrationTokenExists
		}
		row := rec{
			"token": rt.Token, "completed": rt.Completed,
			"expires_at_ms": rt.ExpiresAtMS, "created_at_ms": rt.CreatedAtMS,
		}
		setUsesAllowed(row, rt.UsesAllowed)
		t.put("admin_registration_tokens", row)
		return nil
	})
	switch {
	case errors.Is(err, storage.ErrRegistrationTokenExists):
		return err
	case err != nil:
		return fmt.Errorf("create a registration token: %w", err)
	}
	return nil
}

// setUsesAllowed writes an optional use limit into a row.
func setUsesAllowed(row rec, n *int) {
	if n == nil {
		row["uses_allowed"], row["has_uses_allowed"] = int64(0), false
		return
	}
	row["uses_allowed"], row["has_uses_allowed"] = int64(*n), true
}

// RegistrationToken implements storage.AdminStore.
func (s *Admin) RegistrationToken(ctx context.Context, token string) (storage.RegistrationToken, error) {
	rt, err := scanRegistrationToken(s.db.QueryRowContext(ctx, `SELECT `+registrationTokenColumns+`
		FROM admin_registration_tokens WHERE token = ?`, token).Scan)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.RegistrationToken{}, storage.ErrNoSuchRegistrationToken
	case err != nil:
		return storage.RegistrationToken{}, fmt.Errorf("read a registration token: %w", err)
	}
	return rt, nil
}

// RegistrationTokens implements storage.AdminStore.
func (s *Admin) RegistrationTokens(ctx context.Context) ([]storage.RegistrationToken, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT `+registrationTokenColumns+`
		FROM admin_registration_tokens ORDER BY token`)
	if err != nil {
		return nil, fmt.Errorf("list registration tokens: %w", err)
	}
	defer func() { _ = rows.Close() }()

	out := []storage.RegistrationToken{}
	for rows.Next() {
		rt, err := scanRegistrationToken(rows.Scan)
		if err != nil {
			return nil, fmt.Errorf("list registration tokens: %w", err)
		}
		out = append(out, rt)
	}
	return out, rows.Err()
}

// UpdateRegistrationToken implements storage.AdminStore.
func (s *Admin) UpdateRegistrationToken(ctx context.Context, token string, usesAllowed *int, expiresAtMS int64) error {
	err := s.write(ctx, func(t *tx) error {
		row := rec{"expires_at_ms": expiresAtMS}
		setUsesAllowed(row, usesAllowed)
		n, err := t.setWhere("admin_registration_tokens", row, `token = ?`, token)
		if err != nil {
			return fmt.Errorf("update a registration token: %w", err)
		}
		if n == 0 {
			return storage.ErrNoSuchRegistrationToken
		}
		return nil
	})
	return err
}

// DeleteRegistrationToken implements storage.AdminStore.
func (s *Admin) DeleteRegistrationToken(ctx context.Context, token string) error {
	return s.write(ctx, func(t *tx) error {
		found, err := t.delKey("admin_registration_tokens", token)
		if err != nil {
			return fmt.Errorf("delete a registration token: %w", err)
		}
		if !found {
			return storage.ErrNoSuchRegistrationToken
		}
		return nil
	})
}

// UseRegistrationToken implements storage.AdminStore. The check and the
// spend happen under the write lock, so two registrations cannot both take
// the last use.
func (s *Admin) UseRegistrationToken(ctx context.Context, token string, nowMS int64) error {
	return s.write(ctx, func(t *tx) error {
		var completed int64
		err := t.queryRow(`SELECT completed FROM admin_registration_tokens
			WHERE token = ?
			AND (NOT has_uses_allowed OR completed < uses_allowed)
			AND (expires_at_ms = 0 OR expires_at_ms > ?)`, token, nowMS).Scan(&completed)
		switch {
		case errors.Is(err, sql.ErrNoRows):
			return storage.ErrNoSuchRegistrationToken
		case err != nil:
			return fmt.Errorf("spend a registration token: %w", err)
		}
		t.put("admin_registration_tokens", rec{"token": token, "completed": completed + 1})
		return nil
	})
}

// ReleaseRegistrationToken implements storage.AdminStore. A token deleted in
// the meantime has nothing to give back to, which is not an error.
func (s *Admin) ReleaseRegistrationToken(ctx context.Context, token string) error {
	err := s.write(ctx, func(t *tx) error {
		var completed int64
		err := t.queryRow(`SELECT completed FROM admin_registration_tokens WHERE token = ? AND completed > 0`, token).
			Scan(&completed)
		switch {
		case errors.Is(err, sql.ErrNoRows):
			return nil
		case err != nil:
			return err
		}
		t.put("admin_registration_tokens", rec{"token": token, "completed": completed - 1})
		return nil
	})
	if err != nil {
		return fmt.Errorf("release a registration token: %w", err)
	}
	return nil
}

func scanRegistrationToken(scan func(...any) error) (storage.RegistrationToken, error) {
	var (
		rt   storage.RegistrationToken
		uses sql.NullInt64
	)
	if err := scan(&rt.Token, &uses, &rt.Completed, &rt.ExpiresAtMS, &rt.CreatedAtMS); err != nil {
		return storage.RegistrationToken{}, err
	}
	if uses.Valid {
		n := int(uses.Int64)
		rt.UsesAllowed = &n
	}
	return rt, nil
}

const (
	auditColumns = `audit_id, ts_ms, actor, action, target, details_json`

	defaultAuditPage = 50
	maxAuditPage     = 1000
)

// AppendAudit implements storage.AdminStore.
func (s *Admin) AppendAudit(ctx context.Context, e storage.AuditEntry) (int64, error) {
	details := string(e.Details)
	if details == "" {
		details = "{}"
	}
	var id int64
	err := s.write(ctx, func(t *tx) error {
		var err error
		if id, err = t.next(adminSequences, seqAudit); err != nil {
			return err
		}
		t.put("admin_audit_log", rec{
			"audit_id": id, "ts_ms": e.TSMS, "actor": e.Actor,
			"action": e.Action, "target": e.Target, "details_json": details,
		})
		return nil
	})
	if err != nil {
		return 0, fmt.Errorf("append to the audit log: %w", err)
	}
	return id, nil
}

// AuditLog implements storage.AdminStore.
func (s *Admin) AuditLog(ctx context.Context, f storage.AuditFilter) ([]storage.AuditEntry, error) {
	var (
		conditions []string
		args       []any
	)
	if f.Before > 0 {
		conditions = append(conditions, `audit_id < ?`)
		args = append(args, f.Before)
	}
	if f.Actor != "" {
		conditions = append(conditions, `actor = ?`)
		args = append(args, f.Actor)
	}
	if f.Action != "" {
		conditions = append(conditions, `action = ?`)
		args = append(args, f.Action)
	}
	limit := f.Limit
	switch {
	case limit <= 0:
		limit = defaultAuditPage
	case limit > maxAuditPage:
		limit = maxAuditPage
	}
	args = append(args, limit)

	query := `SELECT ` + auditColumns + ` FROM admin_audit_log`
	if len(conditions) > 0 {
		query += ` WHERE ` + strings.Join(conditions, " AND ")
	}
	query += ` ORDER BY audit_id DESC LIMIT ?`

	rows, err := s.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("read the audit log: %w", err)
	}
	defer func() { _ = rows.Close() }()

	out := []storage.AuditEntry{}
	for rows.Next() {
		var (
			e       storage.AuditEntry
			details string
		)
		if err := rows.Scan(&e.ID, &e.TSMS, &e.Actor, &e.Action, &e.Target, &details); err != nil {
			return nil, fmt.Errorf("read the audit log: %w", err)
		}
		e.Details = []byte(details)
		out = append(out, e)
	}
	return out, rows.Err()
}

// ServerNoticeRoom implements storage.AdminStore.
func (s *Admin) ServerNoticeRoom(ctx context.Context, localpart string) (string, error) {
	var roomID string
	err := s.db.QueryRowContext(ctx, `SELECT room_id FROM admin_server_notice_rooms WHERE localpart = ?`, localpart).
		Scan(&roomID)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return "", storage.ErrNoServerNoticeRoom
	case err != nil:
		return "", fmt.Errorf("read the server notice room of %q: %w", localpart, err)
	}
	return roomID, nil
}

// SetServerNoticeRoom implements storage.AdminStore.
func (s *Admin) SetServerNoticeRoom(ctx context.Context, localpart, roomID string) error {
	err := s.write(ctx, func(t *tx) error {
		t.put("admin_server_notice_rooms", rec{"localpart": localpart, "room_id": roomID})
		return nil
	})
	if err != nil {
		return fmt.Errorf("record the server notice room of %q: %w", localpart, err)
	}
	return nil
}

// AppServices is the PocketBase implementation of storage.AppServiceStore.
type AppServices struct{ *base }

// AppServiceState implements storage.AppServiceStore.
func (s *AppServices) AppServiceState(ctx context.Context, id string, startPos int64) (storage.AppServiceState, error) {
	var st storage.AppServiceState
	err := s.write(ctx, func(t *tx) error {
		err := t.queryRow(`SELECT as_id, stream_pos, txn_id, txn_to FROM appservice_state WHERE as_id = ?`, id).
			Scan(&st.ID, &st.StreamPos, &st.TxnID, &st.TxnTo)
		switch {
		case errors.Is(err, sql.ErrNoRows):
			st = storage.AppServiceState{ID: id, StreamPos: startPos}
			t.put("appservice_state", rec{"as_id": id, "stream_pos": startPos, "txn_id": int64(0), "txn_to": int64(0)})
			return nil
		case err != nil:
			return err
		}
		return nil
	})
	if err != nil {
		return storage.AppServiceState{}, fmt.Errorf("read the cursor of application service %q: %w", id, err)
	}
	return st, nil
}

// BeginTransaction implements storage.AppServiceStore.
func (s *AppServices) BeginTransaction(ctx context.Context, id string, txnID, to int64) error {
	err := s.write(ctx, func(t *tx) error {
		_, err := t.setWhere("appservice_state", rec{"txn_id": txnID, "txn_to": to}, `as_id = ?`, id)
		return err
	})
	if err != nil {
		return fmt.Errorf("begin transaction %d of application service %q: %w", txnID, id, err)
	}
	return nil
}

// CompleteTransaction implements storage.AppServiceStore.
func (s *AppServices) CompleteTransaction(ctx context.Context, id string, to int64) error {
	err := s.write(ctx, func(t *tx) error {
		_, err := t.setWhere("appservice_state", rec{"stream_pos": to, "txn_to": int64(0)}, `as_id = ?`, id)
		return err
	})
	if err != nil {
		return fmt.Errorf("advance the cursor of application service %q: %w", id, err)
	}
	return nil
}

var (
	_ storage.AdminStore      = (*Admin)(nil)
	_ storage.AppServiceStore = (*AppServices)(nil)
)