atvouch / cli / store.go

package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"github.com/bluesky-social/indigo/atproto/auth/oauth"
	"github.com/bluesky-social/indigo/atproto/syntax"
	_ "github.com/mattn/go-sqlite3"
)

type Store struct {
	db *sql.DB
}

var _ oauth.ClientAuthStore = &Store{}

func NewStore() (*Store, error) {
	configDir, err := os.UserConfigDir()
	if err != nil {
		return nil, fmt.Errorf("getting config dir: %w", err)
	}
	dir := filepath.Join(configDir, "atvouch")
	if err := os.MkdirAll(dir, 0700); err != nil {
		return nil, err
	}

	db, err := sql.Open("sqlite3", filepath.Join(dir, "state.db"))
	if err != nil {
		return nil, err
	}

	for _, pragma := range []string{
		"PRAGMA journal_mode = WAL",
		"PRAGMA busy_timeout = 5000",
		"PRAGMA synchronous = NORMAL",
		"PRAGMA cache_size = -6000",
		"PRAGMA foreign_keys = true",
		"PRAGMA temp_store = memory",
	} {
		if _, err := db.Exec(pragma); err != nil {
			db.Close()
			return nil, fmt.Errorf("setting pragma: %w", err)
		}
	}

	if err := migrate(db); err != nil {
		db.Close()
		return nil, err
	}

	return &Store{db: db}, nil
}

var migrations = []struct {
	version int
	sql     string
}{
	{1, `
		CREATE TABLE sessions (
			did TEXT NOT NULL,
			session_id TEXT NOT NULL,
			data TEXT NOT NULL,
			PRIMARY KEY (did, session_id)
		) STRICT;

		CREATE TABLE auth_requests (
			state TEXT NOT NULL PRIMARY KEY,
			data TEXT NOT NULL
		) STRICT;

		CREATE TABLE active_session (
			id INTEGER NOT NULL PRIMARY KEY CHECK (id = 1),
			did TEXT NOT NULL,
			session_id TEXT NOT NULL
		) STRICT;
	`},
}

func migrate(db *sql.DB) error {
	_, err := db.Exec(`
		CREATE TABLE IF NOT EXISTS migration_log (
			version INTEGER NOT NULL PRIMARY KEY,
			applied_at TEXT NOT NULL DEFAULT (datetime('now'))
		) STRICT;
	`)
	if err != nil {
		return fmt.Errorf("creating migration_log: %w", err)
	}

	for _, m := range migrations {
		var exists int
		err := db.QueryRow("SELECT 1 FROM migration_log WHERE version = ?", m.version).Scan(&exists)
		if err == nil {
			continue
		}
		if err != sql.ErrNoRows {
			return fmt.Errorf("checking migration %d: %w", m.version, err)
		}

		tx, err := db.Begin()
		if err != nil {
			return err
		}
		if _, err := tx.Exec(m.sql); err != nil {
			tx.Rollback()
			return fmt.Errorf("migration %d: %w", m.version, err)
		}
		if _, err := tx.Exec("INSERT INTO migration_log (version) VALUES (?)", m.version); err != nil {
			tx.Rollback()
			return fmt.Errorf("recording migration %d: %w", m.version, err)
		}
		if err := tx.Commit(); err != nil {
			return fmt.Errorf("committing migration %d: %w", m.version, err)
		}
	}

	return nil
}

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

func (s *Store) SetActive(did syntax.DID, sessionID string) error {
	_, err := s.db.Exec(
		`INSERT INTO active_session (id, did, session_id) VALUES (1, ?, ?)
		 ON CONFLICT (id) DO UPDATE SET did = excluded.did, session_id = excluded.session_id`,
		did.String(), sessionID,
	)
	return err
}

type activeSession struct {
	DID       syntax.DID
	SessionID string
}

func (s *Store) GetActive() (*activeSession, error) {
	var didStr, sessionID string
	err := s.db.QueryRow("SELECT did, session_id FROM active_session WHERE id = 1").Scan(&didStr, &sessionID)
	if err == sql.ErrNoRows {
		return nil, fmt.Errorf("no active session (run 'atvouch login' first)")
	}
	if err != nil {
		return nil, err
	}
	did, err := syntax.ParseDID(didStr)
	if err != nil {
		return nil, err
	}
	return &activeSession{DID: did, SessionID: sessionID}, nil
}

func (s *Store) GetSession(ctx context.Context, did syntax.DID, sessionID string) (*oauth.ClientSessionData, error) {
	var data string
	err := s.db.QueryRowContext(ctx,
		"SELECT data FROM sessions WHERE did = ? AND session_id = ?",
		did.String(), sessionID,
	).Scan(&data)
	if err == sql.ErrNoRows {
		return nil, fmt.Errorf("session not found for %s", did)
	}
	if err != nil {
		return nil, err
	}
	var sess oauth.ClientSessionData
	if err := json.Unmarshal([]byte(data), &sess); err != nil {
		return nil, err
	}
	return &sess, nil
}

func (s *Store) SaveSession(ctx context.Context, sess oauth.ClientSessionData) error {
	data, err := json.Marshal(sess)
	if err != nil {
		return err
	}
	_, err = s.db.ExecContext(ctx,
		`INSERT INTO sessions (did, session_id, data) VALUES (?, ?, ?)
		 ON CONFLICT (did, session_id) DO UPDATE SET data = excluded.data`,
		sess.AccountDID.String(), sess.SessionID, string(data),
	)
	return err
}

func (s *Store) DeleteSession(ctx context.Context, did syntax.DID, sessionID string) error {
	_, err := s.db.ExecContext(ctx,
		"DELETE FROM sessions WHERE did = ? AND session_id = ?",
		did.String(), sessionID,
	)
	return err
}

func (s *Store) GetAuthRequestInfo(ctx context.Context, state string) (*oauth.AuthRequestData, error) {
	var data string
	err := s.db.QueryRowContext(ctx,
		"SELECT data FROM auth_requests WHERE state = ?",
		state,
	).Scan(&data)
	if err == sql.ErrNoRows {
		return nil, fmt.Errorf("request info not found: %s", state)
	}
	if err != nil {
		return nil, err
	}
	var req oauth.AuthRequestData
	if err := json.Unmarshal([]byte(data), &req); err != nil {
		return nil, err
	}
	return &req, nil
}

func (s *Store) SaveAuthRequestInfo(ctx context.Context, info oauth.AuthRequestData) error {
	data, err := json.Marshal(info)
	if err != nil {
		return err
	}
	_, err = s.db.ExecContext(ctx,
		"INSERT INTO auth_requests (state, data) VALUES (?, ?)",
		info.State, string(data),
	)
	return err
}

func (s *Store) DeleteAuthRequestInfo(ctx context.Context, state string) error {
	_, err := s.db.ExecContext(ctx,
		"DELETE FROM auth_requests WHERE state = ?",
		state,
	)
	return err
}