Import room key bundles received after invite.

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 decrypts the megolm message, but should not be able to.
+        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 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

@@ -2010,6 +2010,7 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
         this.on(ClientEvent.Event, (event) => {
             rustCrypto.onLiveEventFromSync(event);
         });
+        this.on(ClientEvent.ReceivedToDeviceMessage, (payload) => rustCrypto.onReceiveToDeviceMessage(payload));
 
         // re-emit the events emitted by the crypto impl
         this.reEmitter.reEmit(rustCrypto, [
@@ -2408,7 +2409,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 and inviter specified, it will automatically import it.
+     *
+     * @param roomId - The room we were invited to, for which we did not receive a key bundle before accepting the invite.
+     *
+     * @param inviter - The user who invited us to the room and is expected to send us the room key bundle.
      */
-    maybeAcceptKeyBundle(roomId: string, inviter: string): Promise<void>;
+    markRoomAsPendingKeyBundle(roomId: string, inviter: string): void;
 }
 
 /** The methods which crypto implementations should expose to the Sync api

src/rust-crypto/rust-crypto.ts

@@ -129,6 +129,9 @@
     /** mapping of roomId → encryptor class */
     private roomEncryptors: Record<string, RoomEncryptor> = {};
 
+    /** mapping of room ID -> inviter ID for rooms pending MSC4268 key bundles */
+    private roomsPendingKeyBundles: Record<string, string> = {};
+
     private eventDecryptor: EventDecryptor;
     private keyClaimManager: KeyClaimManager;
     private outgoingRequestProcessor: OutgoingRequestProcessor;
@@ -329,7 +332,7 @@
     /**
      * 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)
@@ -352,7 +355,7 @@
         );
         if (!bundleData) {
             logger.info("No key bundle found for user");
-            return;
+            return false;
         }
 
         logger.info(`Fetching key bundle ${bundleData.url}`);
@@ -391,7 +394,17 @@
             logger.warn(`Error receiving encrypted bundle:`, err);
             throw err;
         }
+
+        return true;
+    }
+
+    /**
+     * Implementation of {@link CryptoBackend.markRoomAsPendingKeyBundle}.
+     */
+    public async markRoomAsPendingKeyBundle(roomId: string, inviterId: string): Promise<void> {
+        this.roomsPendingKeyBundles[roomId] = inviterId;
     }
+
     ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
     //
     // CryptoApi implementation
@@ -2138,6 +2151,34 @@
     public async getOwnIdentity(): Promise<RustSdkCryptoJs.OwnUserIdentity | undefined> {
         return await this.olmMachine.getIdentity(new RustSdkCryptoJs.UserId(this.userId));
     }
+    /**
+     * Handles the receipt of a to-device message, specifically for processing
+     * "io.element.msc4268.room_key_bundle" message types.
+     *
+     * @param payload - The received to-device message payload, which includes
+     * the message content and optional encryption information.
+     */
+    public async onReceiveToDeviceMessage(payload: ReceivedToDeviceMessage): Promise<void> {
+        if (payload.message.type != "io.element.msc4268.room_key_bundle") {
+            return;
+        }
+
+        const { message, encryptionInfo } = payload;
+        const claimedSender = encryptionInfo?.sender ?? message.sender;
+
+        // Validate room ID
+        const roomId = message.content.room_id;
+        if (typeof roomId !== "string") {
+            return;
+        }
+
+        // Check if the room is in the map of rooms we expect to receive bundles from, otherwise discard.
+        if (this.roomsPendingKeyBundles[roomId] !== claimedSender) {
+            return;
+        }
+
+        await this.maybeAcceptKeyBundle(roomId, claimedSender);
+    }
 }
 
 class EventDecryptor {