Warp/server/internal/storage/sqlstore/keys.go

package sqlstore

import (
	"cmp"
	"context"
	"database/sql"
	"errors"
	"fmt"
	"slices"
	"strconv"

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

// Keys is the shared SQL implementation of storage.KeyStore.
//
// Callers hand it key JSON already in canonical form. That is what lets
// "the same key uploaded again" be a byte comparison here: a client that
// re-sends a key it already uploaded may well encode it differently, and
// must not be told it is a conflicting one.
type Keys struct {
	db     *sql.DB
	rebind func(string) string
}

// NewKeys builds a key store.
func NewKeys(db *sql.DB, rebind func(string) string) *Keys {
	if rebind == nil {
		rebind = func(query string) string { return query }
	}
	return &Keys{db: db, rebind: rebind}
}

// The keyserver's sequence table and its one counter.
const (
	keySequences  = "keyserver_sequences"
	seqDeviceList = "device_list"
)

// maxInList caps the size of an IN list. SQLite's limit on bound
// parameters is in the tens of thousands and PostgreSQL's 65535, and a
// sync of a user in a large room asks about every member of it.
const maxInList = 500

// claimAttempts bounds the retries of a claim that lost a race for a key.
const claimAttempts = 5

func (s *Keys) q(query string) string { return s.rebind(query) }

// tx runs fn in a transaction, rolling back on any error.
func (s *Keys) tx(ctx context.Context, fn func(*sql.Tx) error) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin transaction: %w", err)
	}
	defer func() { _ = tx.Rollback() }()

	if err := fn(tx); err != nil {
		return err
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit transaction: %w", err)
	}
	return nil
}

// --- the device list stream ---------------------------------------------------

const (
	selectDeviceListPosSQL = `
		SELECT stream_pos FROM keyserver_device_lists WHERE user_id = ?`

	upsertDeviceListPosSQL = `
		INSERT INTO keyserver_device_lists (user_id, stream_pos) VALUES (?, ?)
		ON CONFLICT (user_id) DO UPDATE SET stream_pos = excluded.stream_pos`
)

// bumpDeviceList moves a user's device list once per device named, or once
// for the user when none is, and returns an update per move.
//
// Each device gets a position of its own because a remote server treats a
// stream ID it has already seen as a duplicate: one position shared by the
// deletion of three devices would reach it as one deletion and two
// repeats.
func bumpDeviceList(ctx context.Context, tx *sql.Tx, q func(string) string, userID string, deviceIDs []string) ([]storage.DeviceListUpdate, error) {
	if len(deviceIDs) == 0 {
		deviceIDs = []string{""}
	}
	n := int64(len(deviceIDs))
	last, err := advanceSequence(ctx, tx, q, keySequences, seqDeviceList, n)
	if err != nil {
		return nil, err
	}

	var prev int64
	err = tx.QueryRowContext(ctx, q(selectDeviceListPosSQL), userID).Scan(&prev)
	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		return nil, fmt.Errorf("read the device list position of %q: %w", userID, err)
	}

	out := make([]storage.DeviceListUpdate, 0, n)
	for i, deviceID := range deviceIDs {
		pos := last - n + 1 + int64(i)
		out = append(out, storage.DeviceListUpdate{
			UserID: userID, DeviceID: deviceID, StreamPos: pos, PrevPos: prev,
		})
		prev = pos
	}
	if _, err := tx.ExecContext(ctx, q(upsertDeviceListPosSQL), userID, last); err != nil {
		return nil, fmt.Errorf("move the device list of %q: %w", userID, err)
	}
	return out, nil
}

// MarkDeviceListsChanged implements storage.KeyStore.
func (s *Keys) MarkDeviceListsChanged(ctx context.Context, userIDs []string) (int64, error) {
	var top int64
	err := s.tx(ctx, func(tx *sql.Tx) error {
		for _, userID := range dedupe(userIDs) {
			updates, err := bumpDeviceList(ctx, tx, s.q, userID, nil)
			if err != nil {
				return err
			}
			top = updates[len(updates)-1].StreamPos
		}
		return nil
	})
	return top, err
}

