Rust: Query backup on fail to decrypt similar to libolm (#3711)

* Refactor key backup recovery to prepare for rust * rust backup restore support * map decryption errors correctly from rust * query backup on fail to decrypt
2025-08-07 23:02:56 +03:00 · 2023-09-14 12:10:53 +02:00
parent 6d118008b6
commit 40168d4419
4 changed files with 78 additions and 58 deletions
--- a/spec/integ/crypto/megolm-backup.spec.ts
+++ b/spec/integ/crypto/megolm-backup.spec.ts
@@ -18,61 +18,21 @@ import fetchMock from "fetch-mock-jest";
 import "fake-indexeddb/auto";
 import { IDBFactory } from "fake-indexeddb";

-import { createClient, CryptoEvent, ICreateClientOpts, IEvent, MatrixClient, TypedEventEmitter } from "../../../src";
+import { createClient, CryptoEvent, ICreateClientOpts, MatrixClient, TypedEventEmitter } from "../../../src";
 import { SyncResponder } from "../../test-utils/SyncResponder";
 import { E2EKeyReceiver } from "../../test-utils/E2EKeyReceiver";
 import { E2EKeyResponder } from "../../test-utils/E2EKeyResponder";
 import { mockInitialApiRequests } from "../../test-utils/mockEndpoints";
 import { awaitDecryption, CRYPTO_BACKENDS, InitCrypto, syncPromise } from "../../test-utils/test-utils";
 import * as testData from "../../test-utils/test-data";
-import { KeyBackupInfo, KeyBackupSession } from "../../../src/crypto-api/keybackup";
+import { KeyBackupInfo } from "../../../src/crypto-api/keybackup";
 import { IKeyBackup } from "../../../src/crypto/backup";

-const ROOM_ID = "!ROOM:ID";
+const ROOM_ID = testData.TEST_ROOM_ID;

 /** The homeserver url that we give to the test client, and where we intercept /sync, /keys, etc requests. */
 const TEST_HOMESERVER_URL = "https://alice-server.com";

-const SESSION_ID = "o+21hSjP+mgEmcfdslPsQdvzWnkdt0Wyo00Kp++R8Kc";
-
-const ENCRYPTED_EVENT: Partial<IEvent> = {
-    type: "m.room.encrypted",
-    content: {
-        algorithm: "m.megolm.v1.aes-sha2",
-        sender_key: "SENDER_CURVE25519",
-        session_id: SESSION_ID,
-        ciphertext:
-            "AwgAEjD+VwXZ7PoGPRS/H4kwpAsMp/g+WPvJVtPEKE8fmM9IcT/N" +
-            "CiwPb8PehecDKP0cjm1XO88k6Bw3D17aGiBHr5iBoP7oSw8CXULXAMTkBl" +
-            "mkufRQq2+d0Giy1s4/Cg5n13jSVrSb2q7VTSv1ZHAFjUCsLSfR0gxqcQs",
-    },
-    room_id: "!ROOM:ID",
-    event_id: "$event1",
-    origin_server_ts: 1507753886000,
-};
-
-const CURVE25519_KEY_BACKUP_DATA: KeyBackupSession = {
-    first_message_index: 0,
-    forwarded_count: 0,
-    is_verified: false,
-    session_data: {
-        ciphertext:
-            "2z2M7CZ+azAiTHN1oFzZ3smAFFt+LEOYY6h3QO3XXGdw" +
-            "6YpNn/gpHDO6I/rgj1zNd4FoTmzcQgvKdU8kN20u5BWRHxaHTZ" +
-            "Slne5RxE6vUdREsBgZePglBNyG0AogR/PVdcrv/v18Y6rLM5O9" +
-            "SELmwbV63uV9Kuu/misMxoqbuqEdG7uujyaEKtjlQsJ5MGPQOy" +
-            "Syw7XrnesSwF6XWRMxcPGRV0xZr3s9PI350Wve3EncjRgJ9IGF" +
-            "ru1bcptMqfXgPZkOyGvrphHoFfoK7nY3xMEHUiaTRfRIjq8HNV" +
-            "4o8QY1qmWGnxNBQgOlL8MZlykjg3ULmQ3DtFfQPj/YYGS3jzxv" +
-            "C+EBjaafmsg+52CTeK3Rswu72PX450BnSZ1i3If4xWAUKvjTpe" +
-            "Ug5aDLqttOv1pITolTJDw5W/SD+b5rjEKg1CFCHGEGE9wwV3Nf" +
-            "QHVCQL+dfpd7Or0poy4dqKMAi3g0o3Tg7edIF8d5rREmxaALPy" +
-            "iie8PHD8mj/5Y0GLqrac4CD6+Mop7eUTzVovprjg",
-        mac: "5lxYBHQU80M",
-        ephemeral: "/Bn0A4UMFwJaDDvh0aEk1XZj3k1IfgCxgFY9P9a0b14",
-    },
-};
-
 const TEST_USER_ID = "@alice:localhost";
 const TEST_DEVICE_ID = "xzcvb";

@@ -138,7 +98,8 @@ function mockUploadEmitter(
 describe.each(Object.entries(CRYPTO_BACKENDS))("megolm-keys backup (%s)", (backend: string, initCrypto: InitCrypto) => {
    // oldBackendOnly is an alternative to `it` or `test` which will skip the test if we are running against the
    // Rust backend. Once we have full support in the rust sdk, it will go away.
-    const oldBackendOnly = backend === "rust-sdk" ? test.skip : test;
+    // const oldBackendOnly = backend === "rust-sdk" ? test.skip : test;
+    // const newBackendOnly = backend === "libolm" ? test.skip : test;

    let aliceClient: MatrixClient;
    /** an object which intercepts `/sync` requests on the test homeserver */
@@ -188,21 +149,24 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("megolm-keys backup (%s)", (backe
        return client;
    }

-    oldBackendOnly("Alice checks key backups when receiving a message she can't decrypt", async function () {
+    it("Alice checks key backups when receiving a message she can't decrypt", async function () {
        const syncResponse = {
            next_batch: 1,
            rooms: {
                join: {
                    [ROOM_ID]: {
                        timeline: {
-                            events: [ENCRYPTED_EVENT],
+                            events: [testData.ENCRYPTED_EVENT],
                        },
                    },
                },
            },
        };

-        fetchMock.get("express:/_matrix/client/v3/room_keys/keys/:room_id/:session_id", CURVE25519_KEY_BACKUP_DATA);
+        fetchMock.get(
+            "express:/_matrix/client/v3/room_keys/keys/:room_id/:session_id",
+            testData.CURVE25519_KEY_BACKUP_DATA,
+        );
        fetchMock.get("path:/_matrix/client/v3/room_keys/version", testData.SIGNED_BACKUP_DATA);

        aliceClient = await initTestClient();
@@ -216,7 +180,7 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("megolm-keys backup (%s)", (backe
        // XXX: should we automatically re-check after a device becomes verified?
        await waitForDeviceList();
        await aliceCrypto.setDeviceVerified(testData.TEST_USER_ID, testData.TEST_DEVICE_ID);
-        await aliceClient.checkKeyBackup();
+        await aliceClient.getCrypto()!.checkKeyBackupAndEnable();

        // Now, send Alice a message that she won't be able to decrypt, and check that she fetches the key from the backup.
        syncResponder.sendOrQueueSyncResponse(syncResponse);
@@ -225,7 +189,8 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("megolm-keys backup (%s)", (backe
        const room = aliceClient.getRoom(ROOM_ID)!;
        const event = room.getLiveTimeline().getEvents()[0];
        await awaitDecryption(event, { waitOnDecryptionFailure: true });
-        expect(event.getContent()).toEqual("testytest");
+
+        expect(event.getContent()).toEqual(testData.CLEAR_EVENT.content);
    });

    describe("recover from backup", () => {
--- a/spec/test-utils/test-data/generate-test-data.py
+++ b/spec/test-utils/test-data/generate-test-data.py
@@ -456,7 +456,6 @@ def symetric_ratchet_step_of_megolm_key(previous: dict , megolm_private_key: ed2

    return megolm_export

-
 def encrypt_megolm_key_for_backup(session_data: dict, backup_public_key: x25519.X25519PublicKey) -> dict:

    """
--- a/spec/test-utils/test-utils.ts
+++ b/spec/test-utils/test-utils.ts
@@ -521,10 +521,19 @@ export async function awaitDecryption(
    }

    return new Promise((resolve) => {
+        if (waitOnDecryptionFailure) {
+            event.on(MatrixEventEvent.Decrypted, (ev, err) => {
+                logger.log(`${Date.now()}: MatrixEventEvent.Decrypted for event ${event.getId()}: ${err ?? "success"}`);
+                if (!err) {
+                    resolve(ev);
+                }
+            });
+        } else {
            event.once(MatrixEventEvent.Decrypted, (ev, err) => {
                logger.log(`${Date.now()}: MatrixEventEvent.Decrypted for event ${event.getId()}: ${err ?? "success"}`);
                resolve(ev);
            });
+        }
    });
 }

--- a/src/rust-crypto/rust-crypto.ts
+++ b/src/rust-crypto/rust-crypto.ts
@@ -25,11 +25,11 @@ import { Room } from "../models/room";
 import { RoomMember } from "../models/room-member";
 import { BackupDecryptor, CryptoBackend, OnSyncCompletedData } from "../common-crypto/CryptoBackend";
 import { logger } from "../logger";
-import { IHttpOpts, MatrixHttpApi, Method } from "../http-api";
+import { ClientPrefix, IHttpOpts, MatrixHttpApi, Method } from "../http-api";
 import { RoomEncryptor } from "./RoomEncryptor";
 import { OutgoingRequest, OutgoingRequestProcessor } from "./OutgoingRequestProcessor";
 import { KeyClaimManager } from "./KeyClaimManager";
-import { MapWithDefault } from "../utils";
+import { MapWithDefault, encodeUri } from "../utils";
 import {
    BackupTrustInfo,
    BootstrapCrossSigningOpts,
@@ -47,6 +47,7 @@ import {
    ImportRoomKeysOpts,
    KeyBackupCheck,
    KeyBackupInfo,
+    KeyBackupSession,
    UserVerificationStatus,
    VerificationRequest,
 } from "../crypto-api";
@@ -78,6 +79,8 @@ interface ISignableObject {
    unsigned?: object;
 }

+const KEY_BACKUP_CHECK_RATE_LIMIT = 5000; // ms
+
 /**
 * An implementation of {@link CryptoBackend} using the Rust matrix-sdk-crypto.
 *
@@ -102,6 +105,8 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, RustCryptoEv
    private crossSigningIdentity: CrossSigningIdentity;
    private readonly backupManager: RustBackupManager;

+    private sessionLastCheckAttemptedTime: Record<string, number> = {}; // When did we last try to check the server for a given session id?
+
    private readonly reemitter = new TypedReEmitter<RustCryptoEvents, RustCryptoEventMap>(this);

    public constructor(
@@ -130,7 +135,7 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, RustCryptoEv
        super();
        this.outgoingRequestProcessor = new OutgoingRequestProcessor(olmMachine, http);
        this.keyClaimManager = new KeyClaimManager(olmMachine, this.outgoingRequestProcessor);
-        this.eventDecryptor = new EventDecryptor(olmMachine);
+        this.eventDecryptor = new EventDecryptor(olmMachine, this);

        this.backupManager = new RustBackupManager(olmMachine, http, this.outgoingRequestProcessor);
        this.reemitter.reEmit(this.backupManager, [
@@ -142,6 +147,46 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, RustCryptoEv
        this.crossSigningIdentity = new CrossSigningIdentity(olmMachine, this.outgoingRequestProcessor, secretStorage);
    }

+    /**
+     * Attempts to retrieve a session from a key backup, if enough time
+     * has elapsed since the last check for this session id.
+     */
+    public async queryKeyBackupRateLimited(targetRoomId: string, targetSessionId: string): Promise<void> {
+        const backupKeys: RustSdkCryptoJs.BackupKeys = await this.olmMachine.getBackupKeys();
+        if (!backupKeys.decryptionKey) return;
+        const version = backupKeys.backupVersion;
+
+        const now = new Date().getTime();
+        if (
+            !this.sessionLastCheckAttemptedTime[targetSessionId!] ||
+            now - this.sessionLastCheckAttemptedTime[targetSessionId!] > KEY_BACKUP_CHECK_RATE_LIMIT
+        ) {
+            this.sessionLastCheckAttemptedTime[targetSessionId!] = now;
+
+            const path = encodeUri("/room_keys/keys/$roomId/$sessionId", {
+                $roomId: targetRoomId,
+                $sessionId: targetSessionId,
+            });
+
+            const res = await this.http.authedRequest<KeyBackupSession>(Method.Get, path, { version }, undefined, {
+                prefix: ClientPrefix.V3,
+            });
+
+            if (this.stopped) return;
+
+            const backupDecryptor = new RustBackupDecryptor(backupKeys.decryptionKey);
+            if (res) {
+                const sessionsToImport: Record<string, KeyBackupSession> = {};
+                sessionsToImport[targetSessionId] = res;
+                const keys = await backupDecryptor.decryptSessions(sessionsToImport);
+                for (const k of keys) {
+                    k.room_id = targetRoomId!;
+                }
+                await this.importRoomKeys(keys);
+            }
+        }
+    }
+
    /**
     * Return the OlmMachine only if {@link RustCrypto#stop} has not been called.
     *
@@ -1383,7 +1428,7 @@ class EventDecryptor {
        () => new MapWithDefault<string, Set<MatrixEvent>>(() => new Set()),
    );

-    public constructor(private readonly olmMachine: RustSdkCryptoJs.OlmMachine) {}
+    public constructor(private readonly olmMachine: RustSdkCryptoJs.OlmMachine, private readonly crypto: RustCrypto) {}

    public async attemptEventDecryption(event: MatrixEvent): Promise<IEventDecryptionResult> {
        logger.info("Attempting decryption of event", event);
@@ -1431,6 +1476,7 @@ class EventDecryptor {
                                session: content.sender_key + "|" + content.session_id,
                            },
                        );
+                        this.crypto.queryKeyBackupRateLimited(event.getRoomId()!, event.getWireContent().session_id!);
                        break;
                    }
                    case RustSdkCryptoJs.DecryptionErrorCode.UnknownMessageIndex: {
@@ -1441,6 +1487,7 @@ class EventDecryptor {
                                session: content.sender_key + "|" + content.session_id,
                            },
                        );
+                        this.crypto.queryKeyBackupRateLimited(event.getRoomId()!, event.getWireContent().session_id!);
                        break;
                    }
                    // We don't map MismatchedIdentityKeys for now, as there is no equivalent in legacy.