Element-R: Add support for /discardsession (#3209)

Fixes https://github.com/vector-im/element-web/issues/24431

spec/integ/crypto.spec.ts

@@ -342,6 +342,11 @@ async function expectSendRoomKey(
                 resolve(onSendRoomKey(content));
                 return {};
             },
+            {
+                // append to the list of intercepts on this path (since we have some tests that call
+                // this function multiple times)
+                overwriteRoutes: false,
+            },
         );
     });
 }
@@ -360,12 +365,20 @@ async function expectSendMegolmMessage(
     inboundGroupSessionPromise: Promise<Olm.InboundGroupSession>,
 ): Promise<Partial<IEvent>> {
     const encryptedMessageContent = await new Promise<IContent>((resolve) => {
-        fetchMock.putOnce(new RegExp("/send/m.room.encrypted/"), (url: string, opts: RequestInit): MockResponse => {
+        fetchMock.putOnce(
+            new RegExp("/send/m.room.encrypted/"),
+            (url: string, opts: RequestInit): MockResponse => {
                 resolve(JSON.parse(opts.body as string));
                 return {
                     event_id: "$event_id",
                 };
-        });
+            },
+            {
+                // append to the list of intercepts on this path (since we have some tests that call
+                // this function multiple times)
+                overwriteRoutes: false,
+            },
+        );
     });
 
     // In some of the tests, the room key is sent *after* the actual event, so we may need to wait for it now.
@@ -808,6 +821,40 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("crypto (%s)", (backend: string,
         ]);
     });
 
+    it("We should start a new megolm session after forceDiscardSession", async () => {
+        aliceClient.setGlobalErrorOnUnknownDevices(false);
+        expectAliceKeyQuery({ device_keys: { "@alice:localhost": {} }, failures: {} });
+        await startClientAndAwaitFirstSync();
+
+        // Alice shares a room with Bob
+        syncResponder.sendOrQueueSyncResponse(getSyncResponse(["@bob:xyz"]));
+        await syncPromise(aliceClient);
+
+        // Once we send the message, Alice will check Bob's device list (twice, because reasons) ...
+        expectAliceKeyQuery(getTestKeysQueryResponse("@bob:xyz"));
+        expectAliceKeyQuery(getTestKeysQueryResponse("@bob:xyz"));
+
+        // ... and claim one of his OTKs ...
+        expectAliceKeyClaim(getTestKeysClaimResponse("@bob:xyz"));
+
+        // ... and send an m.room_key message
+        const inboundGroupSessionPromise = expectSendRoomKey("@bob:xyz", testOlmAccount);
+
+        // Send the first message, and check we can decrypt it.
+        await Promise.all([
+            aliceClient.sendTextMessage(ROOM_ID, "test"),
+            expectSendMegolmMessage(inboundGroupSessionPromise),
+        ]);
+
+        // Finally the interesting part: discard the session.
+        aliceClient.forceDiscardSession(ROOM_ID);
+
+        // Now when we send the next message, we should get a *new* megolm session.
+        const inboundGroupSessionPromise2 = expectSendRoomKey("@bob:xyz", testOlmAccount);
+        const p2 = expectSendMegolmMessage(inboundGroupSessionPromise2);
+        await Promise.all([aliceClient.sendTextMessage(ROOM_ID, "test2"), p2]);
+    });
+
     oldBackendOnly("Alice sends a megolm message", async () => {
         // TODO: do something about this for the rust backend.
         //   Currently it fails because we don't respect the default GlobalErrorOnUnknownDevices and

src/client.ts

@@ -3130,13 +3130,14 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
      *
      * @param roomId - The ID of the room to discard the session for
      *
-     * This should not normally be necessary.
+     * @deprecated Prefer {@link CryptoApi.forceDiscardSession | `CryptoApi.forceDiscardSession`}:
+     *
      */
     public forceDiscardSession(roomId: string): void {
-        if (!this.crypto) {
+        if (!this.cryptoBackend) {
             throw new Error("End-to-End encryption disabled");
         }
-        this.crypto.forceDiscardSession(roomId);
+        this.cryptoBackend.forceDiscardSession(roomId);
     }
 
     /**

src/crypto-api.ts

@@ -50,6 +50,19 @@ export interface CryptoApi {
      */
     prepareToEncrypt(room: Room): void;
 
+    /**
+     * Discard any existing megolm session for the given room.
+     *
+     * This will ensure that a new session is created on the next call to {@link prepareToEncrypt},
+     * or the next time a message is sent.
+     *
+     * This should not normally be necessary: it should only be used as a debugging tool if there has been a
+     * problem with encryption.
+     *
+     * @param roomId - the room to discard sessions for
+     */
+    forceDiscardSession(roomId: string): Promise<void>;
+
     /**
      * Get a list containing all of the room keys
      *

src/crypto/index.ts

@@ -2492,13 +2492,14 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
      *
      * This should not normally be necessary.
      */
-    public forceDiscardSession(roomId: string): void {
+    public forceDiscardSession(roomId: string): Promise<void> {
         const alg = this.roomEncryptors.get(roomId);
         if (alg === undefined) throw new Error("Room not encrypted");
         if (alg.forceDiscardSession === undefined) {
             throw new Error("Room encryption algorithm doesn't support session discarding");
         }
         alg.forceDiscardSession();
+        return Promise.resolve();
     }
 
     /**

src/rust-crypto/RoomEncryptor.ts

@@ -116,6 +116,16 @@ export class RoomEncryptor {
         }
     }
 
+    /**
+     * Discard any existing group session for this room
+     */
+    public async forceDiscardSession(): Promise<void> {
+        const r = await this.olmMachine.invalidateGroupSession(new RoomId(this.room.roomId));
+        if (r) {
+            this.prefixedLogger.info("Discarded existing group session");
+        }
+    }
+
     /**
      * Encrypt an event for this room
      *

src/rust-crypto/rust-crypto.ts

@@ -169,6 +169,10 @@ export class RustCrypto implements CryptoBackend {
         }
     }
 
+    public forceDiscardSession(roomId: string): Promise<void> {
+        return this.roomEncryptors[roomId]?.forceDiscardSession();
+    }
+
     public async exportRoomKeys(): Promise<IMegolmSessionData[]> {
         // TODO
         return [];