// DeviceListChanges implements storage.KeyStore.
func (s *Keys) DeviceListChanges(ctx context.Context, userIDs []string, from, to int64) ([]storage.DeviceListChange, error) {
	var out []storage.DeviceListChange
	for chunk := range slices.Chunk(dedupe(userIDs), maxInList) {
		query := `
			SELECT user_id, stream_pos FROM keyserver_device_lists
			WHERE user_id IN (` + placeholders(len(chunk)) + `) AND stream_pos > ?`
		args := make([]any, 0, len(chunk)+2)
		for _, userID := range chunk {
			args = append(args, userID)
		}
		args = append(args, from)
		if to > 0 {
			query += ` AND stream_pos <= ?`
			args = append(args, to)
		}

		changes, err := s.deviceListChanges(ctx, query, args)
		if err != nil {
			return nil, err
		}
		out = append(out, changes...)
	}
	slices.SortFunc(out, func(a, b storage.DeviceListChange) int { return cmp.Compare(a.StreamPos, b.StreamPos) })
	return out, nil
}

// deviceListChanges runs one chunk of DeviceListChanges.
func (s *Keys) deviceListChanges(ctx context.Context, query string, args []any) ([]storage.DeviceListChange, error) {
	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read device list changes: %w", err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.DeviceListChange
	for rows.Next() {
		var c storage.DeviceListChange
		if err := rows.Scan(&c.UserID, &c.StreamPos); err != nil {
			return nil, fmt.Errorf("read device list changes: %w", err)
		}
		out = append(out, c)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("read device list changes: %w", err)
	}
	return out, nil
}

// DeviceListPosition implements storage.KeyStore.
func (s *Keys) DeviceListPosition(ctx context.Context, userID string) (int64, error) {
	var pos int64
	err := s.db.QueryRowContext(ctx, s.q(selectDeviceListPosSQL), userID).Scan(&pos)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return 0, nil
	case err != nil:
		return 0, fmt.Errorf("read the device list position of %q: %w", userID, err)
	}
	return pos, nil
}

// --- device keys --------------------------------------------------------------

const (
	selectDeviceKeyJSONSQL = `
		SELECT key_json FROM keyserver_device_keys WHERE user_id = ? AND device_id = ?`

	upsertDeviceKeysSQL = `
		INSERT INTO keyserver_device_keys (user_id, device_id, key_json) VALUES (?, ?, ?)
		ON CONFLICT (user_id, device_id) DO UPDATE SET key_json = excluded.key_json`
)

// SetDeviceKeys implements storage.KeyStore.
func (s *Keys) SetDeviceKeys(ctx context.Context, keys storage.DeviceKeys) (storage.DeviceListUpdate, error) {
	var update storage.DeviceListUpdate
	err := s.tx(ctx, func(tx *sql.Tx) error {
		var existing string
		err := tx.QueryRowContext(ctx, s.q(selectDeviceKeyJSONSQL), keys.UserID, keys.DeviceID).Scan(&existing)
		switch {
		case err == nil && existing == string(keys.JSON):
			return nil
		case err != nil && !errors.Is(err, sql.ErrNoRows):
			return fmt.Errorf("read the keys of %s/%s: %w", keys.UserID, keys.DeviceID, err)
		}

		if _, err := tx.ExecContext(ctx, s.q(upsertDeviceKeysSQL), keys.UserID, keys.DeviceID, string(keys.JSON)); err != nil {
			return fmt.Errorf("store the keys of %s/%s: %w", keys.UserID, keys.DeviceID, err)
		}
		updates, err := bumpDeviceList(ctx, tx, s.q, keys.UserID, []string{keys.DeviceID})
		if err != nil {
			return err
		}
		update = updates[0]
		return nil
	})
	return update, err
}

// DeviceKeys implements storage.KeyStore.
func (s *Keys) DeviceKeys(ctx context.Context, userID string, deviceIDs []string) ([]storage.DeviceKeys, error) {
	query := `SELECT device_id, key_json FROM keyserver_device_keys WHERE user_id = ?`
	args := []any{userID}
	if len(deviceIDs) > 0 {
		deviceIDs = dedupe(deviceIDs)
		if len(deviceIDs) > maxInList {
			deviceIDs = deviceIDs[:maxInList]
		}
		query += ` AND device_id IN (` + placeholders(len(deviceIDs)) + `)`
		for _, deviceID := range deviceIDs {
			args = append(args, deviceID)
		}
	}
	query += ` ORDER BY device_id`

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

	var out []storage.DeviceKeys
	for rows.Next() {
		k := storage.DeviceKeys{UserID: userID}
		var keyJSON string
		if err := rows.Scan(&k.DeviceID, &keyJSON); err != nil {
			return nil, fmt.Errorf("read the device keys of %q: %w", userID, err)
		}
		k.JSON = []byte(keyJSON)
		out = append(out, k)
	}
	return out, rows.Err()
}

