Element-R: Automatically import keys from backup after gossip (following interactive verification)

spec/integ/crypto/verification.spec.ts

@@ -62,6 +62,7 @@
     BOB_TEST_USER_ID,
     CURVE25519_KEY_BACKUP_DATA,
     MASTER_CROSS_SIGNING_PUBLIC_KEY_BASE64,
+    MEGOLM_SESSION_DATA,
     SIGNED_CROSS_SIGNING_KEYS_DATA,
     SIGNED_TEST_DEVICE_DATA,
     TEST_DEVICE_ID,
@@ -80,7 +81,8 @@
     encryptSecretSend,
     ToDeviceEvent,
 } from "./olm-utils";
 import { KeyBackupInfo } from "../../../src/crypto-api";
+
 import { encodeBase64 } from "../../../src/base64";
 
 // The verification flows use javascript timers to set timeouts. We tell jest to use mock timer implementations
@@ -1259,19 +1261,41 @@
 
             const requestId = await requestPromises.get("m.megolm_backup.v1");
 
-            await sendBackupGossipAndExpectVersion(requestId!, BACKUP_DECRYPTION_KEY_BASE64, matchingBackupInfo);
+            const keyBackupIsCached = new Promise<void>((resolve) => {
+                aliceClient.on(CryptoEvent.KeyBackupPrivateKeyCached, () => {
+                    resolve();
+                });
+            });
 
-            // We are lacking a way to signal that the secret has been received, so we wait a bit..
-            jest.useRealTimers();
-            await new Promise((resolve) => {
-                setTimeout(resolve, 500);
+            const keyBackupStatusEnabled = new Promise<void>((resolve) => {
+                aliceClient.on(CryptoEvent.KeyBackupStatus, (enabled) => {
+                    if (enabled) {
+                        resolve();
+                    }
+                });
+            });
+
+            const autoImportFinished = new Promise<void>((resolve) => {
+                aliceClient.on(CryptoEvent.KeyBackupAutoImportFinished, (total) => {
+                    if (total > 0) {
+                        resolve();
+                    }
+                });
             });
-            jest.useFakeTimers();
+
+            await sendBackupGossipAndExpectVersion(requestId!, BACKUP_DECRYPTION_KEY_BASE64, matchingBackupInfo);
+
+            await keyBackupStatusEnabled;
+            await keyBackupIsCached;
 
             // the backup secret should be cached
             const cachedKey = await aliceClient.getCrypto()!.getSessionBackupPrivateKey();
             expect(cachedKey).toBeTruthy();
             expect(encodeBase64(cachedKey!)).toEqual(BACKUP_DECRYPTION_KEY_BASE64);
+
+            // An auto import should have been triggered
+            await autoImportFinished;
+
         });
 
         newBackendOnly("Should not accept the backup decryption key gossip if private key do not match", async () => {
@@ -1281,14 +1305,17 @@
 
             const requestId = await requestPromises.get("m.megolm_backup.v1");
 
+            const keyBackupStatusUpdate = new Promise<void>((resolve) => {
+                aliceClient.on(CryptoEvent.KeyBackupStatus, (enabled) => {
+                    resolve();
+                });
+            });
+
             await sendBackupGossipAndExpectVersion(requestId!, BACKUP_DECRYPTION_KEY_BASE64, nonMatchingBackupInfo);
 
+            await keyBackupStatusUpdate;
             // We are lacking a way to signal that the secret has been received, so we wait a bit..
-            jest.useRealTimers();
-            await new Promise((resolve) => {
-                setTimeout(resolve, 500);
-            });
-            jest.useFakeTimers();
+            await jest.runAllTimersAsync();
 
             // the backup secret should not be cached
             const cachedKey = await aliceClient.getCrypto()!.getSessionBackupPrivateKey();
@@ -1305,14 +1332,17 @@
             const infoCopy = Object.assign({}, matchingBackupInfo);
             delete infoCopy.auth_data.signatures;
 
+            const keyBackupStatusUpdate = new Promise<boolean>((resolve) => {
+                aliceClient.on(CryptoEvent.KeyBackupStatus, (enabled) => {
+                    resolve(enabled);
+                });
+            });
+
             await sendBackupGossipAndExpectVersion(requestId!, BACKUP_DECRYPTION_KEY_BASE64, infoCopy);
 
-            // We are lacking a way to signal that the secret has been received, so we wait a bit..
-            jest.useRealTimers();
-            await new Promise((resolve) => {
-                setTimeout(resolve, 500);
-            });
-            jest.useFakeTimers();
+            const backupStatus = await keyBackupStatusUpdate;
+            expect(backupStatus).toBe(false);
+            await jest.runAllTimersAsync();
 
             // the backup secret should not be cached
             const cachedKey = await aliceClient.getCrypto()!.getSessionBackupPrivateKey();
@@ -1326,18 +1356,23 @@
 
             const requestId = await requestPromises.get("m.megolm_backup.v1");
 
+
+            const keyBackupStatusUpdate = new Promise<boolean>((resolve) => {
+                aliceClient.on(CryptoEvent.KeyBackupStatus, (enabled) => {
+                    resolve(enabled);
+                });
+            });
+
+
             await sendBackupGossipAndExpectVersion(
                 requestId!,
                 BACKUP_DECRYPTION_KEY_BASE64,
                 unknownAlgorithmBackupInfo,
             );
 
-            // We are lacking a way to signal that the secret has been received, so we wait a bit..
-            jest.useRealTimers();
-            await new Promise((resolve) => {
-                setTimeout(resolve, 500);
-            });
-            jest.useFakeTimers();
+
+            await keyBackupStatusUpdate;
+            await jest.runAllTimersAsync();
 
             // the backup secret should not be cached
             const cachedKey = await aliceClient.getCrypto()!.getSessionBackupPrivateKey();
@@ -1351,14 +1386,17 @@
 
             const requestId = await requestPromises.get("m.megolm_backup.v1");
 
+            const keyBackupStatusUpdate = new Promise<boolean>((resolve) => {
+                aliceClient.on(CryptoEvent.KeyBackupStatus, (enabled) => {
+                    resolve(enabled);
+                });
+            });
+
             await sendBackupGossipAndExpectVersion(requestId!, "InvalidSecret", matchingBackupInfo);
 
-            // We are lacking a way to signal that the secret has been received, so we wait a bit..
-            jest.useRealTimers();
-            await new Promise((resolve) => {
-                setTimeout(resolve, 500);
-            });
-            jest.useFakeTimers();
+            const backupStatus = await keyBackupStatusUpdate;
+            expect(backupStatus).toBe(true);
+            jest.runAllTimersAsync();
 
             // the backup secret should not be cached
             const cachedKey = await aliceClient.getCrypto()!.getSessionBackupPrivateKey();
@@ -1400,8 +1438,19 @@
                 );
             });
 
