Import room key bundles received after invite. (#5080)

* feat: Import room key bundles when received after invite.

* tests: Add spec test for room key bundle arriving after invite accepted.

* chore: Fix code quality issue (unnecessary async function).

* docs: Tidy up comments.

* refactor: Simplify key bundle importing after invite to one entrypoint.

- Remove `onReceiveToDeviceEvent` from `CryptoBackend`.
- Copy old room key bundle importing logic to
  `preprocessToDeviceEvents`.

* refactor: Move late bundle importing to main preprocess loop.

* fix: Use `Map` over `Record` to prevent prototype pollution.

Author: kaylendog

spec/integ/crypto/history-sharing.spec.ts

@@ -30,7 +30,7 @@ import {
 import { E2EKeyReceiver } from "../../test-utils/E2EKeyReceiver.ts";
 import { SyncResponder } from "../../test-utils/SyncResponder.ts";
 import { mockInitialApiRequests, mockSetupCrossSigningRequests } from "../../test-utils/mockEndpoints.ts";
-import { getSyncResponse, mkEventCustom, syncPromise } from "../../test-utils/test-utils.ts";
+import { getSyncResponse, mkEventCustom, syncPromise, waitFor } from "../../test-utils/test-utils.ts";
 import { E2EKeyResponder } from "../../test-utils/E2EKeyResponder.ts";
 import { flushPromises } from "../../test-utils/flushPromises.ts";
 import { E2EOTKClaimResponder } from "../../test-utils/E2EOTKClaimResponder.ts";
@@ -80,6 +80,9 @@ describe("History Sharing", () => {
     let bobSyncResponder: SyncResponder;
 
     beforeEach(async () => {
+        // Reset mocks.
+        fetchMock.reset();
+
         // anything that we don't have a specific matcher for silently returns a 404
         fetchMock.catch(404);
         fetchMock.config.warnOnFallback = false;
@@ -201,6 +204,104 @@ describe("History Sharing", () => {
         expect(event.getContent().body).toEqual("Hi!");
     });
 
+    test("Room keys are imported correctly if invite is accepted before the bundle arrives", async () => {
+        // Alice is in an encrypted room
+        const syncResponse = getSyncResponse([aliceClient.getSafeUserId()], ROOM_ID);
+        aliceSyncResponder.sendOrQueueSyncResponse(syncResponse);
+        await syncPromise(aliceClient);
+
+        // ... and she sends an event
+        const msgProm = expectSendRoomEvent(ALICE_HOMESERVER_URL, "m.room.encrypted");
+        await aliceClient.sendEvent(ROOM_ID, EventType.RoomMessage, { msgtype: MsgType.Text, body: "Hello!" });
+        const sentMessage = await msgProm;
+        debug(`Alice sent encrypted room event: ${JSON.stringify(sentMessage)}`);
+
+        // Now, Alice invites Bob
+        const uploadProm = new Promise<Uint8Array>((resolve) => {
+            fetchMock.postOnce(new URL("/_matrix/media/v3/upload", ALICE_HOMESERVER_URL).toString(), (url, request) => {
+                const body = request.body as Uint8Array;
+                debug(`Alice uploaded blob of length ${body.length}`);
+                resolve(body);
+                return { content_uri: "mxc://alice-server/here" };
+            });
+        });
+        const toDeviceMessageProm = expectSendToDeviceMessage(ALICE_HOMESERVER_URL, "m.room.encrypted");
+        // POST https://alice-server.com/_matrix/client/v3/rooms/!room%3Aexample.com/invite
+        fetchMock.postOnce(`${ALICE_HOMESERVER_URL}/_matrix/client/v3/rooms/${encodeURIComponent(ROOM_ID)}/invite`, {});
+        await aliceClient.invite(ROOM_ID, bobClient.getSafeUserId(), { shareEncryptedHistory: true });
+        const uploadedBlob = await uploadProm;
+        const sentToDeviceRequest = await toDeviceMessageProm;
+        debug(`Alice sent encrypted to-device events: ${JSON.stringify(sentToDeviceRequest)}`);
+        const bobToDeviceMessage = sentToDeviceRequest[bobClient.getSafeUserId()][bobClient.deviceId!];
+        expect(bobToDeviceMessage).toBeDefined();
+
+        // Bob receives the room invite, but not the room key bundle
+        const inviteEvent = mkEventCustom({
+            type: "m.room.member",
+            sender: aliceClient.getSafeUserId(),
+            state_key: bobClient.getSafeUserId(),
+            content: { membership: KnownMembership.Invite },
+        });
+        bobSyncResponder.sendOrQueueSyncResponse({
+            rooms: { invite: { [ROOM_ID]: { invite_state: { events: [inviteEvent] } } } },
+        });
+        await syncPromise(bobClient);
+
+        const room = bobClient.getRoom(ROOM_ID);
+        expect(room).toBeTruthy();
+        expect(room?.getMyMembership()).toEqual(KnownMembership.Invite);
+
+        fetchMock.postOnce(`${BOB_HOMESERVER_URL}/_matrix/client/v3/join/${encodeURIComponent(ROOM_ID)}`, {
+            room_id: ROOM_ID,
+        });
+        await bobClient.joinRoom(ROOM_ID, { acceptSharedHistory: true });
+
+        // Bob receives and attempts to decrypt the megolm message, but should not be able to (yet).
+        const bobSyncResponse = getSyncResponse([aliceClient.getSafeUserId(), bobClient.getSafeUserId()], ROOM_ID);
+        bobSyncResponse.rooms.join[ROOM_ID].timeline.events.push(
+            mkEventCustom({
+                type: "m.room.encrypted",
+                sender: aliceClient.getSafeUserId(),
+                content: sentMessage,
+                event_id: "$event_id",
+            }) as any,
+        );
+        bobSyncResponder.sendOrQueueSyncResponse(bobSyncResponse);
+        await syncPromise(bobClient);
+        const bobRoom = bobClient.getRoom(ROOM_ID);
+        const event = bobRoom!.getLastLiveEvent()!;
+        expect(event.getId()).toEqual("$event_id");
+        await event.getDecryptionPromise();
+        expect(event.isDecryptionFailure()).toBeTruthy();
+
+        // Now the room key bundle message arrives
+        fetchMock.getOnce(
+            `begin:${BOB_HOMESERVER_URL}/_matrix/client/v1/media/download/alice-server/here`,
+            { body: uploadedBlob },
+            { sendAsJson: false },
+        );
+        bobSyncResponder.sendOrQueueSyncResponse({
+            to_device: {
+                events: [
+                    {
+                        type: "m.room.encrypted",
+                        sender: aliceClient.getSafeUserId(),
+                        content: bobToDeviceMessage,
+                    },
+                ],
+            },
+        });
+        await syncPromise(bobClient);
+
+        // Once the room key bundle finishes downloading, we should be able to decrypt the message.
+        await waitFor(async () => {
+            await event.getDecryptionPromise();
+            expect(event.isDecryptionFailure()).toBeFalsy();
+            expect(event.getType()).toEqual("m.room.message");
+            expect(event.getContent().body).toEqual("Hello!");
+        });
+    });
+
     afterEach(async () => {
         bobClient.stopClient();
         aliceClient.stopClient();

src/client.ts

@@ -2406,7 +2406,12 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
 
         const roomId = res.room_id;
         if (opts.acceptSharedHistory && inviter && this.cryptoBackend) {
-            await this.cryptoBackend.maybeAcceptKeyBundle(roomId, inviter);
+            // Try to accept the room key bundle specified in a `m.room_key_bundle` to-device message we (might have) already received.
+            const bundleDownloaded = await this.cryptoBackend.maybeAcceptKeyBundle(roomId, inviter);
+            // If this fails, i.e. we haven't received this message yet, we need to wait until the to-device message arrives.
+            if (!bundleDownloaded) {
+                this.cryptoBackend.markRoomAsPendingKeyBundle(roomId, inviter);
+            }
         }
 
         // In case we were originally given an alias, check the room cache again

src/common-crypto/CryptoBackend.ts

@@ -90,8 +90,20 @@ export interface CryptoBackend extends SyncCryptoCallbacks, CryptoApi {
      *
      * @param inviter - The user who invited us to the room and is expected to have
      *   sent the room key bundle.
+     *
+     * @returns `true` if the key bundle was successfuly downloaded and imported.
+     */
+    maybeAcceptKeyBundle(roomId: string, inviter: string): Promise<boolean>;
+
+    /**
+     * Mark a room as pending a key bundle under MSC4268. The backend will listen for room key bundle messages, and if
+     * it sees one matching the room specified, it will automatically import it as long as the message author's ID matches
+     * the inviter's ID.
+     *
+     * @param roomId - The room we were invited to, for which we did not receive a key bundle before accepting the invite.
+     * @param inviterId - The user who invited us to the room and is expected to send the room key bundle.
      */
-    maybeAcceptKeyBundle(roomId: string, inviter: string): Promise<void>;
+    markRoomAsPendingKeyBundle(roomId: string, inviterId: string): void;
 }
 
 /** The methods which crypto implementations should expose to the Sync api

src/rust-crypto/rust-crypto.ts

@@ -129,6 +129,9 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
     /** mapping of roomId → encryptor class */
     private roomEncryptors: Record<string, RoomEncryptor> = {};
 
+    /** mapping of room ID -> inviter ID for rooms pending MSC4268 key bundles */
+    private readonly roomsPendingKeyBundles: Map<string, string> = new Map();
+
     private eventDecryptor: EventDecryptor;
     private keyClaimManager: KeyClaimManager;
     private outgoingRequestProcessor: OutgoingRequestProcessor;
@@ -329,10 +332,9 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
     /**
      * Implementation of {@link CryptoBackend.maybeAcceptKeyBundle}.
      */
-    public async maybeAcceptKeyBundle(roomId: string, inviter: string): Promise<void> {
+    public async maybeAcceptKeyBundle(roomId: string, inviter: string): Promise<boolean> {
         // TODO: retry this if it gets interrupted or it fails. (https://github.com/matrix-org/matrix-rust-sdk/issues/5112)
         // TODO: do this in the background.
-        // TODO: handle the bundle message arriving after the invite (https://github.com/element-hq/element-web/issues/30740)
 
         const logger = new LogSpan(this.logger, `maybeAcceptKeyBundle(${roomId}, ${inviter})`);
 
@@ -352,7 +354,7 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
         );
         if (!bundleData) {
             logger.info("No key bundle found for user");
-            return;
+            return false;
         }
 
         logger.info(`Fetching key bundle ${bundleData.url}`);
@@ -391,7 +393,17 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
             logger.warn(`Error receiving encrypted bundle:`, err);
             throw err;
         }
+
+        return true;
+    }
+
+    /**
+     * Implementation of {@link CryptoBackend.markRoomAsPendingKeyBundle}.
+     */
+    public markRoomAsPendingKeyBundle(roomId: string, inviter: string): void {
+        this.roomsPendingKeyBundles.set(roomId, inviter);
     }
+
     ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
     //
     // CryptoApi implementation
@@ -1703,6 +1715,37 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
                             senderVerified: encryptionInfo.isSenderVerified(),
                         },
                     });
+
+                    // If we have received a room key bundle message, and have previously marked the room
+                    // IDs it references as pending key bundles, tell the Rust SDK to try and accept it,
+                    // just in case it was received after invite.
+                    //
+                    // We don't actually need to validate the contents of the bundle message, or do
+                    // anything with its contents at all. We simply want to inform the Rust SDK we have
+                    // received a new room key bundle that we might be able to download.
+                    if (
+                        isRoomKeyBundleMessage(parsedMessage) &&
+                        this.roomsPendingKeyBundles.has(parsedMessage.content.room_id)
+                    ) {
+                        // No `await`-ing here, as this is called from inside the `/sync` loop.
+                        this.maybeAcceptKeyBundle(
+                            parsedMessage.content.room_id,
+                            this.roomsPendingKeyBundles.get(parsedMessage.content.room_id)!,
+                        ).then(
+                            (success) => {
+                                if (success) {
+                                    this.roomsPendingKeyBundles.delete(parsedMessage.content.room_id);
+                                }
+                            },
+                            (err) => {
+                                this.logger.error(
+                                    `Error attempting to download key bundle for room ${parsedMessage.content.room_id}`,
+                                );
+                                this.logger.error(err);
+                            },
+                        );
+                    }
+
                     break;
                 }
                 case RustSdkCryptoJs.ProcessedToDeviceEventType.PlainText: {
@@ -2468,5 +2511,25 @@ function rustEncryptionInfoToJsEncryptionInfo(
     return { shieldColour, shieldReason };
 }
 
+interface RoomKeyBundleMessage {
+    type: "io.element.msc4268.room_key_bundle";
+    content: {
+        room_id: string;
+    };
+}
+
+/**
+ * Determines if the given payload is a RoomKeyBundleMessage.
+ *
+ * A RoomKeyBundleMessage is identified by having a specific message type
+ * ("io.element.msc4268.room_key_bundle") and a valid room_id in its content.
+ *
+ * @param message - The received to-device message to check.
+ * @returns True if the payload matches the RoomKeyBundleMessage structure, false otherwise.
+ */
+function isRoomKeyBundleMessage(message: IToDeviceEvent): message is IToDeviceEvent & RoomKeyBundleMessage {
+    return message.type === "io.element.msc4268.room_key_bundle" && typeof message.content.room_id === "string";
+}
+
 type CryptoEvents = (typeof CryptoEvent)[keyof typeof CryptoEvent];
 type RustCryptoEvents = Exclude<CryptoEvents, CryptoEvent.LegacyCryptoStoreMigrationProgress>;