// DeleteDeviceKeys implements storage.KeyStore.
func (s *Keys) DeleteDeviceKeys(ctx context.Context, userID string, deviceIDs []string) ([]storage.DeviceListUpdate, error) {
	var updates []storage.DeviceListUpdate
	err := s.tx(ctx, func(tx *sql.Tx) error {
		where := ` WHERE user_id = ?`
		args := []any{userID}
		if len(deviceIDs) > 0 {
			deviceIDs = dedupe(deviceIDs)
			where += ` AND device_id IN (` + placeholders(len(deviceIDs)) + `)`
			for _, deviceID := range deviceIDs {
				args = append(args, deviceID)
			}
		}

		had, err := s.devicesWithKeys(ctx, tx, userID, where, args)
		if err != nil {
			return err
		}

		for _, table := range []string{"keyserver_device_keys", "keyserver_one_time_keys", "keyserver_fallback_keys"} {
			if _, err := tx.ExecContext(ctx, s.q(`DELETE FROM `+table+where), args...); err != nil {
				return fmt.Errorf("delete from %s for %q: %w", table, userID, err)
			}
		}
		if len(had) == 0 {
			return nil
		}
		updates, err = bumpDeviceList(ctx, tx, s.q, userID, had)
		return err
	})
	return updates, err
}