-            fetchMock.get("express:/_matrix/client/v3/room_keys/keys", CURVE25519_KEY_BACKUP_DATA);
+            const fullBackup = {
+                rooms: {
+                    "!ROOM:ID": {
+                        sessions: {
+                            [MEGOLM_SESSION_DATA.session_id]: CURVE25519_KEY_BACKUP_DATA,
+                        },
+                    },
+                },
+            };
+
+            fetchMock.get("express:/_matrix/client/v3/room_keys/keys", fullBackup);
+
 
             // The dummy device sends the secret
             returnToDeviceMessageFromSync(toDeviceEvent);
 
@@ -1539,6 +1588,7 @@
             const messages = JSON.parse(opts.body as string).messages[TEST_USER_ID];
             // rust crypto broadcasts to all devices, old crypto to a specific device, take the first one
             const content = Object.values(messages)[0] as any;
+
             if (content.action == "request") {
                 const name = content.name;
                 const requestId = content.request_id;

src/client.ts

@@ -961,6 +961,9 @@ type CryptoEvents =
     | CryptoEvent.KeyBackupStatus
     | CryptoEvent.KeyBackupFailed
     | CryptoEvent.KeyBackupSessionsRemaining
+    | CryptoEvent.KeyBackupPrivateKeyCached
+    | CryptoEvent.KeyBackupAutoImportStarted
+    | CryptoEvent.KeyBackupAutoImportFinished
     | CryptoEvent.RoomKeyRequest
     | CryptoEvent.RoomKeyRequestCancellation
     | CryptoEvent.VerificationRequest
@@ -2269,6 +2272,7 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
             CryptoEvent.DeviceVerificationChanged,
             CryptoEvent.UserTrustStatusChanged,
             CryptoEvent.KeysChanged,
+            CryptoEvent.KeyBackupPrivateKeyCached,
         ]);
 
         this.logger.debug("Crypto: initialising crypto object...");
@@ -2350,13 +2354,25 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
             rustCrypto.onLiveEventFromSync(event);
         });
 
+        rustCrypto.on(CryptoEvent.KeyBackupPrivateKeyCached, (info) => {
+            // do not await this, as it can be long.
+            // XXX We need new APIs to do this properly. To avoid having several restore going on at the same time,
+            // there should be proper progress reporting, and a way to cancel a restore.
+            this.emit(CryptoEvent.KeyBackupAutoImportStarted);
+            this.restoreKeyBackupWithCache(undefined, undefined, info).then((result) => {
+                this.emit(CryptoEvent.KeyBackupAutoImportFinished, result.imported);
+                logger.info("Backup restored.");
+            });
+        });
+
         // re-emit the events emitted by the crypto impl
         this.reEmitter.reEmit(rustCrypto, [
             CryptoEvent.VerificationRequestReceived,
             CryptoEvent.UserTrustStatusChanged,
             CryptoEvent.KeyBackupStatus,
             CryptoEvent.KeyBackupSessionsRemaining,
             CryptoEvent.KeyBackupFailed,
+            CryptoEvent.KeyBackupPrivateKeyCached,
         ]);
     }
 

src/crypto/index.ts

