Element-R: implement encryption of outgoing events (#3122)

This PR wires up the Rust-SDK into the event encryption path

package.json

@@ -109,6 +109,7 @@
         "eslint-plugin-unicorn": "^45.0.0",
         "exorcist": "^2.0.0",
         "fake-indexeddb": "^4.0.0",
+        "fetch-mock-jest": "^1.5.1",
         "jest": "^29.0.0",
         "jest-environment-jsdom": "^29.0.0",
         "jest-localstorage-mock": "^2.4.6",

spec/integ/crypto.spec.ts

@@ -633,6 +633,35 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("megolm (%s)", (backend: string,
         expect(event.getContent().body).toEqual("42");
     });
 
+    oldBackendOnly("prepareToEncrypt", async () => {
+        aliceTestClient.expectKeyQuery({ device_keys: { "@alice:localhost": {} }, failures: {} });
+        await aliceTestClient.start();
+        aliceTestClient.client.setGlobalErrorOnUnknownDevices(false);
+
+        // tell alice she is sharing a room with bob
+        aliceTestClient.httpBackend.when("GET", "/sync").respond(200, getSyncResponse(["@bob:xyz"]));
+        await aliceTestClient.flushSync();
+
+        // we expect alice first to query bob's keys...
+        aliceTestClient.httpBackend.when("POST", "/keys/query").respond(200, getTestKeysQueryResponse("@bob:xyz"));
+        aliceTestClient.httpBackend.flush("/keys/query", 1);
+
+        // ... and then claim one of his OTKs
+        aliceTestClient.httpBackend.when("POST", "/keys/claim").respond(200, getTestKeysClaimResponse("@bob:xyz"));
+        aliceTestClient.httpBackend.flush("/keys/claim", 1);
+
+        // fire off the prepare request
+        const room = aliceTestClient.client.getRoom(ROOM_ID);
+        expect(room).toBeTruthy();
+        const p = aliceTestClient.client.prepareToEncrypt(room!);
+
+        // we expect to get a room key message
+        await expectSendRoomKey(aliceTestClient.httpBackend, "@bob:xyz", testOlmAccount);
+
+        // the prepare request should complete successfully.
+        await p;
+    });
+
     oldBackendOnly("Alice sends a megolm message", async () => {
         aliceTestClient.expectKeyQuery({ device_keys: { "@alice:localhost": {} }, failures: {} });
         await aliceTestClient.start();

spec/unit/matrix-client.spec.ts

@@ -1415,7 +1415,7 @@ describe("MatrixClient", function () {
                 expect(getRoomId).toEqual(roomId);
                 return mockRoom;
             };
-            client.crypto = {
+            client.crypto = client["cryptoBackend"] = {
                 // mock crypto
                 encryptEvent: () => new Promise(() => {}),
                 stop: jest.fn(),
@@ -1437,8 +1437,9 @@ describe("MatrixClient", function () {
 
         it("should cancel an event which is encrypting", async () => {
             // @ts-ignore protected method access
-            client.encryptAndSendEvent(null, event);
+            client.encryptAndSendEvent(mockRoom, event);
             await testUtils.emitPromise(event, "Event.status");
+            expect(event.status).toBe(EventStatus.ENCRYPTING);
             client.cancelPendingEvent(event);
             assertCancelled();
         });

spec/unit/rust-crypto/KeyClaimManager.spec.ts

@@ -0,0 +1,158 @@
+/*
+Copyright 2023 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import * as RustSdkCryptoJs from "@matrix-org/matrix-sdk-crypto-js";
+import fetchMock from "fetch-mock-jest";
+import { Mocked } from "jest-mock";
+import { KeysClaimRequest, UserId } from "@matrix-org/matrix-sdk-crypto-js";
+
+import { OutgoingRequestProcessor } from "../../../src/rust-crypto/OutgoingRequestProcessor";
+import { KeyClaimManager } from "../../../src/rust-crypto/KeyClaimManager";
+import { TypedEventEmitter } from "../../../src/models/typed-event-emitter";
+import { HttpApiEvent, HttpApiEventHandlerMap, MatrixHttpApi } from "../../../src";
+
+afterEach(() => {
+    fetchMock.mockReset();
+});
+
+describe("KeyClaimManager", () => {
+    /* for these tests, we connect a KeyClaimManager to a mock OlmMachine, and a real OutgoingRequestProcessor
+     * (which is connected to a mock fetch implementation)
+     */
+
+    /** the KeyClaimManager implementation under test */
+    let keyClaimManager: KeyClaimManager;
+
+    /** a mocked-up OlmMachine which the OutgoingRequestProcessor and KeyClaimManager are connected to */
+    let olmMachine: Mocked<RustSdkCryptoJs.OlmMachine>;
+
+    beforeEach(async () => {
+        const dummyEventEmitter = new TypedEventEmitter<HttpApiEvent, HttpApiEventHandlerMap>();
+        const httpApi = new MatrixHttpApi(dummyEventEmitter, {
+            baseUrl: "https://example.com",
+            prefix: "/_matrix",
+            onlyData: true,
+        });
+
+        olmMachine = {
+            getMissingSessions: jest.fn(),
+            markRequestAsSent: jest.fn(),
+        } as unknown as Mocked<RustSdkCryptoJs.OlmMachine>;
+
+        const outgoingRequestProcessor = new OutgoingRequestProcessor(olmMachine, httpApi);
+
+        keyClaimManager = new KeyClaimManager(olmMachine, outgoingRequestProcessor);
+    });
+
+    /**
+     * Returns a promise which resolve once olmMachine.markRequestAsSent is called.
+     *
+     * The call itself will block initially.
+     *
+     * The promise returned by this function yields a callback function, which should be called to unblock the
+     * markRequestAsSent call.
+     */
+    function awaitCallToMarkRequestAsSent(): Promise<() => void> {
+        return new Promise<() => void>((resolveCalledPromise, _reject) => {
+            olmMachine.markRequestAsSent.mockImplementationOnce(async () => {
+                // the mock implementation returns a promise...
+                const completePromise = new Promise<void>((resolveCompletePromise, _reject) => {
+                    // ...  and we now resolve the original promise with the resolver for that second promise.
+                    resolveCalledPromise(resolveCompletePromise);
+                });
+                return completePromise;
+            });
+        });
+    }
+
+    it("should claim missing keys", async () => {
+        const u1 = new UserId("@alice:example.com");
+        const u2 = new UserId("@bob:example.com");
+
+        // stub out olmMachine.getMissingSessions(), with a result indicating that it needs a keyclaim
+        const keysClaimRequest = new KeysClaimRequest("1234", '{ "k1": "v1" }');
+        olmMachine.getMissingSessions.mockResolvedValueOnce(keysClaimRequest);
+
+        // have the claim request return a 200
+        fetchMock.postOnce("https://example.com/_matrix/client/v3/keys/claim", '{ "k": "v" }');
+
+        // also stub out olmMachine.markRequestAsSent
+        olmMachine.markRequestAsSent.mockResolvedValueOnce(undefined);
+
+        // fire off the request
+        await keyClaimManager.ensureSessionsForUsers([u1, u2]);
+
+        // check that all the calls were made
+        expect(olmMachine.getMissingSessions).toHaveBeenCalledWith([u1, u2]);
+        expect(fetchMock).toHaveFetched("https://example.com/_matrix/client/v3/keys/claim", {
+            method: "POST",
+            body: { k1: "v1" },
+        });
+        expect(olmMachine.markRequestAsSent).toHaveBeenCalledWith("1234", keysClaimRequest.type, '{ "k": "v" }');
+    });
+
+    it("should wait for previous claims to complete before making another", async () => {
+        const u1 = new UserId("@alice:example.com");
+        const u2 = new UserId("@bob:example.com");
+
+        // stub out olmMachine.getMissingSessions(), with a result indicating that it needs a keyclaim
+        const keysClaimRequest = new KeysClaimRequest("1234", '{ "k1": "v1" }');
+        olmMachine.getMissingSessions.mockResolvedValue(keysClaimRequest);
+
+        // have the claim request return a 200
+        fetchMock.post("https://example.com/_matrix/client/v3/keys/claim", '{ "k": "v" }');
+
+        // stub out olmMachine.markRequestAsSent, and have it block
+        let markRequestAsSentPromise = awaitCallToMarkRequestAsSent();
+
+        // fire off two requests, and keep track of whether their promises resolve
+        let req1Resolved = false;
+        keyClaimManager.ensureSessionsForUsers([u1]).then(() => {
+            req1Resolved = true;
+        });
+        let req2Resolved = false;
+        const req2 = keyClaimManager.ensureSessionsForUsers([u2]).then(() => {
+            req2Resolved = true;
+        });
+
+        // now: wait for the (first) call to OlmMachine.markRequestAsSent
+        let resolveMarkRequestAsSentCallback = await markRequestAsSentPromise;
+
+        // at this point, there should have been a single call to getMissingSessions, and a single fetch; and neither
+        // call to ensureSessionsAsUsers should have completed
+        expect(olmMachine.getMissingSessions).toHaveBeenCalledWith([u1]);
+        expect(olmMachine.getMissingSessions).toHaveBeenCalledTimes(1);
+        expect(fetchMock).toHaveBeenCalledTimes(1);
+        expect(req1Resolved).toBe(false);
+        expect(req2Resolved).toBe(false);
+
+        // await the next call to markRequestAsSent, and release the first one
+        markRequestAsSentPromise = awaitCallToMarkRequestAsSent();
+        resolveMarkRequestAsSentCallback();
+        resolveMarkRequestAsSentCallback = await markRequestAsSentPromise;
+
+        // the first request should now have completed, and we should have more calls and fetches
+        expect(olmMachine.getMissingSessions).toHaveBeenCalledWith([u2]);
+        expect(olmMachine.getMissingSessions).toHaveBeenCalledTimes(2);
+        expect(fetchMock).toHaveBeenCalledTimes(2);
+        expect(req1Resolved).toBe(true);
+        expect(req2Resolved).toBe(false);
+
+        // finally, release the second call to markRequestAsSent and check that the second request completes
+        resolveMarkRequestAsSentCallback();
+        await req2;
+    });
+});

src/client.ts

@@ -2185,7 +2185,11 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
         // importing rust-crypto will download the webassembly, so we delay it until we know it will be
         // needed.
         const RustCrypto = await import("./rust-crypto");
-        this.cryptoBackend = await RustCrypto.initRustCrypto(this.http, userId, deviceId);
+        const rustCrypto = await RustCrypto.initRustCrypto(this.http, userId, deviceId);
+        this.cryptoBackend = rustCrypto;
+
+        // attach the event listeners needed by RustCrypto
+        this.on(RoomMemberEvent.Membership, rustCrypto.onRoomMembership.bind(rustCrypto));
     }
 
     /**
@@ -2608,10 +2612,10 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
      * @param room - the room the event is in
      */
     public prepareToEncrypt(room: Room): void {
-        if (!this.crypto) {
+        if (!this.cryptoBackend) {
             throw new Error("End-to-end encryption disabled");
         }
-        this.crypto.prepareToEncrypt(room);
+        this.cryptoBackend.prepareToEncrypt(room);
     }
 
     /**
@@ -4392,11 +4396,11 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
             return null;
         }
 
-        if (!this.isRoomEncrypted(event.getRoomId()!)) {
+        if (!room || !this.isRoomEncrypted(event.getRoomId()!)) {
             return null;
         }
 
-        if (!this.crypto && this.usingExternalCrypto) {
+        if (!this.cryptoBackend && this.usingExternalCrypto) {
             // The client has opted to allow sending messages to encrypted
             // rooms even if the room is encrypted, and we haven't setup
             // crypto. This is useful for users of matrix-org/pantalaimon
@@ -4417,13 +4421,11 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
             return null;
         }
 
-        if (!this.crypto) {
-            throw new Error(
-                "This room is configured to use encryption, but your client does " + "not support encryption.",
-            );
+        if (!this.cryptoBackend) {
+            throw new Error("This room is configured to use encryption, but your client does not support encryption.");
         }
 
-        return this.crypto.encryptEvent(event, room);
+        return this.cryptoBackend.encryptEvent(event, room);
     }
 
     /**

src/common-crypto/CryptoBackend.ts

@@ -18,6 +18,7 @@ import type { IEventDecryptionResult, IMegolmSessionData } from "../@types/crypt
 import type { IToDeviceEvent } from "../sync-accumulator";
 import type { DeviceTrustLevel, UserTrustLevel } from "../crypto/CrossSigning";
 import { MatrixEvent } from "../models/event";
+import { Room } from "../models/room";
 import { IEncryptedEventInfo } from "../crypto/api";
 
 /**
@@ -75,6 +76,26 @@ export interface CryptoBackend extends SyncCryptoCallbacks {
      */
     checkDeviceTrust(userId: string, deviceId: string): DeviceTrustLevel;
 
+    /**
+     * Perform any background tasks that can be done before a message is ready to
+     * send, in order to speed up sending of the message.
+     *
+     * @param room - the room the event is in
+     */
+    prepareToEncrypt(room: Room): void;
+
+    /**
+     * Encrypt an event according to the configuration of the room.
+     *
+     * @param event -  event to be sent
+     *
+     * @param room - destination room.
+     *
+     * @returns Promise which resolves when the event has been
+     *     encrypted, or null if nothing was needed
+     */
+    encryptEvent(event: MatrixEvent, room: Room): Promise<void>;
+
     /**
      * Decrypt a received event
      *
@@ -117,6 +138,20 @@ export interface SyncCryptoCallbacks {
      */
     preprocessToDeviceMessages(events: IToDeviceEvent[]): Promise<IToDeviceEvent[]>;
 
+    /**
+     * Called by the /sync loop whenever an m.room.encryption event is received.
+     *
+     * This is called before RoomStateEvents are emitted for any of the events in the /sync
+     * response (even if the other events technically happened first). This works around a problem
+     * if the client uses a RoomStateEvent (typically a membership event) as a trigger to send a message
+     * in a new room (or one where encryption has been newly enabled): that would otherwise leave the
+     * crypto layer confused because it expects crypto to be set up, but it has not yet been.
+     *
+     * @param room - in which the event was received
+     * @param event - encryption event to be processed
+     */
+    onCryptoEvent(room: Room, event: MatrixEvent): Promise<void>;
+
     /**
      * Called by the /sync loop after each /sync response is processed.
      *

src/crypto/index.ts

@@ -2808,11 +2808,7 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
      * @returns Promise which resolves when the event has been
      *     encrypted, or null if nothing was needed
      */
-    public async encryptEvent(event: MatrixEvent, room?: Room): Promise<void> {
-        if (!room) {
-            throw new Error("Cannot send encrypted messages in unknown rooms");
-        }
-
+    public async encryptEvent(event: MatrixEvent, room: Room): Promise<void> {
         const roomId = event.getRoomId()!;
 
         const alg = this.roomEncryptors.get(roomId);