// devicesWithKeys lists the devices a DeleteDeviceKeys will take identity
// keys from, which are the ones whose deletion moves the stream.
func (s *Keys) devicesWithKeys(ctx context.Context, tx *sql.Tx, userID, where string, args []any) ([]string, error) {
	rows, err := tx.QueryContext(ctx, s.q(`SELECT device_id FROM keyserver_device_keys`+where+` ORDER BY device_id`), args...)
	if err != nil {
		return nil, fmt.Errorf("read the devices of %q: %w", userID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []string
	for rows.Next() {
		var deviceID string
		if err := rows.Scan(&deviceID); err != nil {
			return nil, fmt.Errorf("read the devices of %q: %w", userID, err)
		}
		out = append(out, deviceID)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("read the devices of %q: %w", userID, err)
	}
	return out, nil
}

// --- one-time and fallback keys -----------------------------------------------

const (
	selectOneTimeKeySQL = `
		SELECT key_json FROM keyserver_one_time_keys
		WHERE user_id = ? AND device_id = ? AND key_id = ?`

	insertOneTimeKeySQL = `
		INSERT INTO keyserver_one_time_keys (user_id, device_id, key_id, algorithm, key_json)
		VALUES (?, ?, ?, ?, ?)`

	countOneTimeKeysSQL = `
		SELECT algorithm, COUNT(*) FROM keyserver_one_time_keys
		WHERE user_id = ? AND device_id = ? GROUP BY algorithm`

	selectFallbackKeySQL = `
		SELECT key_id, key_json FROM keyserver_fallback_keys
		WHERE user_id = ? AND device_id = ? AND algorithm = ?`

	upsertFallbackKeySQL = `
		INSERT INTO keyserver_fallback_keys (user_id, device_id, algorithm, key_id, key_json, used)
		VALUES (?, ?, ?, ?, ?, ?)
		ON CONFLICT (user_id, device_id, algorithm) DO UPDATE SET
			key_id = excluded.key_id, key_json = excluded.key_json, used = excluded.used`

	selectUnusedFallbackSQL = `
		SELECT algorithm FROM keyserver_fallback_keys
		WHERE user_id = ? AND device_id = ? AND used = ? ORDER BY algorithm`

	selectClaimableKeySQL = `
		SELECT key_id, key_json FROM keyserver_one_time_keys
		WHERE user_id = ? AND device_id = ? AND algorithm = ?
		ORDER BY key_id LIMIT 1`

	deleteOneTimeKeySQL = `
		DELETE FROM keyserver_one_time_keys WHERE user_id = ? AND device_id = ? AND key_id = ?`

	markFallbackUsedSQL = `
		UPDATE keyserver_fallback_keys SET used = ?
		WHERE user_id = ? AND device_id = ? AND algorithm = ?`
)

// AddOneTimeKeys implements storage.KeyStore.
func (s *Keys) AddOneTimeKeys(ctx context.Context, userID, deviceID string, keys []storage.OneTimeKey) error {
	if len(keys) == 0 {
		return nil
	}
	return s.tx(ctx, func(tx *sql.Tx) error {
		for _, k := range keys {
			var existing string
			err := tx.QueryRowContext(ctx, s.q(selectOneTimeKeySQL), userID, deviceID, k.KeyID).Scan(&existing)
			switch {
			case err == nil && existing == string(k.JSON):
				continue
			case err == nil:
				return fmt.Errorf("%w: %s", storage.ErrKeyConflict, k.KeyID)
			case !errors.Is(err, sql.ErrNoRows):
				return fmt.Errorf("read one-time key %s of %s/%s: %w", k.KeyID, userID, deviceID, err)
			}
			if _, err := tx.ExecContext(ctx, s.q(insertOneTimeKeySQL),
				userID, deviceID, k.KeyID, k.Algorithm(), string(k.JSON)); err != nil {
				return fmt.Errorf("store one-time key %s of %s/%s: %w", k.KeyID, userID, deviceID, err)
			}
		}
		return nil
	})
}

// OneTimeKeyCounts implements storage.KeyStore.
func (s *Keys) OneTimeKeyCounts(ctx context.Context, userID, deviceID string) (map[string]int, error) {
	rows, err := s.db.QueryContext(ctx, s.q(countOneTimeKeysSQL), userID, deviceID)
	if err != nil {
		return nil, fmt.Errorf("count the one-time keys of %s/%s: %w", userID, deviceID, err)
	}
	defer func() { _ = rows.Close() }()

	out := map[string]int{}
	for rows.Next() {
		var (
			algorithm string
			n         int
		)
		if err := rows.Scan(&algorithm, &n); err != nil {
			return nil, fmt.Errorf("count the one-time keys of %s/%s: %w", userID, deviceID, err)
		}
		out[algorithm] = n
	}
	return out, rows.Err()
}

// SetFallbackKeys implements storage.KeyStore.
func (s *Keys) SetFallbackKeys(ctx context.Context, userID, deviceID string, keys []storage.OneTimeKey) error {
	if len(keys) == 0 {
		return nil
	}
	return s.tx(ctx, func(tx *sql.Tx) error {
		for _, k := range keys {
			var existingID, existingJSON string
			err := tx.QueryRowContext(ctx, s.q(selectFallbackKeySQL), userID, deviceID, k.Algorithm()).
				Scan(&existingID, &existingJSON)
			switch {
			case err == nil && existingID == k.KeyID && existingJSON == string(k.JSON):
				// The same key again keeps its used flag.
				continue
			case err != nil && !errors.Is(err, sql.ErrNoRows):
				return fmt.Errorf("read the fallback key of %s/%s: %w", userID, deviceID, err)
			}
			if _, err := tx.ExecContext(ctx, s.q(upsertFallbackKeySQL),
				userID, deviceID, k.Algorithm(), k.KeyID, string(k.JSON), false); err != nil {
				return fmt.Errorf("store the fallback key of %s/%s: %w", userID, deviceID, err)
			}
		}
		return nil
	})
}

// UnusedFallbackKeyTypes implements storage.KeyStore.
func (s *Keys) UnusedFallbackKeyTypes(ctx context.Context, userID, deviceID string) ([]string, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectUnusedFallbackSQL), userID, deviceID, false)
	if err != nil {
		return nil, fmt.Errorf("read the fallback keys of %s/%s: %w", userID, deviceID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []string
	for rows.Next() {
		var algorithm string
		if err := rows.Scan(&algorithm); err != nil {
			return nil, fmt.Errorf("read the fallback keys of %s/%s: %w", userID, deviceID, err)
		}
		out = append(out, algorithm)
	}
	return out, rows.Err()
}

// ClaimKey implements storage.KeyStore.
//
// The one-time key is read and then deleted, and the claim is only won if
// the delete removed the row: two servers claiming a key of the same
// device at the same moment can both read it, and only one of them may
// have it. The loser tries the next key.
func (s *Keys) ClaimKey(ctx context.Context, userID, deviceID, algorithm string) (storage.ClaimedKey, bool, error) {
	for range claimAttempts {
		var (
			claimed storage.ClaimedKey
			found   bool
			lost    bool
		)
		err := s.tx(ctx, func(tx *sql.Tx) error {
			var keyID, keyJSON string
			err := tx.QueryRowContext(ctx, s.q(selectClaimableKeySQL), userID, deviceID, algorithm).Scan(&keyID, &keyJSON)
			switch {
			case err == nil:
				res, err := tx.ExecContext(ctx, s.q(deleteOneTimeKeySQL), userID, deviceID, keyID)
				if err != nil {
					return fmt.Errorf("claim one-time key %s of %s/%s: %w", keyID, userID, deviceID, err)
				}
				if n, _ := res.RowsAffected(); n != 1 {
					lost = true
					return nil
				}
				claimed = storage.ClaimedKey{OneTimeKey: storage.OneTimeKey{KeyID: keyID, JSON: []byte(keyJSON)}}
				found = true
				return nil
			case !errors.Is(err, sql.ErrNoRows):
				return fmt.Errorf("claim a one-time key of %s/%s: %w", userID, deviceID, err)
			}

			err = tx.QueryRowContext(ctx, s.q(selectFallbackKeySQL), userID, deviceID, algorithm).Scan(&keyID, &keyJSON)
			switch {
			case errors.Is(err, sql.ErrNoRows):
				return nil
			case err != nil:
				return fmt.Errorf("claim the fallback key of %s/%s: %w", userID, deviceID, err)
			}
			if _, err := tx.ExecContext(ctx, s.q(markFallbackUsedSQL), true, userID, deviceID, algorithm); err != nil {
				return fmt.Errorf("claim the fallback key of %s/%s: %w", userID, deviceID, err)
			}
			claimed = storage.ClaimedKey{OneTimeKey: storage.OneTimeKey{KeyID: keyID, JSON: []byte(keyJSON)}, Fallback: true}
			found = true
			return nil
		})
		if err != nil {
			return storage.ClaimedKey{}, false, err
		}
		if !lost {
			return claimed, found, nil
		}
	}
	return storage.ClaimedKey{}, false, fmt.Errorf("claim a key of %s/%s: lost the race %d times", userID, deviceID, claimAttempts)
}

// --- cross-signing ------------------------------------------------------------

const (
	upsertCrossSigningKeySQL = `
		INSERT INTO keyserver_cross_signing_keys (user_id, purpose, key_json) VALUES (?, ?, ?)
		ON CONFLICT (user_id, purpose) DO UPDATE SET key_json = excluded.key_json`

	deleteCrossSigningKeySQL = `
		DELETE FROM keyserver_cross_signing_keys WHERE user_id = ? AND purpose = ?`

	selectCrossSigningKeysSQL = `
		SELECT purpose, key_json FROM keyserver_cross_signing_keys WHERE user_id = ?`
)

// SetCrossSigningKeys implements storage.KeyStore.
func (s *Keys) SetCrossSigningKeys(ctx context.Context, userID string, keys map[string][]byte) (storage.DeviceListUpdate, error) {
	var update storage.DeviceListUpdate
	err := s.tx(ctx, func(tx *sql.Tx) error {
		for _, purpose := range sortedKeys(keys) {
			keyJSON := keys[purpose]
			var err error
			if keyJSON == nil {
				_, err = tx.ExecContext(ctx, s.q(deleteCrossSigningKeySQL), userID, purpose)
			} else {
				_, err = tx.ExecContext(ctx, s.q(upsertCrossSigningKeySQL), userID, purpose, string(keyJSON))
			}
			if err != nil {
				return fmt.Errorf("store the %s key of %q: %w", purpose, userID, err)
			}
		}
		updates, err := bumpDeviceList(ctx, tx, s.q, userID, nil)
		if err != nil {
			return err
		}
		update = updates[0]
		return nil
	})
	return update, err
}

// CrossSigningKeys implements storage.KeyStore.
func (s *Keys) CrossSigningKeys(ctx context.Context, userID string) (map[string][]byte, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectCrossSigningKeysSQL), userID)
	if err != nil {
		return nil, fmt.Errorf("read the cross-signing keys of %q: %w", userID, err)
	}
	defer func() { _ = rows.Close() }()

	out := map[string][]byte{}
	for rows.Next() {
		var purpose, keyJSON string
		if err := rows.Scan(&purpose, &keyJSON); err != nil {
			return nil, fmt.Errorf("read the cross-signing keys of %q: %w", userID, err)
		}
		out[purpose] = []byte(keyJSON)
	}
	return out, rows.Err()
}

