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

package pocketbase

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

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

// Users is the PocketBase implementation of storage.UserStore.
//
// The reads are the SQL store's statements. The writes are record
// operations, and where the SQL store lets a primary key or an UPDATE's
// affected-row count answer a question — is this localpart taken, does
// this account exist — this one asks first: the write lock is held from
// the question to the batch, so the answer cannot change in between.
type Users struct{ *base }

const (
	defaultAccountPageSize = 100
	maxAccountPageSize     = 1000
)

// CreateAccount implements storage.UserStore.
func (u *Users) CreateAccount(ctx context.Context, acc storage.NewAccount) (storage.Account, error) {
	err := u.write(ctx, func(t *tx) error {
		taken, err := t.exists("user_accounts", `localpart = ?`, acc.Localpart)
		if err != nil {
			return err
		}
		if taken {
			return storage.ErrUserExists
		}
		t.put("user_accounts", rec{
			"localpart":     acc.Localpart,
			"password_hash": acc.PasswordHash,
			"display_name":  acc.DisplayName,
			"avatar_url":    "",
			"admin":         acc.Admin,
			"deactivated":   false,
			"created_at_ms": acc.CreatedAtMS,
		})
		return nil
	})
	if err != nil {
		return storage.Account{}, fmt.Errorf("create account %q: %w", acc.Localpart, err)
	}
	return storage.Account{
		Localpart:   acc.Localpart,
		DisplayName: acc.DisplayName,
		Admin:       acc.Admin,
		CreatedAtMS: acc.CreatedAtMS,
	}, nil
}

const selectAccountSQL = `
	SELECT localpart, display_name, avatar_url, admin, deactivated, created_at_ms
	FROM user_accounts WHERE localpart = ?`

// Account implements storage.UserStore.
func (u *Users) Account(ctx context.Context, localpart string) (storage.Account, error) {
	var acc storage.Account
	err := u.db.QueryRowContext(ctx, selectAccountSQL, localpart).Scan(
		&acc.Localpart, &acc.DisplayName, &acc.AvatarURL,
		&acc.Admin, &acc.Deactivated, &acc.CreatedAtMS)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.Account{}, fmt.Errorf("account %q: %w", localpart, storage.ErrNoSuchUser)
	case err != nil:
		return storage.Account{}, fmt.Errorf("account %q: %w", localpart, err)
	}
	return acc, nil
}

// PasswordHash implements storage.UserStore.
func (u *Users) PasswordHash(ctx context.Context, localpart string) (string, error) {
	var hash string
	err := u.db.QueryRowContext(ctx, `SELECT password_hash FROM user_accounts WHERE localpart = ?`, localpart).Scan(&hash)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return "", fmt.Errorf("password for %q: %w", localpart, storage.ErrNoSuchUser)
	case err != nil:
		return "", fmt.Errorf("password for %q: %w", localpart, err)
	}
	return hash, nil
}

// SetPasswordHash implements storage.UserStore.
func (u *Users) SetPasswordHash(ctx context.Context, localpart, hash string) error {
	return u.updateAccount(ctx, localpart, rec{"password_hash": hash})
}

// SetAdmin implements storage.UserStore.
func (u *Users) SetAdmin(ctx context.Context, localpart string, admin bool) error {
	return u.updateAccount(ctx, localpart, rec{"admin": admin})
}

// SetDeactivated implements storage.UserStore.
func (u *Users) SetDeactivated(ctx context.Context, localpart string, deactivated bool) error {
	return u.updateAccount(ctx, localpart, rec{"deactivated": deactivated})
}

// SetDisplayName implements storage.UserStore.
func (u *Users) SetDisplayName(ctx context.Context, localpart, displayName string) error {
	return u.updateAccount(ctx, localpart, rec{"display_name": displayName})
}

// SetAvatarURL implements storage.UserStore.
func (u *Users) SetAvatarURL(ctx context.Context, localpart, avatarURL string) error {
	return u.updateAccount(ctx, localpart, rec{"avatar_url": avatarURL})
}

