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.

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<void>;
+    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.
+     */
+    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>;