// --- signatures ---------------------------------------------------------------

const (
	upsertSignatureSQL = `
		INSERT INTO keyserver_signatures
			(target_user_id, target_key_id, signer_user_id, signer_key_id, signature)
		VALUES (?, ?, ?, ?, ?)
		ON CONFLICT (target_user_id, target_key_id, signer_user_id, signer_key_id)
		DO UPDATE SET signature = excluded.signature`

	selectSignaturesSQL = `
		SELECT target_key_id, signer_user_id, signer_key_id, signature
		FROM keyserver_signatures WHERE target_user_id = ?
		ORDER BY target_key_id, signer_user_id, signer_key_id`
)

// AddSignatures implements storage.KeyStore.
func (s *Keys) AddSignatures(ctx context.Context, sigs []storage.KeySignature) ([]storage.DeviceListUpdate, error) {
	if len(sigs) == 0 {
		return nil, nil
	}
	var updates []storage.DeviceListUpdate
	err := s.tx(ctx, func(tx *sql.Tx) error {
		targets := map[string][]string{}
		for _, sig := range sigs {
			if _, err := tx.ExecContext(ctx, s.q(upsertSignatureSQL),
				sig.TargetUserID, sig.TargetKeyID, sig.SignerUserID, sig.SignerKeyID, sig.Signature); err != nil {
				return fmt.Errorf("store a signature of %s/%s: %w", sig.TargetUserID, sig.TargetKeyID, err)
			}
			if !slices.Contains(targets[sig.TargetUserID], sig.TargetKeyID) {
				targets[sig.TargetUserID] = append(targets[sig.TargetUserID], sig.TargetKeyID)
			}
		}
		for _, userID := range sortedKeys(targets) {
			moved, err := bumpDeviceList(ctx, tx, s.q, userID, targets[userID])
			if err != nil {
				return err
			}
			updates = append(updates, moved...)
		}
		return nil
	})
	return updates, err
}