// updateAccount changes columns of one account, returning ErrNoSuchUser for
// an account that is not there.
func (u *Users) updateAccount(ctx context.Context, localpart string, row rec) error {
	err := u.write(ctx, func(t *tx) error {
		n, err := t.setWhere("user_accounts", row, `localpart = ?`, localpart)
		if err != nil {
			return err
		}
		if n == 0 {
			return storage.ErrNoSuchUser
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("update account %q: %w", localpart, err)
	}
	return nil
}

// likeEscape is the escape character the search patterns use.
const likeEscape = '!'

const (
	countAccountsSQL = `
		SELECT COUNT(*) FROM user_accounts
		WHERE LOWER(localpart) LIKE ? ESCAPE '!' OR LOWER(display_name) LIKE ? ESCAPE '!'`

	listAccountsSQL = `
		SELECT localpart, display_name, avatar_url, admin, deactivated, created_at_ms
		FROM user_accounts
		WHERE LOWER(localpart) LIKE ? ESCAPE '!' OR LOWER(display_name) LIKE ? ESCAPE '!'
		ORDER BY localpart
		LIMIT ? OFFSET ?`
)

// ListAccounts implements storage.UserStore.
func (u *Users) ListAccounts(ctx context.Context, filter storage.AccountFilter) ([]storage.Account, int, error) {
	pattern := likeContains(filter.Search)

	limit := filter.Limit
	switch {
	case limit <= 0:
		limit = defaultAccountPageSize
	case limit > maxAccountPageSize:
		limit = maxAccountPageSize
	}
	offset := max(filter.Offset, 0)

	var total int
	if err := u.db.QueryRowContext(ctx, countAccountsSQL, pattern, pattern).Scan(&total); err != nil {
		return nil, 0, fmt.Errorf("count accounts: %w", err)
	}

	rows, err := u.db.QueryContext(ctx, listAccountsSQL, pattern, pattern, limit, offset)
	if err != nil {
		return nil, 0, fmt.Errorf("list accounts: %w", err)
	}
	defer func() { _ = rows.Close() }()

	accounts := []storage.Account{}
	for rows.Next() {
		var acc storage.Account
		if err := rows.Scan(&acc.Localpart, &acc.DisplayName, &acc.AvatarURL,
			&acc.Admin, &acc.Deactivated, &acc.CreatedAtMS); err != nil {
			return nil, 0, fmt.Errorf("scan account: %w", err)
		}
		accounts = append(accounts, acc)
	}
	if err := rows.Err(); err != nil {
		return nil, 0, fmt.Errorf("list accounts: %w", err)
	}
	return accounts, total, nil
}

// likeContains turns a search term into a LIKE pattern matching it
// anywhere, with the wildcards and the escape character escaped.
func likeContains(search string) string {
	search = strings.ToLower(strings.TrimSpace(search))
	var b strings.Builder
	b.Grow(len(search) + 2)
	b.WriteByte('%')
	for _, r := range search {
		if r == '%' || r == '_' || r == likeEscape {
			b.WriteByte(likeEscape)
		}
		b.WriteRune(r)
	}
	b.WriteByte('%')
	return b.String()
}

const statsSQL = `
	SELECT
		(SELECT COUNT(*) FROM user_accounts),
		(SELECT COUNT(*) FROM user_accounts WHERE admin = ?),
		(SELECT COUNT(*) FROM user_accounts WHERE deactivated = ?),
		(SELECT COUNT(*) FROM user_devices)`

// Stats implements storage.UserStore.
func (u *Users) Stats(ctx context.Context) (storage.Stats, error) {
	var s storage.Stats
	err := u.db.QueryRowContext(ctx, statsSQL, true, true).
		Scan(&s.Users, &s.Admins, &s.Deactivated, &s.Devices)
	if err != nil {
		return storage.Stats{}, fmt.Errorf("user stats: %w", err)
	}
	return s, nil
}

// CreateSession implements storage.UserStore.
//
// Logging in on a device ID that already exists keeps the device and its
// display name but revokes the tokens it held, which is what the spec
// requires.
func (u *Users) CreateSession(ctx context.Context, s storage.NewSession) (storage.Device, error) {
	device := storage.Device{
		Localpart:   s.Localpart,
		ID:          s.DeviceID,
		DisplayName: s.DisplayName,
		CreatedAtMS: s.CreatedAtMS,
		LastSeenMS:  s.CreatedAtMS,
		LastSeenIP:  s.LastSeenIP,
		UserAgent:   s.UserAgent,
	}

	err := u.write(ctx, func(t *tx) error {
		found, err := t.exists("user_accounts", `localpart = ?`, s.Localpart)
		if err != nil {
			return fmt.Errorf("create session for %q: %w", s.Localpart, err)
		}
		if !found {
			return fmt.Errorf("create session for %q: %w", s.Localpart, storage.ErrNoSuchUser)
		}

		var existingName string
		var existingCreated int64
		err = t.queryRow(`SELECT display_name, created_at_ms FROM user_devices
			WHERE localpart = ? AND device_id = ?`, s.Localpart, s.DeviceID).
			Scan(&existingName, &existingCreated)
		switch {
		case errors.Is(err, sql.ErrNoRows):
			t.put("user_devices", rec{
				"localpart":     s.Localpart,
				"device_id":     s.DeviceID,
				"display_name":  s.DisplayName,
				"created_at_ms": s.CreatedAtMS,
				"last_seen_ms":  s.CreatedAtMS,
				"last_seen_ip":  s.LastSeenIP,
				"user_agent":    s.UserAgent,
			})
		case err != nil:
			return fmt.Errorf("read device %q: %w", s.DeviceID, err)
		default:
			device.DisplayName = existingName
			device.CreatedAtMS = existingCreated
			t.put("user_devices", rec{
				"localpart":    s.Localpart,
				"device_id":    s.DeviceID,
				"last_seen_ms": s.CreatedAtMS,
				"last_seen_ip": s.LastSeenIP,
				"user_agent":   s.UserAgent,
			})
			if err := revokeDeviceTokens(t, `localpart = ? AND device_id = ?`, s.Localpart, s.DeviceID); err != nil {
				return err
			}
		}

		t.put("user_access_tokens", rec{
			"token_hash":    s.AccessTokenHash,
			"localpart":     s.Localpart,
			"device_id":     s.DeviceID,
			"created_at_ms": s.CreatedAtMS,
			"expires_at_ms": s.AccessTokenExpiresAtMS,
		})
		if s.RefreshTokenHash != "" {
			t.put("user_refresh_tokens", rec{
				"token_hash":    s.RefreshTokenHash,
				"localpart":     s.Localpart,
				"device_id":     s.DeviceID,
				"created_at_ms": s.CreatedAtMS,
				"expires_at_ms": s.RefreshTokenExpiresAtMS,
			})
		}
		return nil
	})
	if err != nil {
		return storage.Device{}, err
	}
	return device, nil
}

// revokeDeviceTokens deletes the access and refresh tokens matching a
// condition on (localpart, device_id).
func revokeDeviceTokens(t *tx, where string, args ...any) error {
	for _, table := range []string{"user_access_tokens", "user_refresh_tokens"} {
		if _, err := t.delWhere(table, where, args...); err != nil {
			return fmt.Errorf("revoke tokens: %w", err)
		}
	}
	return nil
}

// RefreshSession implements storage.UserStore.
//
// The spent refresh token and the access token it had minted both go away
// immediately.
func (u *Users) RefreshSession(ctx context.Context, r storage.RefreshedSession) (storage.TokenInfo, error) {
	var info storage.TokenInfo
	var expired bool

	err := u.write(ctx, func(t *tx) error {
		var expiresAt int64
		err := t.queryRow(`SELECT localpart, device_id, expires_at_ms FROM user_refresh_tokens WHERE token_hash = ?`,
			r.RefreshTokenHash).Scan(&info.Localpart, &info.DeviceID, &expiresAt)
		switch {
		case errors.Is(err, sql.ErrNoRows):
			return fmt.Errorf("refresh token: %w", storage.ErrNoSuchToken)
		case err != nil:
			return fmt.Errorf("refresh token: %w", err)
		}
		if _, err := t.delKey("user_refresh_tokens", r.RefreshTokenHash); err != nil {
			return fmt.Errorf("spend refresh token: %w", err)
		}
		if expiresAt != 0 && expiresAt <= r.RefreshedAtMS {
			// The dead token is cleaned up on the way past: the batch is
			// still sent, and the error reported after it.
			expired = true
			return nil
		}

		if _, err := t.delWhere("user_access_tokens", `localpart = ? AND device_id = ?`,
			info.Localpart, info.DeviceID); err != nil {
			return fmt.Errorf("revoke previous access token: %w", err)
		}
		t.put("user_access_tokens", rec{
			"token_hash":    r.NewAccessTokenHash,
			"localpart":     info.Localpart,
			"device_id":     info.DeviceID,
			"created_at_ms": r.RefreshedAtMS,
			"expires_at_ms": r.AccessTokenExpiresAtMS,
		})
		if r.NewRefreshTokenHash != "" {
			t.put("user_refresh_tokens", rec{
				"token_hash":    r.NewRefreshTokenHash,
				"localpart":     info.Localpart,
				"device_id":     info.DeviceID,
				"created_at_ms": r.RefreshedAtMS,
				"expires_at_ms": r.RefreshTokenExpiresAtMS,
			})
		}

		if err := t.queryRow(`SELECT admin, deactivated FROM user_accounts WHERE localpart = ?`, info.Localpart).
			Scan(&info.Admin, &info.Deactivated); err != nil {
			return fmt.Errorf("read account %q: %w", info.Localpart, err)
		}
		info.ExpiresAtMS = r.AccessTokenExpiresAtMS
		return nil
	})
	if err != nil {
		return storage.TokenInfo{}, err
	}
	if expired {
		return storage.TokenInfo{}, fmt.Errorf("refresh token has expired: %w", storage.ErrNoSuchToken)
	}
	return info, nil
}

const selectTokenInfoSQL = `
	SELECT t.localpart, t.device_id, t.expires_at_ms, a.admin, a.deactivated
	FROM user_access_tokens t
	JOIN user_accounts a ON a.localpart = t.localpart
	WHERE t.token_hash = ?`

// TokenInfo implements storage.UserStore.
func (u *Users) TokenInfo(ctx context.Context, tokenHash string) (storage.TokenInfo, error) {
	var info storage.TokenInfo
	err := u.db.QueryRowContext(ctx, selectTokenInfoSQL, tokenHash).Scan(
		&info.Localpart, &info.DeviceID, &info.ExpiresAtMS, &info.Admin, &info.Deactivated)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.TokenInfo{}, storage.ErrNoSuchToken
	case err != nil:
		return storage.TokenInfo{}, fmt.Errorf("look up access token: %w", err)
	}
	return info, nil
}

// TouchDevice implements storage.UserStore. A device that has gone away
// between authenticating and recording the visit is not an error.
func (u *Users) TouchDevice(ctx context.Context, localpart, deviceID, ip, userAgent string, seenAtMS int64) error {
	err := u.write(ctx, func(t *tx) error {
		_, err := t.setWhere("user_devices", rec{
			"last_seen_ms": seenAtMS,
			"last_seen_ip": ip,
			"user_agent":   userAgent,
		}, `localpart = ? AND device_id = ?`, localpart, deviceID)
		return err
	})
	if err != nil {
		return fmt.Errorf("touch device %q: %w", deviceID, err)
	}
	return nil
}

const deviceColumns = `localpart, device_id, display_name, created_at_ms, last_seen_ms, last_seen_ip, user_agent`

// Device implements storage.UserStore.
func (u *Users) Device(ctx context.Context, localpart, deviceID string) (storage.Device, error) {
	var d storage.Device
	err := u.db.QueryRowContext(ctx, `SELECT `+deviceColumns+`
		FROM user_devices WHERE localpart = ? AND device_id = ?`, localpart, deviceID).Scan(
		&d.Localpart, &d.ID, &d.DisplayName, &d.CreatedAtMS,
		&d.LastSeenMS, &d.LastSeenIP, &d.UserAgent)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.Device{}, fmt.Errorf("device %q: %w", deviceID, storage.ErrNoSuchDevice)
	case err != nil:
		return storage.Device{}, fmt.Errorf("device %q: %w", deviceID, err)
	}
	return d, nil
}

