EncryptionManager: un-deprecating EncryptionManager.getEncryptionKeys (#4912)

* EncryptionManager: should be able to re-emit keys

* fix typo in test file name

* review unneeded cast

* remove bad comment

spec/unit/matrixrtc/RTCEncryptionManager.spec.ts

@@ -25,6 +25,7 @@ import { decodeBase64, TypedEventEmitter } from "../../../src";
 import { RoomAndToDeviceTransport } from "../../../src/matrixrtc/RoomAndToDeviceKeyTransport.ts";
 import { type RoomKeyTransport } from "../../../src/matrixrtc/RoomKeyTransport.ts";
 import type { Logger } from "../../../src/logger.ts";
+import { getParticipantId } from "../../../src/matrixrtc/utils.ts";
 
 describe("RTCEncryptionManager", () => {
     // The manager being tested
@@ -428,6 +429,92 @@ describe("RTCEncryptionManager", () => {
                 "@carol:example.org:CAROLDEVICE",
             );
         });
+
+        it("Should store keys for later retrieval", async () => {
+            jest.useFakeTimers();
+
+            const members = [
+                aCallMembership("@bob:example.org", "BOBDEVICE"),
+                aCallMembership("@bob:example.org", "BOBDEVICE2"),
+                aCallMembership("@carl:example.org", "CARLDEVICE"),
+            ];
+            getMembershipMock.mockReturnValue(members);
+
+            // Let's join
+            encryptionManager.join(undefined);
+            encryptionManager.onMembershipsUpdate(members);
+
+            await jest.advanceTimersByTimeAsync(10);
+
+            mockTransport.emit(
+                KeyTransportEvents.ReceivedKeys,
+                "@carl:example.org",
+                "CARLDEVICE",
+                "BBBBBBBBBBB",
+                0 /* KeyId */,
+                1000,
+            );
+
+            mockTransport.emit(
+                KeyTransportEvents.ReceivedKeys,
+                "@carl:example.org",
+                "CARLDEVICE",
+                "CCCCCCCCCCC",
+                5 /* KeyId */,
+                1000,
+            );
+
+            mockTransport.emit(
+                KeyTransportEvents.ReceivedKeys,
+                "@bob:example.org",
+                "BOBDEVICE2",
+                "DDDDDDDDDDD",
+                0 /* KeyId */,
+                1000,
+            );
+
+            const knownKeys = encryptionManager.getEncryptionKeys();
+
+            // My own key should be there
+            const myRing = knownKeys.get(getParticipantId("@alice:example.org", "DEVICE01"));
+            expect(myRing).toBeDefined();
+            expect(myRing).toHaveLength(1);
+            expect(myRing![0]).toMatchObject(
+                expect.objectContaining({
+                    keyIndex: 0,
+                    key: expect.any(Uint8Array),
+                }),
+            );
+
+            const carlRing = knownKeys.get(getParticipantId("@carl:example.org", "CARLDEVICE"));
+            expect(carlRing).toBeDefined();
+            expect(carlRing).toHaveLength(2);
+            expect(carlRing![0]).toMatchObject(
+                expect.objectContaining({
+                    keyIndex: 0,
+                    key: decodeBase64("BBBBBBBBBBB"),
+                }),
+            );
+            expect(carlRing![1]).toMatchObject(
+                expect.objectContaining({
+                    keyIndex: 5,
+                    key: decodeBase64("CCCCCCCCCCC"),
+                }),
+            );
+
+            const bobRing = knownKeys.get(getParticipantId("@bob:example.org", "BOBDEVICE2"));
+            expect(bobRing).toBeDefined();
+            expect(bobRing).toHaveLength(1);
+            expect(bobRing![0]).toMatchObject(
+                expect.objectContaining({
+                    keyIndex: 0,
+                    key: decodeBase64("DDDDDDDDDDD"),
+                }),
+            );
+
+            const bob1Ring = knownKeys.get(getParticipantId("@bob:example.org", "BOBDEVICE"));
+            expect(bob1Ring).not.toBeDefined();
+        });
     });
 
     it("Should only rotate once again if several membership changes during a rollout", async () => {

src/matrixrtc/EncryptionManager.ts

@@ -5,7 +5,7 @@ import { decodeBase64, encodeUnpaddedBase64 } from "../base64.ts";
 import { safeGetRetryAfterMs } from "../http-api/errors.ts";
 import { type CallMembership } from "./CallMembership.ts";
 import { type KeyTransportEventListener, KeyTransportEvents, type IKeyTransport } from "./IKeyTransport.ts";
-import { isMyMembership, type Statistics } from "./types.ts";
+import { isMyMembership, type ParticipantId, type Statistics } from "./types.ts";
 import { getParticipantId } from "./utils.ts";
 import {
     type EnabledTransports,
@@ -41,14 +41,9 @@ export interface IEncryptionManager {
     /**
      * Retrieves the encryption keys currently managed by the encryption manager.
      *
-     * @returns A map where the keys are identifiers and the values are arrays of
+     * @returns A map of participant IDs to their encryption keys.
-     * objects containing encryption keys and their associated timestamps.
-     * @deprecated This method is used internally for testing. It is also used to re-emit keys when there is a change
-     * of RTCSession (matrixKeyProvider#setRTCSession) -Not clear why/when switch RTCSession would occur-. Note that if we switch focus, we do keep the same RTC session,
-     * so no need to re-emit. But it requires the encryption manager to store all keys of all participants, and this is already done
-     * by the key provider. We don't want to add another layer of key storage.
      */
-    getEncryptionKeys(): Map<string, Array<{ key: Uint8Array; timestamp: number }>>;
+    getEncryptionKeys(): ReadonlyMap<ParticipantId, ReadonlyArray<{ key: Uint8Array; keyIndex: number }>>;
 }
 
 /**
@@ -104,8 +99,16 @@ export class EncryptionManager implements IEncryptionManager {
         this.logger = (parentLogger ?? rootLogger).getChild(`[EncryptionManager]`);
     }
 
-    public getEncryptionKeys(): Map<string, Array<{ key: Uint8Array; timestamp: number }>> {
+    public getEncryptionKeys(): ReadonlyMap<ParticipantId, ReadonlyArray<{ key: Uint8Array; keyIndex: number }>> {
-        return this.encryptionKeys;
+        const keysMap = new Map<ParticipantId, ReadonlyArray<{ key: Uint8Array; keyIndex: number }>>();
+        for (const [userId, userKeys] of this.encryptionKeys) {
+            const keys = userKeys.map((entry, index) => ({
+                key: entry.key,
+                keyIndex: index,
+            }));
+            keysMap.set(userId as ParticipantId, keys);
+        }
+        return keysMap;
     }
 
     private joined = false;
@@ -300,7 +303,6 @@ export class EncryptionManager implements IEncryptionManager {
             await this.transport.sendKey(encodeUnpaddedBase64(keyToSend), keyIndexToSend, targets);
             this.logger.debug(
                 `sendEncryptionKeysEvent participantId=${this.userId}:${this.deviceId} numKeys=${myKeys.length} currentKeyIndex=${this.latestGeneratedKeyIndex} keyIndexToSend=${keyIndexToSend}`,
-                this.encryptionKeys,
             );
         } catch (error) {
             if (this.keysEventUpdateTimeout === undefined) {

src/matrixrtc/MatrixRTCSession.ts

@@ -516,29 +516,13 @@ export class MatrixRTCSession extends TypedEventEmitter<
      * the keys.
      */
     public reemitEncryptionKeys(): void {
-        this.encryptionManager?.getEncryptionKeys().forEach((keys, participantId) => {
+        this.encryptionManager?.getEncryptionKeys().forEach((keyRing, participantId) => {
-            keys.forEach((key, index) => {
+            keyRing.forEach((keyInfo) => {
-                this.emit(MatrixRTCSessionEvent.EncryptionKeyChanged, key.key, index, participantId);
+                this.emit(MatrixRTCSessionEvent.EncryptionKeyChanged, keyInfo.key, keyInfo.keyIndex, participantId);
             });
         });
     }
 
-    /**
-     * A map of keys used to encrypt and decrypt (we are using a symmetric
-     * cipher) given participant's media. This also includes our own key
-     *
-     * @deprecated This will be made private in a future release.
-     */
-    public getEncryptionKeys(): IterableIterator<[string, Array<Uint8Array>]> {
-        const keys =
-            this.encryptionManager?.getEncryptionKeys() ??
-            new Map<string, Array<{ key: Uint8Array; timestamp: number }>>();
-        // the returned array doesn't contain the timestamps
-        return Array.from(keys.entries())
-            .map(([participantId, keys]): [string, Uint8Array[]] => [participantId, keys.map((k) => k.key)])
-            .values();
-    }
-
     /**
      * Sets a timer for the soonest membership expiry
      */

src/matrixrtc/RTCEncryptionManager.ts

@@ -46,6 +46,13 @@ import {
  * XXX In the future we want to distribute a ratcheted key not the current one for new joiners.
  */
 export class RTCEncryptionManager implements IEncryptionManager {
+    /**
+     * Store the key rings for each participant.
+     * The encryption manager stores the keys because the application layer might not be ready yet to handle the keys.
+     * The keys are stored and can be retrieved later when the application layer is ready {@link RTCEncryptionManager#getEncryptionKeys}.
+     */
+    private participantKeyRings = new Map<ParticipantId, Array<{ key: Uint8Array; keyIndex: number }>>();
+
     // The current per-sender media key for this device
     private outboundSession: OutboundEncryptionSession | null = null;
 
@@ -94,9 +101,16 @@ export class RTCEncryptionManager implements IEncryptionManager {
         this.logger = parentLogger?.getChild(`[EncryptionManager]`);
     }
 
-    public getEncryptionKeys(): Map<string, Array<{ key: Uint8Array; timestamp: number }>> {
+    public getEncryptionKeys(): ReadonlyMap<ParticipantId, ReadonlyArray<{ key: Uint8Array; keyIndex: number }>> {
-        // This is deprecated should be ignored. Only used by tests?
+        return new Map(this.participantKeyRings);
-        return new Map();
+    }
+
+    private addKeyToParticipant(key: Uint8Array, keyIndex: number, participantId: ParticipantId): void {
+        if (!this.participantKeyRings.has(participantId)) {
+            this.participantKeyRings.set(participantId, []);
+        }
+        this.participantKeyRings.get(participantId)!.push({ key, keyIndex });
+        this.onEncryptionKeysChanged(key, keyIndex, participantId);
     }
 
     public join(joinConfig: EncryptionConfig | undefined): void {
@@ -114,6 +128,7 @@ export class RTCEncryptionManager implements IEncryptionManager {
     public leave(): void {
         this.transport.off(KeyTransportEvents.ReceivedKeys, this.onNewKeyReceived);
         this.transport.stop();
+        this.participantKeyRings.clear();
     }
 
     // Temporary for backwards compatibility
@@ -138,6 +153,7 @@ export class RTCEncryptionManager implements IEncryptionManager {
             }
         }
     };
+
     /**
      * Will ensure that a new key is distributed and used to encrypt our media.
      * If there is already a key distribution in progress, it will schedule a new distribution round just after the current one is completed.
@@ -181,7 +197,7 @@ export class RTCEncryptionManager implements IEncryptionManager {
 
         const outdated = this.keyBuffer.isOutdated(participantId, candidateInboundSession);
         if (!outdated) {
-            this.onEncryptionKeysChanged(
+            this.addKeyToParticipant(
                 candidateInboundSession.key,
                 candidateInboundSession.keyIndex,
                 candidateInboundSession.participantId,
@@ -215,7 +231,7 @@ export class RTCEncryptionManager implements IEncryptionManager {
                 sharedWith: [],
                 keyId: 0,
             };
-            this.onEncryptionKeysChanged(
+            this.addKeyToParticipant(
                 this.outboundSession.key,
                 this.outboundSession.keyId,
                 getParticipantId(this.userId, this.deviceId),
@@ -303,7 +319,7 @@ export class RTCEncryptionManager implements IEncryptionManager {
                 this.logger?.trace(`Delay Rollout for key:${outboundKey.keyId}...`);
                 await sleep(this.delayRolloutTimeMillis);
                 this.logger?.trace(`...Delayed rollout of index:${outboundKey.keyId} `);
-                this.onEncryptionKeysChanged(
+                this.addKeyToParticipant(
                     outboundKey.key,
                     outboundKey.keyId,
                     getParticipantId(this.userId, this.deviceId),