// Signatures implements storage.KeyStore.
func (s *Keys) Signatures(ctx context.Context, targetUserID string) ([]storage.KeySignature, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectSignaturesSQL), targetUserID)
	if err != nil {
		return nil, fmt.Errorf("read the signatures of %q's keys: %w", targetUserID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.KeySignature
	for rows.Next() {
		sig := storage.KeySignature{TargetUserID: targetUserID}
		if err := rows.Scan(&sig.TargetKeyID, &sig.SignerUserID, &sig.SignerKeyID, &sig.Signature); err != nil {
			return nil, fmt.Errorf("read the signatures of %q's keys: %w", targetUserID, err)
		}
		out = append(out, sig)
	}
	return out, rows.Err()
}

// --- key backup ---------------------------------------------------------------

const (
	selectMaxBackupVersionSQL = `
		SELECT COALESCE(MAX(version), 0) FROM keyserver_backup_versions WHERE user_id = ?`

	insertBackupVersionSQL = `
		INSERT INTO keyserver_backup_versions (user_id, version, algorithm, auth_data, etag, deleted)
		VALUES (?, ?, ?, ?, 0, ?)`

	selectCurrentBackupVersionSQL = `
		SELECT COALESCE(MAX(version), 0) FROM keyserver_backup_versions
		WHERE user_id = ? AND deleted = ?`

	selectBackupVersionSQL = `
		SELECT algorithm, auth_data, etag FROM keyserver_backup_versions
		WHERE user_id = ? AND version = ? AND deleted = ?`

	countBackupKeysSQL = `
		SELECT COUNT(*) FROM keyserver_backup_keys WHERE user_id = ? AND version = ?`

	updateBackupAuthDataSQL = `
		UPDATE keyserver_backup_versions SET auth_data = ?
		WHERE user_id = ? AND version = ? AND deleted = ?`

	deleteBackupVersionSQL = `
		UPDATE keyserver_backup_versions SET deleted = ?
		WHERE user_id = ? AND version = ? AND deleted = ?`

	bumpBackupETagSQL = `
		UPDATE keyserver_backup_versions SET etag = etag + 1 WHERE user_id = ? AND version = ?`

	selectBackupKeySQL = `
		SELECT first_message_index, forwarded_count, is_verified, session_data
		FROM keyserver_backup_keys
		WHERE user_id = ? AND version = ? AND room_id = ? AND session_id = ?`

	upsertBackupKeySQL = `
		INSERT INTO keyserver_backup_keys
			(user_id, version, room_id, session_id, first_message_index, forwarded_count,
			 is_verified, session_data)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?)
		ON CONFLICT (user_id, version, room_id, session_id) DO UPDATE SET
			first_message_index = excluded.first_message_index,
			forwarded_count = excluded.forwarded_count,
			is_verified = excluded.is_verified,
			session_data = excluded.session_data`
)

// rowQuerier is the part of *sql.Tx and *sql.DB a single-row read needs.
type rowQuerier interface {
	QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
}

// parseBackupVersion reads a version number. A version that is not one is
// a version that does not exist, not a malformed request: the spec makes
// versions opaque strings, and a client is entitled to ask about any.
func parseBackupVersion(version string) (int64, bool) {
	n, err := strconv.ParseInt(version, 10, 64)
	return n, err == nil && n > 0
}

// readBackup reads one live version, with its key count.
func (s *Keys) readBackup(ctx context.Context, db rowQuerier, userID, version string) (storage.BackupVersion, int64, error) {
	var n int64
	if version == "" {
		if err := db.QueryRowContext(ctx, s.q(selectCurrentBackupVersionSQL), userID, false).Scan(&n); err != nil {
			return storage.BackupVersion{}, 0, fmt.Errorf("read the backup version of %q: %w", userID, err)
		}
		if n == 0 {
			return storage.BackupVersion{}, 0, storage.ErrNoSuchBackup
		}
	} else {
		var ok bool
		if n, ok = parseBackupVersion(version); !ok {
			return storage.BackupVersion{}, 0, storage.ErrNoSuchBackup
		}
	}

	out := storage.BackupVersion{UserID: userID, Version: strconv.FormatInt(n, 10)}
	var (
		authData string
		etag     int64
	)
	err := db.QueryRowContext(ctx, s.q(selectBackupVersionSQL), userID, n, false).Scan(&out.Algorithm, &authData, &etag)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.BackupVersion{}, 0, storage.ErrNoSuchBackup
	case err != nil:
		return storage.BackupVersion{}, 0, fmt.Errorf("read backup version %d of %q: %w", n, userID, err)
	}
	out.AuthData = []byte(authData)
	out.ETag = strconv.FormatInt(etag, 10)
	if err := db.QueryRowContext(ctx, s.q(countBackupKeysSQL), userID, n).Scan(&out.Count); err != nil {
		return storage.BackupVersion{}, 0, fmt.Errorf("count the keys of backup version %d of %q: %w", n, userID, err)
	}
	return out, n, nil
}

