misaki/db/db.go

package db

import (
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"errors"
	"strings"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

type DB struct{ conn *sql.DB }

type User struct {
	ID        int64
	Username  string
	PassHash  string
	Note      string
	Status    string
	IsAdmin   bool
	CreatedAt time.Time
}

type Session struct {
	Token     string
	UserID    int64
	ExpiresAt time.Time
	AMR       string
}

type Service struct {
	Name    string
	Label   string
	Enabled bool
}

type UserService struct {
	UserID     int64
	Service    string
	Label      string
	Granted    bool
	Applied    bool
	Denied     bool
	Data       string
	DenyReason string
}

type UserServiceView struct {
	Service    string
	Label      string
	Applied    bool
	Granted    bool
	Denied     bool
	Data       string
	DenyReason string
}

var (
	ErrAlreadyApplied = errors.New("already applied")
	ErrNameTaken      = errors.New("name taken")
)

func Open(path string) (*DB, error) {
	conn, err := sql.Open("sqlite3", path+"?_journal=WAL&_foreign_keys=on&_busy_timeout=5000")
	if err != nil {
		return nil, err
	}
	d := &DB{conn}
	return d, d.migrate()
}

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

func (d *DB) migrate() error {
	_, err := d.conn.Exec(`
        CREATE TABLE IF NOT EXISTS users (
            id         INTEGER PRIMARY KEY,
            username   TEXT UNIQUE NOT NULL,
            pass_hash  TEXT NOT NULL,
            note       TEXT DEFAULT '',
            status     TEXT DEFAULT 'pending',
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP
        );
        CREATE TABLE IF NOT EXISTS sessions (
            token      TEXT PRIMARY KEY,
            user_id    INTEGER NOT NULL REFERENCES users(id),
            expires_at DATETIME NOT NULL
        );
        CREATE TABLE IF NOT EXISTS services (
            name    TEXT PRIMARY KEY,
            label   TEXT NOT NULL,
            enabled INTEGER NOT NULL DEFAULT 0
        );
        CREATE TABLE IF NOT EXISTS user_services (
            user_id         INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
            service         TEXT NOT NULL,
            enabled         INTEGER NOT NULL DEFAULT 0,
            provisioned     INTEGER NOT NULL DEFAULT 0,
            provision_error TEXT NOT NULL DEFAULT '',
            provisioned_at  DATETIME,
            PRIMARY KEY (user_id, service)
        );
        CREATE TABLE IF NOT EXISTS recovery_codes (
            user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
            hash    TEXT NOT NULL,
            PRIMARY KEY (user_id, hash)
        );
        CREATE TABLE IF NOT EXISTS former_names (
            name       TEXT PRIMARY KEY COLLATE NOCASE,
            user_id    INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
            renamed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
        );
        CREATE TABLE IF NOT EXISTS reset_tokens (
            hash       TEXT PRIMARY KEY,
            user_id    INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
            expires_at DATETIME NOT NULL
        );
    `)
	if err != nil {
		return err
	}
	for _, stmt := range []string{
		`ALTER TABLE users ADD COLUMN is_admin INTEGER NOT NULL DEFAULT 0`,
		`ALTER TABLE users ADD COLUMN email TEXT NOT NULL DEFAULT ''`,
		`ALTER TABLE users ADD COLUMN totp_secret TEXT NOT NULL DEFAULT ''`,
		`ALTER TABLE users ADD COLUMN totp_last INTEGER NOT NULL DEFAULT 0`,
		`ALTER TABLE sessions ADD COLUMN amr TEXT NOT NULL DEFAULT 'pwd'`,
		`ALTER TABLE user_services ADD COLUMN deny_reason TEXT NOT NULL DEFAULT ''`,
	} {
		if _, err := d.conn.Exec(stmt); err != nil && !strings.Contains(err.Error(), "duplicate column name") {
			return err
		}
	}

	if err := d.hashSessionTokens(); err != nil {
		return err
	}

	_, err = d.conn.Exec(`
        INSERT INTO services (name, label, enabled) VALUES
            ('werc',       'Werc',        0),
            ('transfer',   'Transfer.sh', 0),
            ('akkoma',     'Akkoma',      0),
            ('pixelfed',   'Pixelfed',    0),
            ('ergo',       'IRC (Ergo)',  0),
            ('gancio',     'Gancio',      0),
            ('mumble',     'Mumble',      0),
            ('mail',       'Email',       0),
            ('cerca',      'Cerca',       0),
            ('prosody',    'XMPP',        0),
            ('forgejo',    'Forgejo',     0),
            ('streaming',  'Streaming',   0),
            ('radio',      'Radio',       0),
            ('files',      'Files',       0),
            ('notes',      'HedgeDoc',    0),
            ('paste',      'PrivateBin',  0),
            ('rss',        'RSS Reader',  0)
        ON CONFLICT(name) DO UPDATE SET label=excluded.label
    `)
	return err
}

func (d *DB) CreateUser(username, hash, note string) (int64, error) {
	r, err := d.conn.Exec(
		`INSERT INTO users (username, pass_hash, note) VALUES (?,?,?)`,
		username, hash, note)
	if err != nil {
		return 0, err
	}
	return r.LastInsertId()
}

func (d *DB) GetUserByUsername(username string) (*User, error) {
	u := &User{}
	var isAdmin int
	err := d.conn.QueryRow(
		`SELECT id, username, pass_hash, note, status, is_admin, created_at FROM users WHERE username=?`,
		username).Scan(&u.ID, &u.Username, &u.PassHash, &u.Note, &u.Status, &isAdmin, &u.CreatedAt)
	u.IsAdmin = isAdmin == 1
	if err != nil {
		return nil, err
	}
	return u, nil
}

func (d *DB) GetUserByID(id int64) (*User, error) {
	u := &User{}
	var isAdmin int
	err := d.conn.QueryRow(
		`SELECT id, username, pass_hash, note, status, is_admin, created_at FROM users WHERE id=?`,
		id).Scan(&u.ID, &u.Username, &u.PassHash, &u.Note, &u.Status, &isAdmin, &u.CreatedAt)
	u.IsAdmin = isAdmin == 1
	if err != nil {
		return nil, err
	}
	return u, nil
}

func (d *DB) GetUserEmail(id int64) (string, error) {
	var e string
	err := d.conn.QueryRow(`SELECT email FROM users WHERE id=?`, id).Scan(&e)
	return e, err
}

func (d *DB) SetUserEmail(id int64, email string) error {
	_, err := d.conn.Exec(`UPDATE users SET email=? WHERE id=?`, email, id)
	return err
}

func (d *DB) UpdateUserNote(id int64, note string) error {
	_, err := d.conn.Exec(`UPDATE users SET note=? WHERE id=?`, note, id)
	return err
}

func (d *DB) SetUserStatus(id int64, status string) error {
	_, err := d.conn.Exec(`UPDATE users SET status=? WHERE id=?`, status, id)
	return err
}

func (d *DB) ListUsers(status string) ([]*User, error) {
	rows, err := d.conn.Query(
		`SELECT id, username, note, status, is_admin, created_at FROM users WHERE status=? ORDER BY created_at`,
		status)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var users []*User
	for rows.Next() {
		u := &User{}
		var isAdmin int
		if err := rows.Scan(&u.ID, &u.Username, &u.Note, &u.Status, &isAdmin, &u.CreatedAt); err != nil {
			return nil, err
		}
		u.IsAdmin = isAdmin == 1
		users = append(users, u)
	}
	return users, rows.Err()
}

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

func (d *DB) hashSessionTokens() error {
	var version int
	if err := d.conn.QueryRow(`PRAGMA user_version`).Scan(&version); err != nil {
		return err
	}
	if version >= 1 {
		return nil
	}
	tx, err := d.conn.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	rows, err := tx.Query(`SELECT token FROM sessions`)
	if err != nil {
		return err
	}
	var tokens []string
	for rows.Next() {
		var tok string
		if err := rows.Scan(&tok); err != nil {
			rows.Close()
			return err
		}
		tokens = append(tokens, tok)
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return err
	}
	for _, tok := range tokens {
		if _, err := tx.Exec(`UPDATE sessions SET token=? WHERE token=?`, hashToken(tok), tok); err != nil {
			return err
		}
	}
	if _, err := tx.Exec(`PRAGMA user_version = 1`); err != nil {
		return err
	}
	return tx.Commit()
}

func (d *DB) CreateSession(token string, userID int64, expires time.Time, amr string) error {
	_, err := d.conn.Exec(
		`INSERT INTO sessions (token, user_id, expires_at, amr) VALUES (?,?,?,?)`,
		hashToken(token), userID, expires, amr)
	return err
}

func (d *DB) GetSession(token string) (*Session, error) {
	s := &Session{}
	err := d.conn.QueryRow(
		`SELECT token, user_id, expires_at, amr FROM sessions WHERE token=? AND expires_at > ?`,
		hashToken(token), time.Now()).Scan(&s.Token, &s.UserID, &s.ExpiresAt, &s.AMR)
	if err != nil {
		return nil, err
	}
	return s, nil
}

func (d *DB) DeleteSession(token string) error {
	_, err := d.conn.Exec(`DELETE FROM sessions WHERE token=?`, hashToken(token))
	return err
}

func (d *DB) PruneExpiredSessions() error {
	now := time.Now()
	if _, err := d.conn.Exec(`DELETE FROM reset_tokens WHERE expires_at <= ?`, now); err != nil {
		return err
	}
	_, err := d.conn.Exec(`DELETE FROM sessions WHERE expires_at <= ?`, now)
	return err
}

// DeleteUserSessions logs a user out everywhere except the session keep ("" = all).
func (d *DB) DeleteUserSessions(userID int64, keep string) error {
	_, err := d.conn.Exec(`DELETE FROM sessions WHERE user_id=? AND token != ?`, userID, hashToken(keep))
	return err
}

func (d *DB) ListServices() ([]*Service, error) {
	rows, err := d.conn.Query(`SELECT name, label, enabled FROM services ORDER BY name`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var svcs []*Service
	for rows.Next() {
		s := &Service{}
		var enabled int
		if err := rows.Scan(&s.Name, &s.Label, &enabled); err != nil {
			return nil, err
		}
		s.Enabled = enabled == 1
		svcs = append(svcs, s)
	}
	return svcs, rows.Err()
}

func (d *DB) SetServiceEnabled(name string, on bool) error {
	val := 0
	if on {
		val = 1
	}
	_, err := d.conn.Exec(`UPDATE services SET enabled=? WHERE name=?`, val, name)
	return err
}

func (d *DB) GetServicesForUser(userID int64) ([]*UserServiceView, error) {
	rows, err := d.conn.Query(`
        SELECT s.name, s.label,
               COALESCE(us.provisioned, 0),
               COALESCE(us.enabled, 0),
               COALESCE(us.provision_error, ''),
               COALESCE(us.deny_reason, '')
        FROM services s
        LEFT JOIN user_services us ON us.service = s.name AND us.user_id = ?
        WHERE s.enabled = 1
          AND s.name != 'mumble'
        ORDER BY s.name
    `, userID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var svcs []*UserServiceView
	for rows.Next() {
		v := &UserServiceView{}
		var applied, granted int
		if err := rows.Scan(&v.Service, &v.Label, &applied, &granted, &v.Data, &v.DenyReason); err != nil {
			return nil, err
		}
		v.Applied = applied == 1
		v.Denied = applied == 2
		v.Granted = granted == 1
		svcs = append(svcs, v)
	}
	return svcs, rows.Err()
}

func (d *DB) ApplyForService(userID int64, service string) error {
	var applied int
	err := d.conn.QueryRow(
		`SELECT provisioned FROM user_services WHERE user_id=? AND service=?`,
		userID, service).Scan(&applied)
	if err != nil && err != sql.ErrNoRows {
		return err
	}
	if applied == 1 {
		return ErrAlreadyApplied
	}

	_, err = d.conn.Exec(`
        INSERT INTO user_services (user_id, service, provisioned, provisioned_at)
        VALUES (?, ?, 1, CURRENT_TIMESTAMP)
        ON CONFLICT(user_id, service) DO UPDATE
            SET provisioned=1, provisioned_at=CURRENT_TIMESTAMP, deny_reason=''
    `, userID, service)
	return err
}

// DenyService turns an open request into a denial; provisioned=2 marks it.
func (d *DB) DenyService(userID int64, service, reason string) error {
	_, err := d.conn.Exec(`
        UPDATE user_services SET provisioned=2, deny_reason=?
        WHERE user_id=? AND service=? AND provisioned=1 AND enabled=0
    `, reason, userID, service)
	return err
}

func (d *DB) SetUserServiceData(userID int64, service, data string) error {
	_, err := d.conn.Exec(`
        INSERT INTO user_services (user_id, service, provision_error)
        VALUES (?, ?, ?)
        ON CONFLICT(user_id, service) DO UPDATE SET provision_error=excluded.provision_error
    `, userID, service, data)
	return err
}

func (d *DB) ListUserServices(userID int64) ([]*UserService, error) {
	rows, err := d.conn.Query(`
        SELECT us.user_id, us.service, s.label, us.enabled, us.provisioned, us.provision_error, us.deny_reason
        FROM user_services us
        JOIN services s ON s.name = us.service
        WHERE us.user_id=?
        ORDER BY us.service
    `, userID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var svcs []*UserService
	for rows.Next() {
		us := &UserService{}
		var granted, applied int
		if err := rows.Scan(&us.UserID, &us.Service, &us.Label, &granted, &applied, &us.Data, &us.DenyReason); err != nil {
			return nil, err
		}
		us.Granted = granted == 1
		us.Applied = applied == 1
		us.Denied = applied == 2
		svcs = append(svcs, us)
	}
	return svcs, rows.Err()
}

func (d *DB) SetUserServiceEnabled(userID int64, service string, on bool) error {
	val := 0
	if on {
		val = 1
	}
	_, err := d.conn.Exec(`
        INSERT INTO user_services (user_id, service, enabled)
        VALUES (?, ?, ?)
        ON CONFLICT(user_id, service) DO UPDATE SET enabled=excluded.enabled,
            provisioned=CASE WHEN excluded.enabled=1 AND provisioned=2 THEN 1 ELSE provisioned END,
            deny_reason=CASE WHEN excluded.enabled=1 THEN '' ELSE deny_reason END
    `, userID, service, val)
	return err
}

func (d *DB) GetUserService(userID int64, service string) (enabled bool, data string, err error) {
	var en int
	err = d.conn.QueryRow(
		`SELECT COALESCE(enabled,0), COALESCE(provision_error,'')
         FROM user_services WHERE user_id=? AND service=?`,
		userID, service).Scan(&en, &data)
	if err == sql.ErrNoRows {
		return false, "", nil
	}
	return en == 1, data, err
}

func (d *DB) GetUserByUsernameFold(username string) (*User, error) {
	u := &User{}
	var isAdmin int
	err := d.conn.QueryRow(
		`SELECT id, username, pass_hash, note, status, is_admin, created_at FROM users
         WHERE lower(username)=lower(?) ORDER BY id LIMIT 1`,
		username).Scan(&u.ID, &u.Username, &u.PassHash, &u.Note, &u.Status, &isAdmin, &u.CreatedAt)
	u.IsAdmin = isAdmin == 1
	if err != nil {
		return nil, err
	}
	return u, nil
}

func (d *DB) ListServiceUsernames(service string) ([]string, error) {
	rows, err := d.conn.Query(`
        SELECT u.username FROM users u
        JOIN user_services us ON us.user_id = u.id AND us.service = ?
        WHERE u.status = 'approved' AND us.enabled = 1
        ORDER BY u.username
    `, service)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	names := []string{}
	for rows.Next() {
		var n string
		if err := rows.Scan(&n); err != nil {
			return nil, err
		}
		names = append(names, n)
	}
	return names, rows.Err()
}

func (d *DB) SetUserPassword(id int64, hash string) error {
	_, err := d.conn.Exec(`UPDATE users SET pass_hash=? WHERE id=?`, hash, id)
	return err
}

func (d *DB) SetUserAdmin(id int64, admin bool) error {
	val := 0
	if admin {
		val = 1
	}
	_, err := d.conn.Exec(`UPDATE users SET is_admin=? WHERE id=?`, val, id)
	return err
}

func (d *DB) DeleteUser(id int64) error {
	if _, err := d.conn.Exec(`DELETE FROM sessions WHERE user_id=?`, id); err != nil {
		return err
	}
	_, err := d.conn.Exec(`DELETE FROM users WHERE id=?`, id)
	return err
}

func affected(r sql.Result, err error) (bool, error) {
	if err != nil {
		return false, err
	}
	n, err := r.RowsAffected()
	return n == 1, err
}

// GetTOTP returns the sealed TOTP secret ("" = 2FA off) and the last used time step.
func (d *DB) GetTOTP(userID int64) (secret string, last int64, err error) {
	err = d.conn.QueryRow(`SELECT totp_secret, totp_last FROM users WHERE id=?`, userID).Scan(&secret, &last)
	return
}

// EnableTOTP stores the sealed secret and replaces the recovery codes (sha256 hex).
func (d *DB) EnableTOTP(userID int64, secret string, last int64, codeHashes []string) error {
	tx, err := d.conn.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if _, err := tx.Exec(`UPDATE users SET totp_secret=?, totp_last=? WHERE id=?`, secret, last, userID); err != nil {
		return err
	}
	if _, err := tx.Exec(`DELETE FROM recovery_codes WHERE user_id=?`, userID); err != nil {
		return err
	}
	for _, h := range codeHashes {
		if _, err := tx.Exec(`INSERT INTO recovery_codes (user_id, hash) VALUES (?,?)`, userID, h); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func (d *DB) DisableTOTP(userID int64) error {
	return d.EnableTOTP(userID, "", 0, nil)
}

// UseTOTPStep records step as used; false if it (or a later one) was already used.
func (d *DB) UseTOTPStep(userID, step int64) (bool, error) {
	return affected(d.conn.Exec(`UPDATE users SET totp_last=? WHERE id=? AND totp_last < ? AND totp_secret != ''`, step, userID, step))
}

// UseRecoveryCode burns a recovery code by its sha256 hex; false if unknown or used.
func (d *DB) UseRecoveryCode(userID int64, hash string) (bool, error) {
	return affected(d.conn.Exec(`DELETE FROM recovery_codes WHERE user_id=? AND hash=?`, userID, hash))
}

func (d *DB) CountRecoveryCodes(userID int64) (int, error) {
	var n int
	err := d.conn.QueryRow(`SELECT COUNT(*) FROM recovery_codes WHERE user_id=?`, userID).Scan(&n)
	return n, err
}

// CreateResetToken returns a fresh single-use password reset token; only its hash is stored.
func (d *DB) CreateResetToken(userID int64, ttl time.Duration) (string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	tok := hex.EncodeToString(b)
	_, err := d.conn.Exec(`INSERT INTO reset_tokens (hash, user_id, expires_at) VALUES (?,?,?)`,
		hashToken(tok), userID, time.Now().Add(ttl))
	return tok, err
}

func (d *DB) GetResetToken(token string) (int64, error) {
	var id int64
	err := d.conn.QueryRow(`SELECT user_id FROM reset_tokens WHERE hash=? AND expires_at > ?`,
		hashToken(token), time.Now()).Scan(&id)
	return id, err
}

// ConsumeResetToken deletes a valid token; false if it was already used or expired.
func (d *DB) ConsumeResetToken(token string) (bool, error) {
	return affected(d.conn.Exec(`DELETE FROM reset_tokens WHERE hash=? AND expires_at > ?`, hashToken(token), time.Now()))
}

func (d *DB) DeleteResetTokens(userID int64) error {
	_, err := d.conn.Exec(`DELETE FROM reset_tokens WHERE user_id=?`, userID)
	return err
}

const nameTakenSQL = `SELECT EXISTS(SELECT 1 FROM users WHERE lower(username)=lower(?) AND id!=?)
    OR EXISTS(SELECT 1 FROM former_names WHERE name=? AND user_id!=?)`

// NameTaken reports whether name (any case) belongs to, or once belonged to, a user other than uid.
func (d *DB) NameTaken(name string, uid int64) (bool, error) {
	var taken bool
	err := d.conn.QueryRow(nameTakenSQL, name, uid, name, uid).Scan(&taken)
	return taken, err
}

// LastRename returns when uid last changed their username (zero if never).
func (d *DB) LastRename(uid int64) (time.Time, error) {
	var t time.Time
	err := d.conn.QueryRow(`SELECT renamed_at FROM former_names WHERE user_id=? ORDER BY renamed_at DESC LIMIT 1`, uid).Scan(&t)
	if err == sql.ErrNoRows {
		return time.Time{}, nil
	}
	return t, err
}

// RenameUser changes a username and keeps the old one reserved for this user.
func (d *DB) RenameUser(uid int64, newName string) error {
	tx, err := d.conn.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	var old string
	var taken bool
	if err := tx.QueryRow(`SELECT username FROM users WHERE id=?`, uid).Scan(&old); err != nil {
		return err
	}
	if err := tx.QueryRow(nameTakenSQL, newName, uid, newName, uid).Scan(&taken); err != nil {
		return err
	}
	if taken {
		return ErrNameTaken
	}
	for _, q := range []struct {
		sql  string
		args []any
	}{
		{`DELETE FROM former_names WHERE name=?`, []any{newName}},
		{`INSERT INTO former_names (name, user_id) VALUES (?, ?)
          ON CONFLICT(name) DO UPDATE SET renamed_at=CURRENT_TIMESTAMP`, []any{old, uid}},
		{`UPDATE users SET username=? WHERE id=?`, []any{newName, uid}},
	} {
		if _, err := tx.Exec(q.sql, q.args...); err != nil {
			return err
		}
	}
	return tx.Commit()
}