// Devices implements storage.UserStore.
func (u *Users) Devices(ctx context.Context, localpart string) ([]storage.Device, error) {
	rows, err := u.db.QueryContext(ctx, `SELECT `+deviceColumns+`
		FROM user_devices WHERE localpart = ? ORDER BY device_id`, localpart)
	if err != nil {
		return nil, fmt.Errorf("list devices of %q: %w", localpart, err)
	}
	defer func() { _ = rows.Close() }()

	devices := []storage.Device{}
	for rows.Next() {
		var d storage.Device
		if err := rows.Scan(&d.Localpart, &d.ID, &d.DisplayName, &d.CreatedAtMS,
			&d.LastSeenMS, &d.LastSeenIP, &d.UserAgent); err != nil {
			return nil, fmt.Errorf("scan device: %w", err)
		}
		devices = append(devices, d)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("list devices of %q: %w", localpart, err)
	}
	return devices, nil
}

// SetDeviceDisplayName implements storage.UserStore.
func (u *Users) SetDeviceDisplayName(ctx context.Context, localpart, deviceID, displayName string) error {
	err := u.write(ctx, func(t *tx) error {
		n, err := t.setWhere("user_devices", rec{"display_name": displayName},
			`localpart = ? AND device_id = ?`, localpart, deviceID)
		if err != nil {
			return err
		}
		if n == 0 {
			return storage.ErrNoSuchDevice
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("rename device %q: %w", deviceID, err)
	}
	return nil
}

// DeleteDevices implements storage.UserStore.
func (u *Users) DeleteDevices(ctx context.Context, localpart string, deviceIDs []string) (int, error) {
	if len(deviceIDs) == 0 {
		return 0, nil
	}
	var deleted int
	err := u.write(ctx, func(t *tx) error {
		where := `localpart = ? AND device_id IN (SELECT value FROM json_each(?))`
		list := jsonList(deviceIDs)
		if err := revokeDeviceTokens(t, where, localpart, list); err != nil {
			return err
		}
		n, err := t.delWhere("user_devices", where, localpart, list)
		if err != nil {
			return fmt.Errorf("delete devices: %w", err)
		}
		deleted = n
		return nil
	})
	return deleted, err
}

// DeleteAllDevices implements storage.UserStore.
func (u *Users) DeleteAllDevices(ctx context.Context, localpart string) (int, error) {
	var deleted int
	err := u.write(ctx, func(t *tx) error {
		if err := revokeDeviceTokens(t, `localpart = ?`, localpart); err != nil {
			return err
		}
		n, err := t.delWhere("user_devices", `localpart = ?`, localpart)
		if err != nil {
			return fmt.Errorf("delete devices of %q: %w", localpart, err)
		}
		deleted = n
		return nil
	})
	return deleted, err
}

// AccountData implements storage.UserStore.
func (u *Users) AccountData(ctx context.Context, localpart, roomID, dataType string) ([]byte, error) {
	var content string
	err := u.db.QueryRowContext(ctx, `SELECT content FROM user_account_data
		WHERE localpart = ? AND room_id = ? AND data_type = ?`, localpart, roomID, dataType).Scan(&content)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return nil, fmt.Errorf("account data %q: %w", dataType, storage.ErrNoSuchAccountData)
	case err != nil:
		return nil, fmt.Errorf("account data %q: %w", dataType, err)
	}
	return []byte(content), nil
}

// SetAccountData implements storage.UserStore. The write stamps the row with
// a position in the account data stream, in the same batch as the content.
func (u *Users) SetAccountData(ctx context.Context, localpart, roomID, dataType string, content []byte) error {
	err := u.write(ctx, func(t *tx) error {
		pos, err := t.next(syncSequences, seqAccountData)
		if err != nil {
			return err
		}
		t.put("user_account_data", rec{
			"localpart":  localpart,
			"room_id":    roomID,
			"data_type":  dataType,
			"content":    string(content),
			"stream_pos": pos,
		})
		return nil
	})
	if err != nil {
		return fmt.Errorf("set account data %q: %w", dataType, err)
	}
	return nil
}

var _ storage.UserStore = (*Users)(nil)