// CreateBackup implements storage.KeyStore.
func (s *Keys) CreateBackup(ctx context.Context, userID, algorithm string, authData []byte) (string, error) {
	var version int64
	err := s.tx(ctx, func(tx *sql.Tx) error {
		if err := tx.QueryRowContext(ctx, s.q(selectMaxBackupVersionSQL), userID).Scan(&version); err != nil {
			return fmt.Errorf("read the backup versions of %q: %w", userID, err)
		}
		version++
		if _, err := tx.ExecContext(ctx, s.q(insertBackupVersionSQL), userID, version, algorithm, string(authData), false); err != nil {
			return fmt.Errorf("create a backup version for %q: %w", userID, err)
		}
		return nil
	})
	if err != nil {
		return "", err
	}
	return strconv.FormatInt(version, 10), nil
}

// Backup implements storage.KeyStore.
func (s *Keys) Backup(ctx context.Context, userID, version string) (storage.BackupVersion, error) {
	out, _, err := s.readBackup(ctx, s.db, userID, version)
	return out, err
}

// UpdateBackup implements storage.KeyStore.
func (s *Keys) UpdateBackup(ctx context.Context, userID, version string, authData []byte) error {
	n, ok := parseBackupVersion(version)
	if !ok {
		return storage.ErrNoSuchBackup
	}
	res, err := s.db.ExecContext(ctx, s.q(updateBackupAuthDataSQL), string(authData), userID, n, false)
	if err != nil {
		return fmt.Errorf("update backup version %d of %q: %w", n, userID, err)
	}
	if affected, _ := res.RowsAffected(); affected == 0 {
		return storage.ErrNoSuchBackup
	}
	return nil
}

// DeleteBackup implements storage.KeyStore.
func (s *Keys) DeleteBackup(ctx context.Context, userID, version string) error {
	n, ok := parseBackupVersion(version)
	if !ok {
		return storage.ErrNoSuchBackup
	}
	return s.tx(ctx, func(tx *sql.Tx) error {
		res, err := tx.ExecContext(ctx, s.q(deleteBackupVersionSQL), true, userID, n, false)
		if err != nil {
			return fmt.Errorf("delete backup version %d of %q: %w", n, userID, err)
		}
		if affected, _ := res.RowsAffected(); affected == 0 {
			return storage.ErrNoSuchBackup
		}
		if _, err := tx.ExecContext(ctx, s.q(`DELETE FROM keyserver_backup_keys WHERE user_id = ? AND version = ?`), userID, n); err != nil {
			return fmt.Errorf("delete the keys of backup version %d of %q: %w", n, userID, err)
		}
		return nil
	})
}

