From 9c6d5a6c550efce66988fbc0d8c8cd7b3ba3879d Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>
Date: Tue, 20 Jun 2023 10:29:41 +0100
Subject: [PATCH] Element-R: wait for OlmMachine on startup (#3487)

* Element-R: wait for OlmMachine on startup

Previously, if you called `CryptoApi.getUserDeviceInfo()` before the first
`/sync` request happened, it would return an empty list, which made a bunch of
the tests racy. Add a hack to get the OlmMachine to think about its device
lists during startup.

* add a test
---
 spec/unit/rust-crypto/rust-crypto.spec.ts | 42 ++++++++++++++++++++++-
 src/rust-crypto/index.ts                  | 11 ++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/spec/unit/rust-crypto/rust-crypto.spec.ts b/spec/unit/rust-crypto/rust-crypto.spec.ts
index bfa16c62b73..c2826130ba9 100644
--- a/spec/unit/rust-crypto/rust-crypto.spec.ts
+++ b/spec/unit/rust-crypto/rust-crypto.spec.ts
@@ -19,10 +19,19 @@ import { IDBFactory } from "fake-indexeddb";
 import * as RustSdkCryptoJs from "@matrix-org/matrix-sdk-crypto-js";
 import { KeysQueryRequest, OlmMachine } from "@matrix-org/matrix-sdk-crypto-js";
 import { Mocked } from "jest-mock";
+import fetchMock from "fetch-mock-jest";
 
 import { RustCrypto } from "../../../src/rust-crypto/rust-crypto";
 import { initRustCrypto } from "../../../src/rust-crypto";
-import { IHttpOpts, IToDeviceEvent, MatrixClient, MatrixHttpApi } from "../../../src";
+import {
+    HttpApiEvent,
+    HttpApiEventHandlerMap,
+    IHttpOpts,
+    IToDeviceEvent,
+    MatrixClient,
+    MatrixHttpApi,
+    TypedEventEmitter,
+} from "../../../src";
 import { mkEvent } from "../../test-utils/test-utils";
 import { CryptoBackend } from "../../../src/common-crypto/CryptoBackend";
 import { IEventDecryptionResult } from "../../../src/@types/crypto";
@@ -421,6 +430,37 @@ describe("RustCrypto", () => {
             expect(recoveryKey.keyInfo?.passphrase?.iterations).toBe(500000);
         });
     });
+
+    it("should wait for a keys/query before returning devices", async () => {
+        jest.useFakeTimers();
+
+        const mockHttpApi = new MatrixHttpApi(new TypedEventEmitter<HttpApiEvent, HttpApiEventHandlerMap>(), {
+            baseUrl: "http://server/",
+            prefix: "",
+            onlyData: true,
+        });
+        fetchMock.post("path:/_matrix/client/v3/keys/upload", { one_time_key_counts: {} });
+        fetchMock.post("path:/_matrix/client/v3/keys/query", {
+            device_keys: {
+                [testData.TEST_USER_ID]: {
+                    [testData.TEST_DEVICE_ID]: testData.SIGNED_TEST_DEVICE_DATA,
+                },
+            },
+        });
+
+        const rustCrypto = await makeTestRustCrypto(mockHttpApi, testData.TEST_USER_ID);
+
+        // an attempt to fetch the device list should block
+        const devicesPromise = rustCrypto.getUserDeviceInfo([testData.TEST_USER_ID]);
+
+        // ... until a /sync completes, and we trigger the outgoingRequests.
+        rustCrypto.onSyncCompleted({});
+
+        const deviceMap = (await devicesPromise).get(testData.TEST_USER_ID)!;
+        expect(deviceMap.has(TEST_DEVICE_ID)).toBe(true);
+        expect(deviceMap.has(testData.TEST_DEVICE_ID)).toBe(true);
+        rustCrypto.stop();
+    });
 });
 
 /** build a basic RustCrypto instance for testing
diff --git a/src/rust-crypto/index.ts b/src/rust-crypto/index.ts
index 8f4aefdaaba..1fa8a2e93d3 100644
--- a/src/rust-crypto/index.ts
+++ b/src/rust-crypto/index.ts
@@ -57,6 +57,17 @@ export async function initRustCrypto(
         rustCrypto.onRoomKeysUpdated(sessions),
     );
 
+    // Tell the OlmMachine to think about its outgoing requests before we hand control back to the application.
+    //
+    // This is primarily a fudge to get it to correctly populate the `users_for_key_query` list, so that future
+    // calls to getIdentity (etc) block until the key queries are performed.
+    //
+    // Note that we don't actually need to *make* any requests here; it is sufficient to tell the Rust side to think
+    // about them.
+    //
+    // XXX: find a less hacky way to do this.
+    await olmMachine.outgoingRequests();
+
     logger.info("Completed rust crypto-sdk setup");
     return rustCrypto;
 }