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

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", () => {

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:
 
     """

spec/test-utils/test-utils.ts

@@ -521,10 +521,19 @@ export async function awaitDecryption(
     }
 
     return new Promise((resolve) => {
-        event.once(MatrixEventEvent.Decrypted, (ev, err) => {
+        if (waitOnDecryptionFailure) {
-            logger.log(`${Date.now()}: MatrixEventEvent.Decrypted for event ${event.getId()}: ${err ?? "success"}`);
+            event.on(MatrixEventEvent.Decrypted, (ev, err) => {
-            resolve(ev);
+                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);
+            });
+        }
     });
 }
 

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.