// PutBackupKeys implements storage.KeyStore.
func (s *Keys) PutBackupKeys(ctx context.Context, userID, version string, keys []storage.BackupKey) (storage.BackupVersion, error) {
	var out storage.BackupVersion
	err := s.tx(ctx, func(tx *sql.Tx) error {
		_, n, err := s.readBackup(ctx, tx, userID, version)
		if err != nil {
			return err
		}

		changed := false
		for _, k := range keys {
			var (
				existing    storage.BackupKey
				sessionData string
			)
			err := tx.QueryRowContext(ctx, s.q(selectBackupKeySQL), userID, n, k.RoomID, k.SessionID).
				Scan(&existing.FirstMessageIndex, &existing.ForwardedCount, &existing.IsVerified, &sessionData)
			switch {
			case err == nil:
				if !k.Replaces(existing) {
					continue
				}
			case !errors.Is(err, sql.ErrNoRows):
				return fmt.Errorf("read backed-up key %s/%s: %w", k.RoomID, k.SessionID, err)
			}
			if _, err := tx.ExecContext(ctx, s.q(upsertBackupKeySQL), userID, n, k.RoomID, k.SessionID,
				k.FirstMessageIndex, k.ForwardedCount, k.IsVerified, string(k.SessionData)); err != nil {
				return fmt.Errorf("back up key %s/%s: %w", k.RoomID, k.SessionID, err)
			}
			changed = true
		}
		if changed {
			if _, err := tx.ExecContext(ctx, s.q(bumpBackupETagSQL), userID, n); err != nil {
				return fmt.Errorf("move the etag of backup version %d of %q: %w", n, userID, err)
			}
		}
		out, _, err = s.readBackup(ctx, tx, userID, version)
		return err
	})
	return out, err
}

// backupKeyScope is the WHERE clause the key reads and deletes share.
func backupKeyScope(userID string, version int64, roomID, sessionID string) (string, []any) {
	where := ` WHERE user_id = ? AND version = ?`
	args := []any{userID, version}
	if roomID != "" {
		where += ` AND room_id = ?`
		args = append(args, roomID)
		if sessionID != "" {
			where += ` AND session_id = ?`
			args = append(args, sessionID)
		}
	}
	return where, args
}

// BackupKeys implements storage.KeyStore.
func (s *Keys) BackupKeys(ctx context.Context, userID, version, roomID, sessionID string) ([]storage.BackupKey, error) {
	_, n, err := s.readBackup(ctx, s.db, userID, version)
	if err != nil {
		return nil, err
	}
	where, args := backupKeyScope(userID, n, roomID, sessionID)
	rows, err := s.db.QueryContext(ctx, s.q(`
		SELECT room_id, session_id, first_message_index, forwarded_count, is_verified, session_data
		FROM keyserver_backup_keys`+where+` ORDER BY room_id, session_id`), args...)
	if err != nil {
		return nil, fmt.Errorf("read the keys of backup version %d of %q: %w", n, userID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.BackupKey
	for rows.Next() {
		var (
			k           storage.BackupKey
			sessionData string
		)
		if err := rows.Scan(&k.RoomID, &k.SessionID, &k.FirstMessageIndex, &k.ForwardedCount,
			&k.IsVerified, &sessionData); err != nil {
			return nil, fmt.Errorf("read the keys of backup version %d of %q: %w", n, userID, err)
		}
		k.SessionData = []byte(sessionData)
		out = append(out, k)
	}
	return out, rows.Err()
}

// DeleteBackupKeys implements storage.KeyStore.
func (s *Keys) DeleteBackupKeys(ctx context.Context, userID, version, roomID, sessionID string) (storage.BackupVersion, error) {
	var out storage.BackupVersion
	err := s.tx(ctx, func(tx *sql.Tx) error {
		_, n, err := s.readBackup(ctx, tx, userID, version)
		if err != nil {
			return err
		}
		where, args := backupKeyScope(userID, n, roomID, sessionID)
		res, err := tx.ExecContext(ctx, s.q(`DELETE FROM keyserver_backup_keys`+where), args...)
		if err != nil {
			return fmt.Errorf("delete keys from backup version %d of %q: %w", n, userID, err)
		}
		if affected, _ := res.RowsAffected(); affected > 0 {
			if _, err := tx.ExecContext(ctx, s.q(bumpBackupETagSQL), userID, n); err != nil {
				return fmt.Errorf("move the etag of backup version %d of %q: %w", n, userID, err)
			}
		}
		out, _, err = s.readBackup(ctx, tx, userID, version)
		return err
	})
	return out, err
}

// --- helpers ------------------------------------------------------------------

// dedupe returns the distinct values of a list, sorted, so that a caller
// naming a user twice does not move their stream twice.
func dedupe(values []string) []string {
	out := slices.Clone(values)
	slices.Sort(out)
	return slices.Compact(out)
}

// sortedKeys is the keys of a map in order, which makes the order of
// writes — and so of the positions handed out — deterministic.
func sortedKeys[V any](m map[string]V) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	slices.Sort(out)
	return out
}

var _ storage.KeyStore = (*Keys)(nil)