@@ -231,6 +231,34 @@ export enum CryptoEvent {
     KeyBackupStatus = "crypto.keyBackupStatus",
     KeyBackupFailed = "crypto.keyBackupFailed",
     KeyBackupSessionsRemaining = "crypto.keyBackupSessionsRemaining",
+    /**
+     * Fires when a new valid backup decryption key is in cache.
+     *
+     * The payload is the `KeyBackupInfo`.
+     *
+     * This event is only fired by the rust crypto backend.
+     */
+    KeyBackupPrivateKeyCached = "crypto.KeyBackupPrivateKeyCached",
+
+    /**
+     * Fires when an automatic key import is started.
+     *
+     * This will happen after a sucessful verification when the private backup key
+     * is stored in cache.
+     *
+     * This event is only fired by the rust crypto backend.
+     */
+    KeyBackupAutoImportStarted = "crypto.KeyBackupAutoImportStarted",
+
+    /**
+     * Fires when an automatic key import is finished.
+     *
+     * This will happen after a sucessful verification when the private backup key
+     * is stored in cache.
+     *
+     * This event is only fired by the rust crypto backend.
+     */
+    KeyBackupAutoImportFinished = "crypto.KeyBackupAutoImportFinished",
     KeySignatureUploadFailure = "crypto.keySignatureUploadFailure",
     /** @deprecated Use `VerificationRequestReceived`. */
     VerificationRequest = "crypto.verification.request",
@@ -348,6 +376,9 @@ export type CryptoEventHandlerMap = {
      */
     [CryptoEvent.DevicesUpdated]: (users: string[], initialFetch: boolean) => void;
     [CryptoEvent.UserCrossSigningUpdated]: (userId: string) => void;
+    [CryptoEvent.KeyBackupPrivateKeyCached]: (info: KeyBackupInfo) => void;
+    [CryptoEvent.KeyBackupAutoImportStarted]: () => void;
+    [CryptoEvent.KeyBackupAutoImportFinished]: (totalKeys: number) => void;
 };
 
 export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap> implements CryptoBackend {

src/rust-crypto/backup.ts

@@ -156,6 +156,9 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
             );
 
             await this.olmMachine.saveBackupDecryptionKey(backupDecryptionKey, backupCheck.backupInfo.version);
+            // Emit an event that we have a new backup decryption key, so that client can automatically
+            // start importing all keys from the backup.
+            this.emit(CryptoEvent.KeyBackupPrivateKeyCached, backupCheck.backupInfo);
             return true;
         } catch (e) {
             logger.warn("handleBackupSecretReceived: Invalid backup decryption key", e);
@@ -202,6 +205,8 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
                 await this.disableKeyBackup();
             } else {
                 logger.log("Key backup present on server but not trusted: not enabling key backup");
+                // No backup was actually enabled, but still emit the KeyBackupStatus event.
+                this.emit(CryptoEvent.KeyBackupStatus, false);
             }
         } else {
             if (activeVersion === null) {
@@ -486,10 +491,15 @@ export class RustBackupDecryptor implements BackupDecryptor {
 export type RustBackupCryptoEvents =
     | CryptoEvent.KeyBackupStatus
     | CryptoEvent.KeyBackupSessionsRemaining
-    | CryptoEvent.KeyBackupFailed;
+    | CryptoEvent.KeyBackupFailed
+    | CryptoEvent.KeyBackupPrivateKeyCached;
 
 export type RustBackupCryptoEventMap = {
     [CryptoEvent.KeyBackupStatus]: (enabled: boolean) => void;
     [CryptoEvent.KeyBackupSessionsRemaining]: (remaining: number) => void;
     [CryptoEvent.KeyBackupFailed]: (errCode: string) => void;
+    /**
+     * Fired when the backup decryption key is received via secret sharing and stored in cache.
+     */
+    [CryptoEvent.KeyBackupPrivateKeyCached]: (info: KeyBackupInfo) => void;
 };

src/rust-crypto/rust-crypto.ts

@@ -152,6 +152,7 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, RustCryptoEv
             CryptoEvent.KeyBackupStatus,
             CryptoEvent.KeyBackupSessionsRemaining,
             CryptoEvent.KeyBackupFailed,
+            CryptoEvent.KeyBackupPrivateKeyCached,
         ]);
 
         this.crossSigningIdentity = new CrossSigningIdentity(olmMachine, this.outgoingRequestProcessor, secretStorage);
@@ -1779,7 +1780,8 @@ function rustEncryptionInfoToJsEncryptionInfo(
 type RustCryptoEvents =
     | CryptoEvent.VerificationRequestReceived
     | CryptoEvent.UserTrustStatusChanged
-    | RustBackupCryptoEvents;
+    | RustBackupCryptoEvents
+    | CryptoEvent.KeyBackupPrivateKeyCached;
 
 type RustCryptoEventMap = {
     /**
@@ -1791,4 +1793,9 @@ type RustCryptoEventMap = {
      * Fires when the trust status of a user changes.
      */
     [CryptoEvent.UserTrustStatusChanged]: (userId: string, userTrustLevel: UserVerificationStatus) => void;
+
+    /**
+     * Fires when the backup decryption key is received via secret sharing.
+     */
+    [CryptoEvent.KeyBackupPrivateKeyCached]: (info: KeyBackupInfo) => void;
 } & RustBackupCryptoEventMap;