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

package pocketbase

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

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

// Keys is the PocketBase implementation of storage.KeyStore.
//
// Callers hand it key JSON already in canonical form, so "the same key
// uploaded again" is a byte comparison. Where the SQL store reads back what
// it has written inside its transaction — a device list position moved
// earlier in the call, a backup's key count and etag after a change — this
// one carries the value in the call instead.
type Keys struct{ *base }

// maxInList caps an IN list, as the SQL store does.
const maxInList = 500

// keysWrite is one write call's device list positions, as it has left them.
type keysWrite struct {
	*tx
	lists map[string]int64
}

func newKeysWrite(t *tx) *keysWrite { return &keysWrite{tx: t, lists: map[string]int64{}} }

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

// 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.
func (w *keysWrite) bumpDeviceList(userID string, deviceIDs []string) ([]storage.DeviceListUpdate, error) {
	if len(deviceIDs) == 0 {
		deviceIDs = []string{""}
	}
	n := int64(len(deviceIDs))
	last, err := w.advance(keySequences, seqDeviceList, n)
	if err != nil {
		return nil, err
	}

	prev, ok := w.lists[userID]
	if !ok {
		err = w.queryRow(`SELECT stream_pos FROM keyserver_device_lists WHERE user_id = ?`, 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
	}
	w.put("keyserver_device_lists", rec{"user_id": userID, "stream_pos": last})
	w.lists[userID] = last
	return out, nil
}

// MarkDeviceListsChanged implements storage.KeyStore.
func (s *Keys) MarkDeviceListsChanged(ctx context.Context, userIDs []string) (int64, error) {
	var top int64
	err := s.write(ctx, func(t *tx) error {
		w := newKeysWrite(t)
		for _, userID := range dedupe(userIDs) {
			updates, err := w.bumpDeviceList(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 ` + inList + ` AND stream_pos > ?`
		args := []any{jsonList(chunk), from}
		if to > 0 {
			query += ` AND stream_pos <= ?`
			args = append(args, to)
		}

		rows, err := s.db.QueryContext(ctx, query, args...)
		if err != nil {
			return nil, fmt.Errorf("read device list changes: %w", err)
		}
		for rows.Next() {
			var c storage.DeviceListChange
			if err := rows.Scan(&c.UserID, &c.StreamPos); err != nil {
				_ = rows.Close()
				return nil, fmt.Errorf("read device list changes: %w", err)
			}
			out = append(out, c)
		}
		_ = rows.Close()
		if err := rows.Err(); err != nil {
			return nil, fmt.Errorf("read device list changes: %w", err)
		}
	}
	slices.SortFunc(out, func(a, b storage.DeviceListChange) int { return cmp.Compare(a.StreamPos, b.StreamPos) })
	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, `SELECT stream_pos FROM keyserver_device_lists WHERE user_id = ?`, 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 --------------------------------------------------------------

// SetDeviceKeys implements storage.KeyStore.
func (s *Keys) SetDeviceKeys(ctx context.Context, keys storage.DeviceKeys) (storage.DeviceListUpdate, error) {
	var update storage.DeviceListUpdate
	err := s.write(ctx, func(t *tx) error {
		var existing string
		err := t.queryRow(`SELECT key_json FROM keyserver_device_keys WHERE user_id = ? AND device_id = ?`,
			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)
		}

		t.put("keyserver_device_keys", rec{"user_id": keys.UserID, "device_id": keys.DeviceID, "key_json": string(keys.JSON)})
		updates, err := newKeysWrite(t).bumpDeviceList(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 ` + inList
		args = append(args, jsonList(deviceIDs))
	}
	query += ` ORDER BY device_id`

	rows, err := s.db.QueryContext(ctx, 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.write(ctx, func(t *tx) error {
		where := `user_id = ?`
		args := []any{userID}
		if len(deviceIDs) > 0 {
			where += ` AND device_id IN ` + inList
			args = append(args, jsonList(dedupe(deviceIDs)))
		}

		// The devices that lose identity keys are the ones whose deletion
		// moves the stream.
		had, err := s.queryStringsTx(t, `SELECT device_id FROM keyserver_device_keys WHERE `+where+` ORDER BY device_id`, args...)
		if err != nil {
			return fmt.Errorf("read the devices of %q: %w", userID, err)
		}

		for _, table := range []string{"keyserver_device_keys", "keyserver_one_time_keys", "keyserver_fallback_keys"} {
			if _, err := t.delWhere(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 = newKeysWrite(t).bumpDeviceList(userID, had)
		return err
	})
	return updates, err
}

// queryStringsTx runs a one-column query inside a write.
func (s *Keys) queryStringsTx(t *tx, query string, args ...any) ([]string, error) {
	rows, err := t.query(query, args...)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()
	var out []string
	for rows.Next() {
		var v string
		if err := rows.Scan(&v); err != nil {
			return nil, err
		}
		out = append(out, v)
	}
	return out, rows.Err()
}

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

// AddOneTimeKeys implements storage.KeyStore. The same key again is not an
// error; another key under the same ID is a conflict, and nothing is
// stored.
func (s *Keys) AddOneTimeKeys(ctx context.Context, userID, deviceID string, keys []storage.OneTimeKey) error {
	if len(keys) == 0 {
		return nil
	}
	return s.write(ctx, func(t *tx) error {
		ids := make([]string, len(keys))
		for i, k := range keys {
			ids[i] = k.KeyID
		}
		held := map[string]string{}
		rows, err := t.query(`SELECT key_id, key_json FROM keyserver_one_time_keys
			WHERE user_id = ? AND device_id = ? AND key_id IN `+inList, userID, deviceID, jsonList(ids))
		if err != nil {
			return fmt.Errorf("read the one-time keys of %s/%s: %w", userID, deviceID, err)
		}
		for rows.Next() {
			var id, keyJSON string
			if err := rows.Scan(&id, &keyJSON); err != nil {
				_ = rows.Close()
				return fmt.Errorf("read the one-time keys of %s/%s: %w", userID, deviceID, err)
			}
			held[id] = keyJSON
		}
		_ = rows.Close()
		if err := rows.Err(); err != nil {
			return fmt.Errorf("read the one-time keys of %s/%s: %w", userID, deviceID, err)
		}

		for _, k := range keys {
			if existing, ok := held[k.KeyID]; ok {
				if existing == string(k.JSON) {
					continue
				}
				return fmt.Errorf("%w: %s", storage.ErrKeyConflict, k.KeyID)
			}
			t.put("keyserver_one_time_keys", rec{
				"user_id": userID, "device_id": deviceID, "key_id": k.KeyID,
				"algorithm": k.Algorithm(), "key_json": string(k.JSON),
			})
			held[k.KeyID] = string(k.JSON)
		}
		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, `SELECT algorithm, COUNT(*) FROM keyserver_one_time_keys
		WHERE user_id = ? AND device_id = ? GROUP BY algorithm`, 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()
}

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

// SetFallbackKeys implements storage.KeyStore. The same key again keeps its
// used flag; a new one starts unused.
func (s *Keys) SetFallbackKeys(ctx context.Context, userID, deviceID string, keys []storage.OneTimeKey) error {
	if len(keys) == 0 {
		return nil
	}
	return s.write(ctx, func(t *tx) error {
		for _, k := range keys {
			var existingID, existingJSON string
			err := t.queryRow(selectFallbackKeySQL, userID, deviceID, k.Algorithm()).Scan(&existingID, &existingJSON)
			switch {
			case err == nil && existingID == k.KeyID && existingJSON == string(k.JSON):
				continue
			case err != nil && !errors.Is(err, sql.ErrNoRows):
				return fmt.Errorf("read the fallback key of %s/%s: %w", userID, deviceID, err)
			}
			t.put("keyserver_fallback_keys", rec{
				"user_id": userID, "device_id": deviceID, "algorithm": k.Algorithm(),
				"key_id": k.KeyID, "key_json": string(k.JSON), "used": false,
			})
		}
		return nil
	})
}

// UnusedFallbackKeyTypes implements storage.KeyStore.
func (s *Keys) UnusedFallbackKeyTypes(ctx context.Context, userID, deviceID string) ([]string, error) {
	out, err := s.queryStrings(ctx, `SELECT algorithm FROM keyserver_fallback_keys
		WHERE user_id = ? AND device_id = ? AND used = ? ORDER BY algorithm`, userID, deviceID, false)
	if err != nil {
		return nil, fmt.Errorf("read the fallback keys of %s/%s: %w", userID, deviceID, err)
	}
	return out, nil
}

// ClaimKey implements storage.KeyStore. The read and the delete happen under
// the write lock, so a claim cannot lose the key it read to another one.
func (s *Keys) ClaimKey(ctx context.Context, userID, deviceID, algorithm string) (storage.ClaimedKey, bool, error) {
	var (
		claimed storage.ClaimedKey
		found   bool
	)
	err := s.write(ctx, func(t *tx) error {
		var keyID, keyJSON string
		err := t.queryRow(`SELECT key_id, key_json FROM keyserver_one_time_keys
			WHERE user_id = ? AND device_id = ? AND algorithm = ?
			ORDER BY key_id LIMIT 1`, userID, deviceID, algorithm).Scan(&keyID, &keyJSON)
		switch {
		case err == nil:
			t.del("keyserver_one_time_keys", rowID("keyserver_one_time_keys", userID, deviceID, keyID))
			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 = t.queryRow(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)
		}
		t.put("keyserver_fallback_keys", rec{"user_id": userID, "device_id": deviceID, "algorithm": algorithm, "used": true})
		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
	}
	return claimed, found, nil
}

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

// SetCrossSigningKeys implements storage.KeyStore. A nil key deletes that
// purpose.
func (s *Keys) SetCrossSigningKeys(ctx context.Context, userID string, keys map[string][]byte) (storage.DeviceListUpdate, error) {
	var update storage.DeviceListUpdate
	err := s.write(ctx, func(t *tx) error {
		for _, purpose := range sortedKeys(keys) {
			keyJSON := keys[purpose]
			if keyJSON == nil {
				if _, err := t.delKey("keyserver_cross_signing_keys", userID, purpose); err != nil {
					return fmt.Errorf("store the %s key of %q: %w", purpose, userID, err)
				}
				continue
			}
			t.put("keyserver_cross_signing_keys", rec{"user_id": userID, "purpose": purpose, "key_json": string(keyJSON)})
		}
		updates, err := newKeysWrite(t).bumpDeviceList(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, `SELECT purpose, key_json FROM keyserver_cross_signing_keys WHERE user_id = ?`, 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 ---------------------------------------------------------------

// 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.write(ctx, func(t *tx) error {
		w := newKeysWrite(t)
		targets := map[string][]string{}
		for _, sig := range sigs {
			t.put("keyserver_signatures", rec{
				"target_user_id": sig.TargetUserID, "target_key_id": sig.TargetKeyID,
				"signer_user_id": sig.SignerUserID, "signer_key_id": sig.SignerKeyID,
				"signature": sig.Signature,
			})
			if !slices.Contains(targets[sig.TargetUserID], sig.TargetKeyID) {
				targets[sig.TargetUserID] = append(targets[sig.TargetUserID], sig.TargetKeyID)
			}
		}
		for _, userID := range sortedKeys(targets) {
			moved, err := w.bumpDeviceList(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, `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`, 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 ---------------------------------------------------------------

// parseBackupVersion reads a version number. A version that is not one is a
// version that does not exist.
func parseBackupVersion(version string) (int64, bool) {
	n, err := strconv.ParseInt(version, 10, 64)
	return n, err == nil && n > 0
}

// rowQuerier is what readBackup reads through: the pool, or a write.
type rowQuerier func(query string, args ...any) *sql.Row

func (s *Keys) poolRow(ctx context.Context) rowQuerier {
	return func(query string, args ...any) *sql.Row { return s.db.QueryRowContext(ctx, query, args...) }
}

// readBackup reads one live version, with its key count.
func readBackup(row rowQuerier, userID, version string) (storage.BackupVersion, int64, int64, error) {
	var n int64
	if version == "" {
		if err := row(`SELECT COALESCE(MAX(version), 0) FROM keyserver_backup_versions
			WHERE user_id = ? AND deleted = ?`, userID, false).Scan(&n); err != nil {
			return storage.BackupVersion{}, 0, 0, fmt.Errorf("read the backup version of %q: %w", userID, err)
		}
		if n == 0 {
			return storage.BackupVersion{}, 0, 0, storage.ErrNoSuchBackup
		}
	} else {
		var ok bool
		if n, ok = parseBackupVersion(version); !ok {
			return storage.BackupVersion{}, 0, 0, storage.ErrNoSuchBackup
		}
	}

	out := storage.BackupVersion{UserID: userID, Version: strconv.FormatInt(n, 10)}
	var (
		authData string
		etag     int64
	)
	err := row(`SELECT algorithm, auth_data, etag FROM keyserver_backup_versions
		WHERE user_id = ? AND version = ? AND deleted = ?`, userID, n, false).Scan(&out.Algorithm, &authData, &etag)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.BackupVersion{}, 0, 0, storage.ErrNoSuchBackup
	case err != nil:
		return storage.BackupVersion{}, 0, 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 := row(`SELECT COUNT(*) FROM keyserver_backup_keys WHERE user_id = ? AND version = ?`, userID, n).
		Scan(&out.Count); err != nil {
		return storage.BackupVersion{}, 0, 0, fmt.Errorf("count the keys of backup version %d of %q: %w", n, userID, err)
	}
	return out, n, etag, nil
}

// CreateBackup implements storage.KeyStore.
func (s *Keys) CreateBackup(ctx context.Context, userID, algorithm string, authData []byte) (string, error) {
	var version int64
	err := s.write(ctx, func(t *tx) error {
		if err := t.queryRow(`SELECT COALESCE(MAX(version), 0) FROM keyserver_backup_versions WHERE user_id = ?`,
			userID).Scan(&version); err != nil {
			return fmt.Errorf("read the backup versions of %q: %w", userID, err)
		}
		version++
		t.put("keyserver_backup_versions", rec{
			"user_id": userID, "version": version, "algorithm": algorithm,
			"auth_data": string(authData), "etag": int64(0), "deleted": false,
		})
		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 := readBackup(s.poolRow(ctx), 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
	}
	return s.write(ctx, func(t *tx) error {
		affected, err := t.setWhere("keyserver_backup_versions", rec{"auth_data": string(authData)},
			`user_id = ? AND version = ? AND deleted = ?`, userID, n, false)
		if err != nil {
			return fmt.Errorf("update backup version %d of %q: %w", n, userID, err)
		}
		if 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.write(ctx, func(t *tx) error {
		affected, err := t.setWhere("keyserver_backup_versions", rec{"deleted": true},
			`user_id = ? AND version = ? AND deleted = ?`, userID, n, false)
		if err != nil {
			return fmt.Errorf("delete backup version %d of %q: %w", n, userID, err)
		}
		if affected == 0 {
			return storage.ErrNoSuchBackup
		}
		if _, err := t.delWhere("keyserver_backup_keys", `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.write(ctx, func(t *tx) error {
		current, n, etag, err := readBackup(t.queryRow, userID, version)
		if err != nil {
			return err
		}

		changed := false
		written := map[string]storage.BackupKey{}
		for _, k := range keys {
			slot := k.RoomID + "\x00" + k.SessionID
			existing, have := written[slot]
			if !have {
				var sessionData string
				err := t.queryRow(`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 = ?`, userID, n, k.RoomID, k.SessionID).
					Scan(&existing.FirstMessageIndex, &existing.ForwardedCount, &existing.IsVerified, &sessionData)
				switch {
				case err == nil:
					have = true
				case !errors.Is(err, sql.ErrNoRows):
					return fmt.Errorf("read backed-up key %s/%s: %w", k.RoomID, k.SessionID, err)
				}
				if !have {
					current.Count++
				}
			}
			if have && !k.Replaces(existing) {
				written[slot] = existing
				continue
			}
			t.put("keyserver_backup_keys", rec{
				"user_id": userID, "version": n, "room_id": k.RoomID, "session_id": k.SessionID,
				"first_message_index": k.FirstMessageIndex, "forwarded_count": k.ForwardedCount,
				"is_verified": k.IsVerified, "session_data": string(k.SessionData),
			})
			written[slot] = k
			changed = true
		}
		if changed {
			etag++
			t.put("keyserver_backup_versions", rec{"user_id": userID, "version": n, "etag": etag})
		}
		current.ETag = strconv.FormatInt(etag, 10)
		out = current
		return nil
	})
	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 := `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 := readBackup(s.poolRow(ctx), userID, version)
	if err != nil {
		return nil, err
	}
	where, args := backupKeyScope(userID, n, roomID, sessionID)
	rows, err := s.db.QueryContext(ctx, `
		SELECT room_id, session_id, first_message_index, forwarded_count, is_verified, session_data
		FROM keyserver_backup_keys WHERE `+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.write(ctx, func(t *tx) error {
		current, n, etag, err := readBackup(t.queryRow, userID, version)
		if err != nil {
			return err
		}
		where, args := backupKeyScope(userID, n, roomID, sessionID)
		deleted, err := t.delWhere("keyserver_backup_keys", where, args...)
		if err != nil {
			return fmt.Errorf("delete keys from backup version %d of %q: %w", n, userID, err)
		}
		if deleted > 0 {
			etag++
			t.put("keyserver_backup_versions", rec{"user_id": userID, "version": n, "etag": etag})
		}
		current.Count -= deleted
		current.ETag = strconv.FormatInt(etag, 10)
		out = current
		return nil
	})
	return out, err
}

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

// dedupe returns the distinct values of a list, sorted.
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)