keyserver: cache remote users' device and cross-signing keys

POST /keys/query for a remote user who shares a room with a local user is now answered from a cache, fetched whole with GET /user/devices/{userId} on a miss. Incoming m.device_list_update EDUs are applied when their prev_id matches the cached stream_id; anything else marks the copy stale and it is fetched again in the background, and a stale copy is never served. m.signing_key_update replaces the cached cross-signing keys. A user's copy is dropped once no local user shares a room with them, by a keyserver listener on the roomserver's emit stage. Migration 0009 adds keyserver_remote_users, keyserver_remote_device_keys and keyserver_remote_cross_signing_keys for sqlite and postgres. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-17 10:36:08 +02:00
1 parent fc49f56339
commit 74f7bc8805
@@ -0,0 +1,315 @@
+package sqlstore
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"fmt"
+	"slices"
+
+	"github.com/dariuszjaniak/warp/server/internal/storage"
+)
+
+// The cache of remote users' keys: storage.KeyStore's RemoteKeys and the
+// methods after it.
+//
+// None of these writes touch the device list stream. A remote user's
+// position in it is moved by MarkDeviceListsChanged, which the keyserver
+// calls for every update that concerns a local client whether or not the
+// cache could take the update in: the stream says "ask again", and the
+// cache is only where the answer comes from.
+
+const (
+	selectRemoteUserSQL = `
+		SELECT stream_id, stale FROM keyserver_remote_users WHERE user_id = ?`
+
+	upsertRemoteUserSQL = `
+		INSERT INTO keyserver_remote_users (user_id, stream_id, stale) VALUES (?, ?, ?)
+		ON CONFLICT (user_id) DO UPDATE SET stream_id = excluded.stream_id, stale = excluded.stale`
+
+	selectRemoteDevicesSQL = `
+		SELECT device_id, display_name, key_json FROM keyserver_remote_device_keys
+		WHERE user_id = ? ORDER BY device_id`
+
+	upsertRemoteDeviceSQL = `
+		INSERT INTO keyserver_remote_device_keys (user_id, device_id, display_name, key_json)
+		VALUES (?, ?, ?, ?)
+		ON CONFLICT (user_id, device_id) DO UPDATE SET
+			display_name = excluded.display_name, key_json = excluded.key_json`
+
+	deleteRemoteDeviceSQL = `
+		DELETE FROM keyserver_remote_device_keys WHERE user_id = ? AND device_id = ?`
+
+	selectRemoteCrossSigningSQL = `
+		SELECT purpose, key_json FROM keyserver_remote_cross_signing_keys WHERE user_id = ?`
+
+	insertRemoteCrossSigningSQL = `
+		INSERT INTO keyserver_remote_cross_signing_keys (user_id, purpose, key_json) VALUES (?, ?, ?)`
+
+	deleteRemoteCrossSigningSQL = `
+		DELETE FROM keyserver_remote_cross_signing_keys WHERE user_id = ?`
+
+	selectRemoteUsersSQL = `
+		SELECT user_id FROM keyserver_remote_users ORDER BY user_id`
+)
+
+// remoteKeyTables is every table a cached user has rows in.
+var remoteKeyTables = []string{
+	"keyserver_remote_users", "keyserver_remote_device_keys", "keyserver_remote_cross_signing_keys",
+}
+
+// remoteMarker reads a cached user's marker row.
+func (s *Keys) remoteMarker(ctx context.Context, db rowQuerier, userID string) (streamID int64, stale, ok bool, err error) {
+	err = db.QueryRowContext(ctx, s.q(selectRemoteUserSQL), userID).Scan(&streamID, &stale)
+	switch {
+	case errors.Is(err, sql.ErrNoRows):
+		return 0, false, false, nil
+	case err != nil:
+		return 0, false, false, fmt.Errorf("read the cached keys of %q: %w", userID, err)
+	}
+	return streamID, stale, true, nil
+}
+
+// RemoteKeys implements storage.KeyStore.
+func (s *Keys) RemoteKeys(ctx context.Context, userID string) (storage.RemoteKeys, bool, error) {
+	out := storage.RemoteKeys{UserID: userID, CrossSigning: map[string][]byte{}}
+
+	// One transaction, so that the marker and the keys are one snapshot:
+	// a copy read halfway through its replacement would be served as
+	// fresh with half of another one's devices.
+	var found bool
+	err := s.tx(ctx, func(tx *sql.Tx) error {
+		var err error
+		out.StreamID, out.Stale, found, err = s.remoteMarker(ctx, tx, userID)
+		if err != nil || !found {
+			return err
+		}
+		if out.Devices, err = s.remoteDevices(ctx, tx, userID); err != nil {
+			return err
+		}
+		out.CrossSigning, err = s.remoteCrossSigning(ctx, tx, userID)
+		return err
+	})
+	if err != nil || !found {
+		return storage.RemoteKeys{}, false, err
+	}
+	return out, true, nil
+}
+
+func (s *Keys) remoteDevices(ctx context.Context, tx *sql.Tx, userID string) ([]storage.RemoteDevice, error) {
+	rows, err := tx.QueryContext(ctx, s.q(selectRemoteDevicesSQL), userID)
+	if err != nil {
+		return nil, fmt.Errorf("read the cached devices of %q: %w", userID, err)
+	}
+	defer func() { _ = rows.Close() }()
+
+	var out []storage.RemoteDevice
+	for rows.Next() {
+		var d storage.RemoteDevice
+		var keyJSON string
+		if err := rows.Scan(&d.DeviceID, &d.DisplayName, &keyJSON); err != nil {
+			return nil, fmt.Errorf("read the cached devices of %q: %w", userID, err)
+		}
+		d.JSON = []byte(keyJSON)
+		out = append(out, d)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("read the cached devices of %q: %w", userID, err)
+	}
+	return out, nil
+}
+
+func (s *Keys) remoteCrossSigning(ctx context.Context, tx *sql.Tx, userID string) (map[string][]byte, error) {
+	rows, err := tx.QueryContext(ctx, s.q(selectRemoteCrossSigningSQL), userID)
+	if err != nil {
+		return nil, fmt.Errorf("read the cached 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 cached cross-signing keys of %q: %w", userID, err)
+		}
+		out[purpose] = []byte(keyJSON)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("read the cached cross-signing keys of %q: %w", userID, err)
+	}
+	return out, nil
+}
+
+// SetRemoteKeys implements storage.KeyStore.
+func (s *Keys) SetRemoteKeys(ctx context.Context, keys storage.RemoteKeys) (bool, error) {
+	stored := false
+	err := s.tx(ctx, func(tx *sql.Tx) error {
+		held, _, ok, err := s.remoteMarker(ctx, tx, keys.UserID)
+		if err != nil {
+			return err
+		}
+		// A snapshot older than what is held was fetched while an update
+		// was on its way in. Storing it as fresh would lose that update
+		// until the next one showed the gap; leaving the copy as it is —
+		// stale, as the update left it — has the next query fetch again.
+		if ok && held > keys.StreamID {
+			return nil
+		}
+
+		if err := s.deleteRemote(ctx, tx, keys.UserID); err != nil {
+			return err
+		}
+		if _, err := tx.ExecContext(ctx, s.q(upsertRemoteUserSQL), keys.UserID, keys.StreamID, false); err != nil {
+			return fmt.Errorf("cache the keys of %q: %w", keys.UserID, err)
+		}
+		for _, d := range keys.Devices {
+			if _, err := tx.ExecContext(ctx, s.q(upsertRemoteDeviceSQL), keys.UserID, d.DeviceID, d.DisplayName, string(d.JSON)); err != nil {
+				return fmt.Errorf("cache the keys of %s/%s: %w", keys.UserID, d.DeviceID, err)
+			}
+		}
+		if err := s.insertRemoteCrossSigning(ctx, tx, keys.UserID, keys.CrossSigning); err != nil {
+			return err
+		}
+		stored = true
+		return nil
+	})
+	return stored, err
+}
+
+// UpdateRemoteDevice implements storage.KeyStore.
+func (s *Keys) UpdateRemoteDevice(ctx context.Context, u storage.RemoteDeviceUpdate) (storage.RemoteUpdateResult, error) {
+	var result storage.RemoteUpdateResult
+	err := s.tx(ctx, func(tx *sql.Tx) error {
+		held, stale, ok, err := s.remoteMarker(ctx, tx, u.UserID)
+		if err != nil {
+			return err
+		}
+
+		// A stale copy keeps the highest position it has heard of, so that
+		// a snapshot fetched before this update cannot be stored over it.
+		markStale := func(r storage.RemoteUpdateResult) error {
+			result = r
+			if _, err := tx.ExecContext(ctx, s.q(upsertRemoteUserSQL), u.UserID, max(held, u.StreamID), true); err != nil {
+				return fmt.Errorf("mark the cached keys of %q stale: %w", u.UserID, err)
+			}
+			return nil
+		}
+
+		switch {
+		case !ok:
+			return markStale(storage.RemoteUpdateNotCached)
+		case stale:
+			return markStale(storage.RemoteUpdateStale)
+		case u.StreamID <= held:
+			result = storage.RemoteUpdateDuplicate
+			return nil
+		case !followsOn(held, u.PrevIDs):
+			return markStale(storage.RemoteUpdateStale)
+		case !u.Deleted && u.JSON == nil:
+			// An update that says a device changed without saying how
+			// cannot be applied; the copy has to be fetched again.
+			return markStale(storage.RemoteUpdateStale)
+		}
+
+		if u.Deleted {
+			_, err = tx.ExecContext(ctx, s.q(deleteRemoteDeviceSQL), u.UserID, u.DeviceID)
+		} else {
+			_, err = tx.ExecContext(ctx, s.q(upsertRemoteDeviceSQL), u.UserID, u.DeviceID, u.DisplayName, string(u.JSON))
+		}
+		if err != nil {
+			return fmt.Errorf("apply an update to the cached keys of %s/%s: %w", u.UserID, u.DeviceID, err)
+		}
+		if _, err := tx.ExecContext(ctx, s.q(upsertRemoteUserSQL), u.UserID, u.StreamID, false); err != nil {
+			return fmt.Errorf("move the cached keys of %q: %w", u.UserID, err)
+		}
+		result = storage.RemoteUpdateApplied
+		return nil
+	})
+	return result, err
+}
+
+// followsOn reports whether an update with the given prev_id comes right
+// after the position held. An update with no prev_id is the first its
+// server ever sent for the user, which only follows on from nothing.
+func followsOn(held int64, prevIDs []int64) bool {
+	if len(prevIDs) == 0 {
+		return held == 0
+	}
+	return slices.Contains(prevIDs, held)
+}
+
+// SetRemoteCrossSigningKeys implements storage.KeyStore.
+func (s *Keys) SetRemoteCrossSigningKeys(ctx context.Context, userID string, keys map[string][]byte) (bool, error) {
+	stored := false
+	err := s.tx(ctx, func(tx *sql.Tx) error {
+		_, stale, ok, err := s.remoteMarker(ctx, tx, userID)
+		if err != nil || !ok || stale {
+			return err
+		}
+		if _, err := tx.ExecContext(ctx, s.q(deleteRemoteCrossSigningSQL), userID); err != nil {
+			return fmt.Errorf("replace the cached cross-signing keys of %q: %w", userID, err)
+		}
+		if err := s.insertRemoteCrossSigning(ctx, tx, userID, keys); err != nil {
+			return err
+		}
+		stored = true
+		return nil
+	})
+	return stored, err
+}
+
+func (s *Keys) insertRemoteCrossSigning(ctx context.Context, tx *sql.Tx, userID string, keys map[string][]byte) error {
+	for _, purpose := range sortedKeys(keys) {
+		if keys[purpose] == nil {
+			continue
+		}
+		if _, err := tx.ExecContext(ctx, s.q(insertRemoteCrossSigningSQL), userID, purpose, string(keys[purpose])); err != nil {
+			return fmt.Errorf("cache the %s key of %q: %w", purpose, userID, err)
+		}
+	}
+	return nil
+}
+
+// RemoteKeyUsers implements storage.KeyStore.
+func (s *Keys) RemoteKeyUsers(ctx context.Context) ([]string, error) {
+	rows, err := s.db.QueryContext(ctx, s.q(selectRemoteUsersSQL))
+	if err != nil {
+		return nil, fmt.Errorf("read the users with cached keys: %w", err)
+	}
+	defer func() { _ = rows.Close() }()
+
+	var out []string
+	for rows.Next() {
+		var userID string
+		if err := rows.Scan(&userID); err != nil {
+			return nil, fmt.Errorf("read the users with cached keys: %w", err)
+		}
+		out = append(out, userID)
+	}
+	return out, rows.Err()
+}
+
+// DeleteRemoteKeys implements storage.KeyStore.
+func (s *Keys) DeleteRemoteKeys(ctx context.Context, userIDs []string) error {
+	if len(userIDs) == 0 {
+		return nil
+	}
+	return s.tx(ctx, func(tx *sql.Tx) error {
+		for _, userID := range dedupe(userIDs) {
+			if err := s.deleteRemote(ctx, tx, userID); err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+}
+
+// deleteRemote drops every row cached for a user.
+func (s *Keys) deleteRemote(ctx context.Context, tx *sql.Tx, userID string) error {
+	for _, table := range remoteKeyTables {
+		if _, err := tx.ExecContext(ctx, s.q(`DELETE FROM `+table+` WHERE user_id = ?`), userID); err != nil {
+			return fmt.Errorf("delete from %s for %q: %w", table, userID, err)
+		}
+	}
+	return nil
+}