Element-R: handle events which arrive before their keys (#3230)

* minor cleanups to the crypto tests

mostly, this is about using `testUtils.awaitDecryption` rather than custom
code. Some other cleanups too.

* Keep a record of events which are missing their keys

* Retry event decryption when we receive megolm keys

package.json

@@ -55,7 +55,7 @@
     ],
     "dependencies": {
         "@babel/runtime": "^7.12.5",
-        "@matrix-org/matrix-sdk-crypto-js": "^0.1.0-alpha.5",
+        "@matrix-org/matrix-sdk-crypto-js": "^0.1.0-alpha.6",
         "another-json": "^0.2.0",
         "bs58": "^5.0.0",
         "content-type": "^1.0.4",

spec/integ/crypto.spec.ts

@@ -624,10 +624,8 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("crypto (%s)", (backend: string,
         expect(decryptedEvent.getContent().body).toEqual("42");
     });
 
-    oldBackendOnly("Alice receives a megolm message before the session keys", async () => {
+    it("Alice receives a megolm message before the session keys", async () => {
         expectAliceKeyQuery({ device_keys: { "@alice:localhost": {} }, failures: {} });
-
-        // https://github.com/vector-im/element-web/issues/2273
         await startClientAndAwaitFirstSync();
 
         // if we're using the old crypto impl, stub out some methods in the device manager.
@@ -667,7 +665,11 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("crypto (%s)", (backend: string,
         await syncPromise(aliceClient);
 
         const room = aliceClient.getRoom(ROOM_ID)!;
-        expect(room.getLiveTimeline().getEvents()[0].getContent().msgtype).toEqual("m.bad.encrypted");
+        const event = room.getLiveTimeline().getEvents()[0];
+
+        // wait for a first attempt at decryption: should fail
+        await testUtils.awaitDecryption(event);
+        expect(event.getContent().msgtype).toEqual("m.bad.encrypted");
 
         // now she gets the room_key event
         syncResponder.sendOrQueueSyncResponse({
@@ -678,20 +680,8 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("crypto (%s)", (backend: string,
         });
         await syncPromise(aliceClient);
 
-        const event = room.getLiveTimeline().getEvents()[0];
+        await testUtils.awaitDecryption(event, { waitOnDecryptionFailure: true });
-
+        expect(event.getContent().body).toEqual("42");
-        let decryptedEvent: MatrixEvent;
-        if (event.getContent().msgtype != "m.bad.encrypted") {
-            decryptedEvent = event;
-        } else {
-            decryptedEvent = await new Promise<MatrixEvent>((resolve) => {
-                event.once(MatrixEventEvent.Decrypted, (ev) => {
-                    logger.log(`${Date.now()} event ${event.getId()} now decrypted`);
-                    resolve(ev);
-                });
-            });
-        }
-        expect(decryptedEvent.getContent().body).toEqual("42");
     });
 
     it("Alice gets a second room_key message", async () => {

spec/test-utils/test-utils.ts

@@ -375,17 +375,17 @@ export async function awaitDecryption(
     // already
     if (event.getClearContent() !== null) {
         if (waitOnDecryptionFailure && event.isDecryptionFailure()) {
-            logger.log(`${Date.now()} event ${event.getId()} got decryption error; waiting`);
+            logger.log(`${Date.now()}: event ${event.getId()} got decryption error; waiting`);
         } else {
             return event;
         }
     } else {
-        logger.log(`${Date.now()} event ${event.getId()} is not yet decrypted; waiting`);
+        logger.log(`${Date.now()}: event ${event.getId()} is not yet decrypted; waiting`);
     }
 
     return new Promise((resolve) => {
-        event.once(MatrixEventEvent.Decrypted, (ev) => {
+        event.once(MatrixEventEvent.Decrypted, (ev, err) => {
-            logger.log(`${Date.now()} event ${event.getId()} now decrypted`);
+            logger.log(`${Date.now()}: MatrixEventEvent.Decrypted for event ${event.getId()}: ${err ?? "success"}`);
             resolve(ev);
         });
     });

src/rust-crypto/index.ts

@@ -39,6 +39,9 @@ export async function initRustCrypto(
     // TODO: use the pickle key for the passphrase
     const olmMachine = await RustSdkCryptoJs.OlmMachine.initialize(u, d, RUST_SDK_STORE_PREFIX, "test pass");
     const rustCrypto = new RustCrypto(olmMachine, http, userId, deviceId);
+    await olmMachine.registerRoomKeyUpdatedCallback((sessions: RustSdkCryptoJs.RoomKeyInfo[]) =>
+        rustCrypto.onRoomKeysUpdated(sessions),
+    );
 
     logger.info("Completed rust crypto-sdk setup");
     return rustCrypto;

src/rust-crypto/rust-crypto.ts

@@ -1,5 +1,5 @@
 /*
-Copyright 2022 The Matrix.org Foundation C.I.C.
+Copyright 2022-2023 The Matrix.org Foundation C.I.C.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -29,6 +29,7 @@ import { DeviceTrustLevel, UserTrustLevel } from "../crypto/CrossSigning";
 import { RoomEncryptor } from "./RoomEncryptor";
 import { OutgoingRequest, OutgoingRequestProcessor } from "./OutgoingRequestProcessor";
 import { KeyClaimManager } from "./KeyClaimManager";
+import { MapWithDefault } from "../utils";
 
 /**
  * An implementation of {@link CryptoBackend} using the Rust matrix-sdk-crypto.
@@ -45,6 +46,7 @@ export class RustCrypto implements CryptoBackend {
     /** mapping of roomId → encryptor class */
     private roomEncryptors: Record<string, RoomEncryptor> = {};
 
+    private eventDecryptor: EventDecryptor;
     private keyClaimManager: KeyClaimManager;
     private outgoingRequestProcessor: OutgoingRequestProcessor;
 
@@ -56,6 +58,7 @@ export class RustCrypto implements CryptoBackend {
     ) {
         this.outgoingRequestProcessor = new OutgoingRequestProcessor(olmMachine, http);
         this.keyClaimManager = new KeyClaimManager(olmMachine, this.outgoingRequestProcessor);
+        this.eventDecryptor = new EventDecryptor(olmMachine);
     }
 
     ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -101,23 +104,7 @@ export class RustCrypto implements CryptoBackend {
             // through decryptEvent and hence get rid of this case.
             throw new Error("to-device event was not decrypted in preprocessToDeviceMessages");
         }
-        const res = (await this.olmMachine.decryptRoomEvent(
+        return await this.eventDecryptor.attemptEventDecryption(event);
-            JSON.stringify({
-                event_id: event.getId(),
-                type: event.getWireType(),
-                sender: event.getSender(),
-                state_key: event.getStateKey(),
-                content: event.getWireContent(),
-                origin_server_ts: event.getTs(),
-            }),
-            new RustSdkCryptoJs.RoomId(event.getRoomId()!),
-        )) as RustSdkCryptoJs.DecryptedRoomEvent;
-        return {
-            clearEvent: JSON.parse(res.event),
-            claimedEd25519Key: res.senderClaimedEd25519Key,
-            senderCurve25519Key: res.senderCurve25519Key,
-            forwardingCurve25519KeyChain: res.forwardingCurve25519KeyChain,
-        };
     }
 
     public getEventEncryptionInfo(event: MatrixEvent): IEncryptedEventInfo {
@@ -303,6 +290,43 @@ export class RustCrypto implements CryptoBackend {
         enc.onRoomMembership(member);
     }
 
+    /** Callback for OlmMachine.registerRoomKeyUpdatedCallback
+     *
+     * Called by the rust-sdk whenever there is an update to (megolm) room keys. We
+     * check if we have any events waiting for the given keys, and schedule them for
+     * a decryption retry if so.
+     *
+     * @param keys - details of the updated keys
+     */
+    public async onRoomKeysUpdated(keys: RustSdkCryptoJs.RoomKeyInfo[]): Promise<void> {
+        for (const key of keys) {
+            this.onRoomKeyUpdated(key);
+        }
+    }
+
+    private onRoomKeyUpdated(key: RustSdkCryptoJs.RoomKeyInfo): void {
+        logger.debug(`Got update for session ${key.senderKey.toBase64()}|${key.sessionId} in ${key.roomId.toString()}`);
+        const pendingList = this.eventDecryptor.getEventsPendingRoomKey(key);
+        if (pendingList.length === 0) return;
+
+        logger.debug(
+            "Retrying decryption on events:",
+            pendingList.map((e) => `${e.getId()}`),
+        );
+
+        // Have another go at decrypting events with this key.
+        //
+        // We don't want to end up blocking the callback from Rust, which could otherwise end up dropping updates,
+        // so we don't wait for the decryption to complete. In any case, there is no need to wait:
+        // MatrixEvent.attemptDecryption ensures that there is only one decryption attempt happening at once,
+        // and deduplicates repeated attempts for the same event.
+        for (const ev of pendingList) {
+            ev.attemptDecryption(this, { isRetry: true }).catch((_e) => {
+                logger.info(`Still unable to decrypt event ${ev.getId()} after receiving key`);
+            });
+        }
+    }
+
     ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
     //
     // Outgoing requests
@@ -332,3 +356,104 @@ export class RustCrypto implements CryptoBackend {
         }
     }
 }
+
+class EventDecryptor {
+    /**
+     * Events which we couldn't decrypt due to unknown sessions / indexes.
+     *
+     * Map from senderKey to sessionId to Set of MatrixEvents
+     */
+    private eventsPendingKey = new MapWithDefault<string, MapWithDefault<string, Set<MatrixEvent>>>(
+        () => new MapWithDefault<string, Set<MatrixEvent>>(() => new Set()),
+    );
+
+    public constructor(private readonly olmMachine: RustSdkCryptoJs.OlmMachine) {}
+
+    public async attemptEventDecryption(event: MatrixEvent): Promise<IEventDecryptionResult> {
+        logger.info("Attempting decryption of event", event);
+        // add the event to the pending list *before* attempting to decrypt.
+        // then, if the key turns up while decryption is in progress (and
+        // decryption fails), we will schedule a retry.
+        // (fixes https://github.com/vector-im/element-web/issues/5001)
+        this.addEventToPendingList(event);
+
+        const res = (await this.olmMachine.decryptRoomEvent(
+            JSON.stringify({
+                event_id: event.getId(),
+                type: event.getWireType(),
+                sender: event.getSender(),
+                state_key: event.getStateKey(),
+                content: event.getWireContent(),
+                origin_server_ts: event.getTs(),
+            }),
+            new RustSdkCryptoJs.RoomId(event.getRoomId()!),
+        )) as RustSdkCryptoJs.DecryptedRoomEvent;
+
+        // Success. We can remove the event from the pending list, if
+        // that hasn't already happened.
+        this.removeEventFromPendingList(event);
+
+        return {
+            clearEvent: JSON.parse(res.event),
+            claimedEd25519Key: res.senderClaimedEd25519Key,
+            senderCurve25519Key: res.senderCurve25519Key,
+            forwardingCurve25519KeyChain: res.forwardingCurve25519KeyChain,
+        };
+    }
+
+    /**
+     * Look for events which are waiting for a given megolm session
+     *
+     * Returns a list of events which were encrypted by `session` and could not be decrypted
+     *
+     * @param session -
+     */
+    public getEventsPendingRoomKey(session: RustSdkCryptoJs.RoomKeyInfo): MatrixEvent[] {
+        const senderPendingEvents = this.eventsPendingKey.get(session.senderKey.toBase64());
+        if (!senderPendingEvents) return [];
+
+        const sessionPendingEvents = senderPendingEvents.get(session.sessionId);
+        if (!sessionPendingEvents) return [];
+
+        const roomId = session.roomId.toString();
+        return [...sessionPendingEvents].filter((ev) => ev.getRoomId() === roomId);
+    }
+
+    /**
+     * Add an event to the list of those awaiting their session keys.
+     */
+    private addEventToPendingList(event: MatrixEvent): void {
+        const content = event.getWireContent();
+        const senderKey = content.sender_key;
+        const sessionId = content.session_id;
+
+        const senderPendingEvents = this.eventsPendingKey.getOrCreate(senderKey);
+        const sessionPendingEvents = senderPendingEvents.getOrCreate(sessionId);
+        sessionPendingEvents.add(event);
+    }
+
+    /**
+     * Remove an event from the list of those awaiting their session keys.
+     */
+    private removeEventFromPendingList(event: MatrixEvent): void {
+        const content = event.getWireContent();
+        const senderKey = content.sender_key;
+        const sessionId = content.session_id;
+
+        const senderPendingEvents = this.eventsPendingKey.get(senderKey);
+        if (!senderPendingEvents) return;
+
+        const sessionPendingEvents = senderPendingEvents.get(sessionId);
+        if (!sessionPendingEvents) return;
+
+        sessionPendingEvents.delete(event);
+
+        // also clean up the higher-level maps if they are now empty
+        if (sessionPendingEvents.size === 0) {
+            senderPendingEvents.delete(sessionId);
+            if (senderPendingEvents.size === 0) {
+                this.eventsPendingKey.delete(senderKey);
+            }
+        }
+    }
+}

yarn.lock

@@ -1538,10 +1538,10 @@
   dependencies:
     lodash "^4.17.21"
 
-"@matrix-org/matrix-sdk-crypto-js@^0.1.0-alpha.5":
+"@matrix-org/matrix-sdk-crypto-js@^0.1.0-alpha.6":
-  version "0.1.0-alpha.5"
+  version "0.1.0-alpha.6"
-  resolved "https://registry.yarnpkg.com/@matrix-org/matrix-sdk-crypto-js/-/matrix-sdk-crypto-js-0.1.0-alpha.5.tgz#60ede2c43b9d808ba8cf46085a3b347b290d9658"
+  resolved "https://registry.yarnpkg.com/@matrix-org/matrix-sdk-crypto-js/-/matrix-sdk-crypto-js-0.1.0-alpha.6.tgz#c0bdb9ab0d30179b8ef744d1b4010b0ad0ab9c3a"
-  integrity sha512-2KjAgWNGfuGLNjJwsrs6gGX157vmcTfNrA4u249utgnMPbJl7QwuUqh1bGxQ0PpK06yvZjgPlkna0lTbuwtuQw==
+  integrity sha512-7hMffzw7KijxDyyH/eUyTfrLeCQHuyU3kaPOKGhcl3DZ3vx7bCncqjGMGTnxNPoP23I6gosvKSbO+3wYOT24Xg==
 
 "@matrix-org/olm@https://gitlab.matrix.org/api/v4/projects/27/packages/npm/@matrix-org/olm/-/@matrix-org/olm-3.2.14.tgz":
   version "3.2.14"