MatrixRTC: Refactor | Introduce a new Encryption manager (used with experimental to device transport) (#4799)

* refactor: New encryption manager BasicEncryptionManager for todevice

fixup: bad do not commit

* fix: ToDevice transport not setting the sent_ts

* test: BasicEncryptionManager add statistics tests

* code review

* feat: Encryption manager just reshare on new joiner

* refactor: Rename BasicEncryptionManger to RTCEncryptionManager

* fixup: RTC experimental todevice should use new encryption mgr

* fixup: use proper logger hierarchy

* fixup: RTC rollout first key asap even if no members to send to

* fixup: RTC add test for first key use

* fixup! emitting outbound key before anyone registered

* fix: quick patch for transport switch, need test

* test: RTC encryption manager, add test for transport switch

* post rebase fix

* Remove bad corepack commit

* review: cleaning, renaming

* review: cleaning and renaming

* stop using root logger in favor of a parent logger

* post merge fix broken test

* remove corepack again

* fix reverted changes after a merge

* review: Properly deprecate getEncryptionKeys

* review: rename ensureMediaKeyDistribution to ensureKeyDistribution

* review: use OutdatedKeyFilter instead of KeyBuffer

spec/unit/matrixrtc/MatrixRTCSession.spec.ts

@@ -21,6 +21,7 @@ import { MatrixRTCSession, MatrixRTCSessionEvent } from "../../../src/matrixrtc/
 import { type EncryptionKeysEventContent } from "../../../src/matrixrtc/types";
 import { secureRandomString } from "../../../src/randomstring";
 import { makeMockEvent, makeMockRoom, makeMockRoomState, membershipTemplate, makeKey } from "./mocks";
+import { RTCEncryptionManager } from "../../../src/matrixrtc/RTCEncryptionManager.ts";
 
 const mockFocus = { type: "mock" };
 
@@ -745,11 +746,27 @@ describe("MatrixRTCSession", () => {
                     expect(sendKeySpy).toHaveBeenCalledTimes(1);
                     // check that we send the key with index 1 even though the send gets delayed when leaving.
                     // this makes sure we do not use an index that is one too old.
-                    expect(sendKeySpy).toHaveBeenLastCalledWith(expect.any(String), 1, sess.memberships);
+                    expect(sendKeySpy).toHaveBeenLastCalledWith(
+                        expect.any(String),
+                        1,
+                        sess.memberships.map((m) => ({
+                            userId: m.sender,
+                            deviceId: m.deviceId,
+                            membershipTs: m.createdTs(),
+                        })),
+                    );
                     // fake a condition in which we send another encryption key event.
                     // this could happen do to someone joining the call.
                     (sess as unknown as any).encryptionManager.sendEncryptionKeysEvent();
-                    expect(sendKeySpy).toHaveBeenLastCalledWith(expect.any(String), 1, sess.memberships);
+                    expect(sendKeySpy).toHaveBeenLastCalledWith(
+                        expect.any(String),
+                        1,
+                        sess.memberships.map((m) => ({
+                            userId: m.sender,
+                            deviceId: m.deviceId,
+                            membershipTs: m.createdTs(),
+                        })),
+                    );
                     jest.advanceTimersByTime(7000);
 
                     const secondKeysPayload = await keysSentPromise2;
@@ -862,10 +879,14 @@ describe("MatrixRTCSession", () => {
                         manageMediaKeys: true,
                         useExperimentalToDeviceTransport: true,
                     });
+                    sess.onRTCSessionMemberUpdate();
 
                     await keySentPromise;
 
                     expect(sendToDeviceMock).toHaveBeenCalled();
+
+                    // Access private to test
+                    expect(sess["encryptionManager"]).toBeInstanceOf(RTCEncryptionManager);
                 } finally {
                     jest.useRealTimers();
                 }

spec/unit/matrixrtc/OutdatedKeyFilter.spec.ts

@@ -0,0 +1,43 @@
+/*
+Copyright 2025 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 { OutdatedKeyFilter } from "../../../src/matrixrtc/utils.ts";
+import { type InboundEncryptionSession } from "../../../src/matrixrtc";
+
+describe("OutdatedKeyFilter Test", () => {
+    it("Should buffer and disambiguate keys by timestamp", () => {
+        const filter = new OutdatedKeyFilter();
+
+        const aKey = fakeInboundSessionWithTimestamp(1000);
+        const olderKey = fakeInboundSessionWithTimestamp(300);
+        // Simulate receiving out of order keys
+
+        expect(filter.isOutdated(aKey.participantId, aKey)).toBe(false);
+        // Then we receive the most recent key out of order
+        const isOutdated = filter.isOutdated(aKey.participantId, olderKey);
+        // this key is older and should be ignored even if received after
+        expect(isOutdated).toBe(true);
+    });
+
+    function fakeInboundSessionWithTimestamp(ts: number): InboundEncryptionSession {
+        return {
+            keyIndex: 0,
+            creationTS: ts,
+            participantId: "@alice:localhost|ABCDE",
+            key: new Uint8Array(16),
+        };
+    }
+});

spec/unit/matrixrtc/RTCEncrytionManager.spec.ts

@@ -0,0 +1,598 @@
+/*
+Copyright 2025 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 { type Mocked } from "jest-mock";
+
+import { RTCEncryptionManager } from "../../../src/matrixrtc/RTCEncryptionManager.ts";
+import { type CallMembership, type Statistics } from "../../../src/matrixrtc";
+import { type ToDeviceKeyTransport } from "../../../src/matrixrtc/ToDeviceKeyTransport.ts";
+import { KeyTransportEvents, type KeyTransportEventsHandlerMap } from "../../../src/matrixrtc/IKeyTransport.ts";
+import { membershipTemplate, mockCallMembership } from "./mocks.ts";
+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";
+
+describe("RTCEncryptionManager", () => {
+    // The manager being tested
+    let encryptionManager: RTCEncryptionManager;
+    let getMembershipMock: jest.Mock;
+    let mockTransport: Mocked<ToDeviceKeyTransport>;
+    let statistics: Statistics;
+    let onEncryptionKeysChanged: jest.Mock;
+
+    beforeEach(() => {
+        statistics = {
+            counters: {
+                roomEventEncryptionKeysSent: 0,
+                roomEventEncryptionKeysReceived: 0,
+            },
+            totals: {
+                roomEventEncryptionKeysReceivedTotalAge: 0,
+            },
+        };
+        getMembershipMock = jest.fn().mockReturnValue([]);
+        onEncryptionKeysChanged = jest.fn();
+        mockTransport = {
+            start: jest.fn(),
+            stop: jest.fn(),
+            sendKey: jest.fn().mockResolvedValue(undefined),
+            on: jest.fn(),
+            off: jest.fn(),
+        } as unknown as Mocked<ToDeviceKeyTransport>;
+
+        encryptionManager = new RTCEncryptionManager(
+            "@alice:example.org",
+            "DEVICE01",
+            getMembershipMock,
+            mockTransport,
+            statistics,
+            onEncryptionKeysChanged,
+        );
+    });
+
+    it("should start and stop the transport properly", () => {
+        encryptionManager.join(undefined);
+
+        expect(mockTransport.start).toHaveBeenCalledTimes(1);
+        expect(mockTransport.on).toHaveBeenCalledTimes(1);
+        expect(mockTransport.on).toHaveBeenCalledWith(KeyTransportEvents.ReceivedKeys, expect.any(Function));
+        encryptionManager.leave();
+        expect(mockTransport.stop).toHaveBeenCalledTimes(1);
+        expect(mockTransport.off).toHaveBeenCalledWith(KeyTransportEvents.ReceivedKeys, expect.any(Function));
+    });
+
+    describe("Sharing Keys", () => {
+        it("Set up my key asap even if no key distribution is needed", () => {
+            getMembershipMock.mockReturnValue([]);
+
+            encryptionManager.join(undefined);
+            // After join it is too early, key might be lost as no one is listening yet
+            expect(onEncryptionKeysChanged).not.toHaveBeenCalled();
+            encryptionManager.onMembershipsUpdate([]);
+            // The key should have been rolled out immediately
+            expect(onEncryptionKeysChanged).toHaveBeenCalled();
+        });
+
+        it("Should distribute keys to members on join", async () => {
+            const members = [
+                aCallMembership("@bob:example.org", "BOBDEVICE"),
+                aCallMembership("@bob:example.org", "BOBDEVICE2"),
+                aCallMembership("@carl:example.org", "CARLDEVICE"),
+            ];
+            getMembershipMock.mockReturnValue(members);
+
+            encryptionManager.join(undefined);
+            encryptionManager.onMembershipsUpdate([]);
+
+            expect(mockTransport.sendKey).toHaveBeenCalledTimes(1);
+            expect(mockTransport.sendKey).toHaveBeenCalledWith(
+                expect.any(String),
+                // It is the first key
+                0,
+                members.map((m) => ({ userId: m.sender, deviceId: m.deviceId, membershipTs: m.createdTs() })),
+            );
+            await jest.runOnlyPendingTimersAsync();
+            // The key should have been rolled out immediately
+            expect(onEncryptionKeysChanged).toHaveBeenCalled();
+            expect(onEncryptionKeysChanged).toHaveBeenCalledWith(
+                expect.any(Uint8Array<ArrayBufferLike>),
+                0,
+                "@alice:example.org:DEVICE01",
+            );
+        });
+
+        it("Should re-distribute keys to members whom callMemberhsip ts has changed", async () => {
+            let members = [aCallMembership("@bob:example.org", "BOBDEVICE", 1000)];
+            getMembershipMock.mockReturnValue(members);
+
+            encryptionManager.join(undefined);
+            encryptionManager.onMembershipsUpdate([]);
+
+            expect(mockTransport.sendKey).toHaveBeenCalledTimes(1);
+            expect(mockTransport.sendKey).toHaveBeenCalledWith(
+                expect.any(String),
+                // It is the first key
+                0,
+                [
+                    {
+                        userId: "@bob:example.org",
+                        deviceId: "BOBDEVICE",
+                        membershipTs: 1000,
+                    },
+                ],
+            );
+            await jest.runOnlyPendingTimersAsync();
+            // The key should have been rolled out immediately
+            expect(onEncryptionKeysChanged).toHaveBeenCalled();
+
+            mockTransport.sendKey.mockClear();
+            onEncryptionKeysChanged.mockClear();
+
+            members = [aCallMembership("@bob:example.org", "BOBDEVICE", 2000)];
+            getMembershipMock.mockReturnValue(members);
+
+            // There are no membership change but the callMembership ts has changed (reset?)
+            // Resend the key
+            encryptionManager.onMembershipsUpdate(members);
+            await jest.runOnlyPendingTimersAsync();
+
+            expect(mockTransport.sendKey).toHaveBeenCalledTimes(1);
+            expect(mockTransport.sendKey).toHaveBeenCalledWith(
+                expect.any(String),
+                // Re send the same key to that user
+                0,
+                [
+                    {
+                        userId: "@bob:example.org",
+                        deviceId: "BOBDEVICE",
+                        membershipTs: 2000,
+                    },
+                ],
+            );
+        });
+
+        it("Should not rotate key when a user join", async () => {
+            jest.useFakeTimers();
+
+            const members = [
+                aCallMembership("@bob:example.org", "BOBDEVICE"),
+                aCallMembership("@bob:example.org", "BOBDEVICE2"),
+            ];
+            getMembershipMock.mockReturnValue(members);
+
+            // initial rollout
+            encryptionManager.join(undefined);
+            encryptionManager.onMembershipsUpdate([]);
+            await jest.runOnlyPendingTimersAsync();
+
+            expect(mockTransport.sendKey).toHaveBeenCalledTimes(1);
+            expect(mockTransport.sendKey).toHaveBeenCalledWith(
+                expect.any(String),
+                // It is the first key
+                0,
+                members.map((m) => ({ userId: m.sender, deviceId: m.deviceId, membershipTs: m.createdTs() })),
+            );
+            onEncryptionKeysChanged.mockClear();
+            mockTransport.sendKey.mockClear();
+
+            const updatedMembers = [
+                aCallMembership("@bob:example.org", "BOBDEVICE"),
+                aCallMembership("@bob:example.org", "BOBDEVICE2"),
+                aCallMembership("@carl:example.org", "CARLDEVICE"),
+            ];
+            getMembershipMock.mockReturnValue(updatedMembers);
+
+            encryptionManager.onMembershipsUpdate(updatedMembers);
+
+            await jest.runOnlyPendingTimersAsync();
+
+            expect(mockTransport.sendKey).toHaveBeenCalledWith(
+                expect.any(String),
+                // It should not have incremented the key index
+                0,
+                // And send it to the newly joined only
+                [{ userId: "@carl:example.org", deviceId: "CARLDEVICE", membershipTs: 1000 }],
+            );
+
+            expect(onEncryptionKeysChanged).not.toHaveBeenCalled();
+            await jest.advanceTimersByTimeAsync(1000);
+
+            expect(statistics.counters.roomEventEncryptionKeysSent).toBe(2);
+        });
+
+        it("Should not resend keys when no changes", async () => {
+            jest.useFakeTimers();
+
+            const members = [
+                aCallMembership("@bob:example.org", "BOBDEVICE"),
+                aCallMembership("@bob:example.org", "BOBDEVICE2"),
+            ];
+            getMembershipMock.mockReturnValue(members);
+
+            // initial rollout
+            encryptionManager.join(undefined);
+            encryptionManager.onMembershipsUpdate([]);
+            await jest.runOnlyPendingTimersAsync();
+
+            expect(mockTransport.sendKey).toHaveBeenCalledTimes(1);
+            onEncryptionKeysChanged.mockClear();
+            mockTransport.sendKey.mockClear();
+
+            encryptionManager.onMembershipsUpdate(members);
+            await jest.advanceTimersByTimeAsync(200);
+            encryptionManager.onMembershipsUpdate(members);
+            await jest.advanceTimersByTimeAsync(100);
+            encryptionManager.onMembershipsUpdate(members);
+            await jest.advanceTimersByTimeAsync(50);
+            encryptionManager.onMembershipsUpdate(members);
+            await jest.advanceTimersByTimeAsync(100);
+
+            expect(mockTransport.sendKey).not.toHaveBeenCalled();
+        });
+
+        it("Should rotate key when a user leaves and delay the rollout", async () => {
+            jest.useFakeTimers();
+
+            const members = [
+                aCallMembership("@bob:example.org", "BOBDEVICE"),
+                aCallMembership("@bob:example.org", "BOBDEVICE2"),
+                aCallMembership("@carl:example.org", "CARLDEVICE"),
+            ];
+            getMembershipMock.mockReturnValue(members);
+
+            encryptionManager.join(undefined);
+            encryptionManager.onMembershipsUpdate([]);
+            await jest.advanceTimersByTimeAsync(10);
+
+            expect(mockTransport.sendKey).toHaveBeenCalledTimes(1);
+            expect(mockTransport.sendKey).toHaveBeenCalledWith(
+                expect.any(String),
+                // It is the first key
+                0,
+                members.map((m) => ({ userId: m.sender, deviceId: m.deviceId, membershipTs: m.createdTs() })),
+            );
+            // initial rollout
+            expect(mockTransport.sendKey).toHaveBeenCalled();
+            expect(onEncryptionKeysChanged).toHaveBeenCalledTimes(1);
+            onEncryptionKeysChanged.mockClear();
+
+            const updatedMembers = [
+                aCallMembership("@bob:example.org", "BOBDEVICE"),
+                aCallMembership("@bob:example.org", "BOBDEVICE2"),
+            ];
+            getMembershipMock.mockReturnValue(updatedMembers);
+
+            encryptionManager.onMembershipsUpdate(updatedMembers);
+
+            await jest.advanceTimersByTimeAsync(200);
+            // The is rotated but not rolled out yet to give time for the key to be sent
+            expect(mockTransport.sendKey).toHaveBeenCalledWith(
+                expect.any(String),
+                // It should have incremented the key index
+                1,
+                // And send it to the updated members
+                updatedMembers.map((m) => ({ userId: m.sender, deviceId: m.deviceId, membershipTs: m.createdTs() })),
+            );
+
+            expect(onEncryptionKeysChanged).not.toHaveBeenCalled();
+            await jest.advanceTimersByTimeAsync(1000);
+
+            // now should be rolled out
+            expect(onEncryptionKeysChanged).toHaveBeenCalledWith(
+                expect.any(Uint8Array<ArrayBufferLike>),
+                1,
+                "@alice:example.org:DEVICE01",
+            );
+
+            expect(statistics.counters.roomEventEncryptionKeysSent).toBe(2);
+        });
+    });
+
+    describe("Receiving Keys", () => {
+        beforeEach(() => {
+            const emitter = new TypedEventEmitter<KeyTransportEvents, KeyTransportEventsHandlerMap>();
+            mockTransport = {
+                start: jest.fn(),
+                stop: jest.fn(),
+                sendKey: jest.fn().mockResolvedValue(undefined),
+                on: emitter.on.bind(emitter),
+                off: emitter.off.bind(emitter),
+                emit: emitter.emit.bind(emitter),
+            } as unknown as Mocked<ToDeviceKeyTransport>;
+            encryptionManager = new RTCEncryptionManager(
+                "@alice:example.org",
+                "DEVICE01",
+                getMembershipMock,
+                mockTransport,
+                statistics,
+                onEncryptionKeysChanged,
+            );
+        });
+
+        it("should accept keys from transport", async () => {
+            jest.useFakeTimers();
+
+            const members = [
+                aCallMembership("@bob:example.org", "BOBDEVICE"),
+                aCallMembership("@bob:example.org", "BOBDEVICE2"),
+                aCallMembership("@carl:example.org", "CARLDEVICE"),
+            ];
+            getMembershipMock.mockReturnValue(members);
+
+            encryptionManager.join(undefined);
+            encryptionManager.onMembershipsUpdate([]);
+            await jest.advanceTimersByTimeAsync(10);
+
+            mockTransport.emit(
+                KeyTransportEvents.ReceivedKeys,
+                "@bob:example.org",
+                "BOBDEVICE",
+                "AAAAAAAAAAA",
+                0 /* KeyId */,
+                0 /* Timestamp */,
+            );
+            mockTransport.emit(
+                KeyTransportEvents.ReceivedKeys,
+                "@bob:example.org",
+                "BOBDEVICE2",
+                "BBBBBBBBBBB",
+                4 /* KeyId */,
+                0 /* Timestamp */,
+            );
+            mockTransport.emit(
+                KeyTransportEvents.ReceivedKeys,
+                "@carl:example.org",
+                "CARLDEVICE",
+                "CCCCCCCCCC",
+                8 /* KeyId */,
+                0 /* Timestamp */,
+            );
+
+            expect(onEncryptionKeysChanged).toHaveBeenCalledTimes(4);
+            expect(onEncryptionKeysChanged).toHaveBeenCalledWith(
+                decodeBase64("AAAAAAAAAAA"),
+                0,
+                "@bob:example.org:BOBDEVICE",
+            );
+
+            expect(onEncryptionKeysChanged).toHaveBeenCalledWith(
+                decodeBase64("BBBBBBBBBBB"),
+                4,
+                "@bob:example.org:BOBDEVICE2",
+            );
+
+            expect(onEncryptionKeysChanged).toHaveBeenCalledWith(
+                decodeBase64("CCCCCCCCCC"),
+                8,
+                "@carl:example.org:CARLDEVICE",
+            );
+
+            expect(statistics.counters.roomEventEncryptionKeysReceived).toBe(3);
+        });
+
+        it("Should support quick re-joiner if keys received out of order", async () => {
+            jest.useFakeTimers();
+
+            const members = [aCallMembership("@carl:example.org", "CARLDEVICE")];
+            getMembershipMock.mockReturnValue(members);
+
+            // Let's join
+            encryptionManager.join(undefined);
+            await jest.advanceTimersByTimeAsync(10);
+
+            // Simulate Carl leaving then joining back, and key received out of order
+            const initialKey0TimeStamp = 1000;
+            const newKey0TimeStamp = 2000;
+
+            mockTransport.emit(
+                KeyTransportEvents.ReceivedKeys,
+                "@carol:example.org",
+                "CAROLDEVICE",
+                "BBBBBBBBBBB",
+                0 /* KeyId */,
+                newKey0TimeStamp,
+            );
+
+            await jest.advanceTimersByTimeAsync(20);
+
+            mockTransport.emit(
+                KeyTransportEvents.ReceivedKeys,
+                "@carol:example.org",
+                "CAROLDEVICE",
+                "AAAAAAAAAAA",
+                0 /* KeyId */,
+                initialKey0TimeStamp,
+            );
+
+            await jest.advanceTimersByTimeAsync(20);
+
+            // The latest key used for carol should be the one with the latest timestamp
+
+            expect(onEncryptionKeysChanged).toHaveBeenLastCalledWith(
+                decodeBase64("BBBBBBBBBBB"),
+                0,
+                "@carol:example.org:CAROLDEVICE",
+            );
+        });
+    });
+
+    it("Should only rotate once again if several membership changes during a rollout", async () => {
+        jest.useFakeTimers();
+
+        let 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([]);
+        await jest.advanceTimersByTimeAsync(10);
+
+        // The initial rollout
+        expect(onEncryptionKeysChanged).toHaveBeenCalledWith(
+            expect.any(Uint8Array<ArrayBufferLike>),
+            0,
+            "@alice:example.org:DEVICE01",
+        );
+        onEncryptionKeysChanged.mockClear();
+
+        // Trigger a key rotation with a leaver
+        members = [aCallMembership("@bob:example.org", "BOBDEVICE"), aCallMembership("@bob:example.org", "BOBDEVICE2")];
+        getMembershipMock.mockReturnValue(members);
+
+        // This should start a new key rollout
+        encryptionManager.onMembershipsUpdate(members);
+        await jest.advanceTimersByTimeAsync(10);
+
+        // Now simulate a new leaver
+        members = [aCallMembership("@bob:example.org", "BOBDEVICE")];
+        getMembershipMock.mockReturnValue(members);
+
+        // The key `1` rollout is in progress
+        encryptionManager.onMembershipsUpdate(members);
+        await jest.advanceTimersByTimeAsync(10);
+
+        // And another one ( plus a joiner)
+        const lastMembership = [aCallMembership("@bob:example.org", "BOBDEVICE3")];
+        getMembershipMock.mockReturnValue(lastMembership);
+        // The key `1` rollout is still in progress
+        encryptionManager.onMembershipsUpdate(lastMembership);
+        await jest.advanceTimersByTimeAsync(10);
+
+        // Let all rollouts finish
+        await jest.advanceTimersByTimeAsync(2000);
+
+        // There should 2 rollout. The `1` rollout, then just one additional one
+        // that has "buffered" the 2 membership changes with leavers
+        expect(onEncryptionKeysChanged).toHaveBeenCalledTimes(2);
+        expect(onEncryptionKeysChanged).toHaveBeenCalledWith(
+            expect.any(Uint8Array<ArrayBufferLike>),
+            1,
+            "@alice:example.org:DEVICE01",
+        );
+        expect(onEncryptionKeysChanged).toHaveBeenCalledWith(
+            expect.any(Uint8Array<ArrayBufferLike>),
+            2,
+            "@alice:example.org:DEVICE01",
+        );
+
+        // Key `2` should only be distributed to the last membership
+        expect(mockTransport.sendKey).toHaveBeenLastCalledWith(
+            expect.any(String),
+            2,
+            // And send only to the last membership
+            [
+                {
+                    userId: "@bob:example.org",
+                    deviceId: "BOBDEVICE3",
+                    membershipTs: 1000,
+                },
+            ],
+        );
+    });
+
+    it("Should re-distribute key on transport switch", async () => {
+        const toDeviceEmitter = new TypedEventEmitter<KeyTransportEvents, KeyTransportEventsHandlerMap>();
+        const mockToDeviceTransport = {
+            start: jest.fn(),
+            stop: jest.fn(),
+            sendKey: jest.fn().mockResolvedValue(undefined),
+            on: toDeviceEmitter.on.bind(toDeviceEmitter),
+            off: toDeviceEmitter.off.bind(toDeviceEmitter),
+            emit: toDeviceEmitter.emit.bind(toDeviceEmitter),
+            setParentLogger: jest.fn(),
+        } as unknown as Mocked<ToDeviceKeyTransport>;
+
+        const roomEmitter = new TypedEventEmitter<KeyTransportEvents, KeyTransportEventsHandlerMap>();
+        const mockRoomTransport = {
+            start: jest.fn(),
+            stop: jest.fn(),
+            sendKey: jest.fn().mockResolvedValue(undefined),
+            on: roomEmitter.on.bind(roomEmitter),
+            off: roomEmitter.off.bind(roomEmitter),
+            emit: roomEmitter.emit.bind(roomEmitter),
+            setParentLogger: jest.fn(),
+        } as unknown as Mocked<RoomKeyTransport>;
+
+        const mockLogger = {
+            debug: jest.fn(),
+            warn: jest.fn(),
+        } as unknown as Mocked<Logger>;
+
+        const transport = new RoomAndToDeviceTransport(mockToDeviceTransport, mockRoomTransport, {
+            getChild: jest.fn().mockReturnValue(mockLogger),
+        } as unknown as Mocked<Logger>);
+
+        encryptionManager = new RTCEncryptionManager(
+            "@alice:example.org",
+            "DEVICE01",
+            getMembershipMock,
+            transport,
+            statistics,
+            onEncryptionKeysChanged,
+        );
+
+        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([]);
+        await jest.advanceTimersByTimeAsync(10);
+
+        // Should have sent the key to the toDevice transport
+        expect(mockToDeviceTransport.sendKey).toHaveBeenCalledTimes(1);
+        expect(mockRoomTransport.sendKey).not.toHaveBeenCalled();
+
+        // Simulate receiving a key by room transport
+        roomEmitter.emit(
+            KeyTransportEvents.ReceivedKeys,
+            "@bob:example.org",
+            "BOBDEVICE",
+            "AAAAAAAAAAA",
+            0 /* KeyId */,
+            0 /* Timestamp */,
+        );
+
+        await jest.runOnlyPendingTimersAsync();
+
+        // The key should have beed re-distributed to the room transport
+        expect(mockRoomTransport.sendKey).toHaveBeenCalled();
+        expect(mockToDeviceTransport.sendKey).toHaveBeenCalledWith(
+            expect.any(String),
+            // It is the first key re-distributed
+            0,
+            // to all the members
+            members.map((m) => ({ userId: m.sender, deviceId: m.deviceId, membershipTs: m.createdTs() })),
+        );
+    });
+
+    function aCallMembership(userId: string, deviceId: string, ts: number = 1000): CallMembership {
+        return mockCallMembership(
+            Object.assign({}, membershipTemplate, { device_id: deviceId, created_ts: ts }),
+            "!room:id",
+            userId,
+        );
+    }
+});

spec/unit/matrixrtc/RoomAndToDeviceTransport.spec.ts

@@ -16,7 +16,7 @@ limitations under the License.
 
 import { type Mocked } from "jest-mock";
 
-import { makeKey, makeMockEvent, makeMockRoom, membershipTemplate, mockCallMembership } from "./mocks";
+import { makeKey, makeMockEvent, makeMockRoom } from "./mocks";
 import { EventType, type IRoomTimelineData, type Room, RoomEvent, type MatrixClient } from "../../../src";
 import { ToDeviceKeyTransport } from "../../../src/matrixrtc/ToDeviceKeyTransport.ts";
 import {
@@ -24,7 +24,7 @@ import {
     mockClientMethodsEvents,
     mockClientMethodsUser,
 } from "../../test-utils/client.ts";
-import { type Statistics } from "../../../src/matrixrtc";
+import { type ParticipantDeviceInfo, type Statistics } from "../../../src/matrixrtc";
 import { KeyTransportEvents } from "../../../src/matrixrtc/IKeyTransport.ts";
 import { type Logger } from "../../../src/logger.ts";
 import { RoomAndToDeviceEvents, RoomAndToDeviceTransport } from "../../../src/matrixrtc/RoomAndToDeviceKeyTransport.ts";
@@ -88,7 +88,9 @@ describe("RoomAndToDeviceTransport", () => {
     });
     it("only sends to device keys when sending a key", async () => {
         transport.start();
-        await transport.sendKey("1235", 0, [mockCallMembership(membershipTemplate, roomId, "@alice:example.org")]);
+        await transport.sendKey("1235", 0, [
+            { userId: "@alice:example.org", deviceId: "ALICEDEVICE", membershipTs: 1234 },
+        ]);
         expect(toDeviceSendKeySpy).toHaveBeenCalledTimes(1);
         expect(roomSendKeySpy).toHaveBeenCalledTimes(0);
         expect(transport.enabled.room).toBeFalsy();
@@ -118,7 +120,9 @@ describe("RoomAndToDeviceTransport", () => {
         expect(transport.enabled.room).toBeTruthy();
         expect(transport.enabled.toDevice).toBeFalsy();
 
-        await transport.sendKey("1235", 0, [mockCallMembership(membershipTemplate, roomId, "@alice:example.org")]);
+        await transport.sendKey("1235", 0, [
+            { userId: "@alice:example.org", deviceId: "AlICEDEV", membershipTs: 1234 },
+        ]);
         expect(sendEventMock).toHaveBeenCalledTimes(1);
         expect(roomSendKeySpy).toHaveBeenCalledTimes(1);
         expect(toDeviceSendKeySpy).toHaveBeenCalledTimes(0);
@@ -132,7 +136,11 @@ describe("RoomAndToDeviceTransport", () => {
         });
 
         transport.start();
-        const membership = mockCallMembership(membershipTemplate, roomId, "@alice:example.org");
+        const membership: ParticipantDeviceInfo = {
+            userId: "@alice:example.org",
+            deviceId: "ALICEDEVICE",
+            membershipTs: 1234,
+        };
         const onTransportEnabled = jest.fn();
         transport.on(RoomAndToDeviceEvents.EnabledTransportsChanged, onTransportEnabled);
 

spec/unit/matrixrtc/ToDeviceKeyTransport.spec.ts

@@ -16,7 +16,7 @@ limitations under the License.
 
 import { type Mocked } from "jest-mock";
 
-import { makeMockEvent, membershipTemplate, mockCallMembership } from "./mocks";
+import { makeMockEvent } from "./mocks";
 import { ClientEvent, EventType, type MatrixClient } from "../../../src";
 import { ToDeviceKeyTransport } from "../../../src/matrixrtc/ToDeviceKeyTransport.ts";
 import { getMockClientWithEventEmitter } from "../../test-utils/client.ts";
@@ -61,21 +61,9 @@ describe("ToDeviceKeyTransport", () => {
         const keyBase64Encoded = "ABCDEDF";
         const keyIndex = 2;
         await transport.sendKey(keyBase64Encoded, keyIndex, [
-            mockCallMembership(
-                Object.assign({}, membershipTemplate, { device_id: "BOBDEVICE" }),
-                roomId,
-                "@bob:example.org",
-            ),
-            mockCallMembership(
-                Object.assign({}, membershipTemplate, { device_id: "CARLDEVICE" }),
-                roomId,
-                "@carl:example.org",
-            ),
-            mockCallMembership(
-                Object.assign({}, membershipTemplate, { device_id: "MATDEVICE" }),
-                roomId,
-                "@mat:example.org",
-            ),
+            { userId: "@bob:example.org", deviceId: "BOBDEVICE", membershipTs: 1234 },
+            { userId: "@carl:example.org", deviceId: "CARLDEVICE", membershipTs: 1234 },
+            { userId: "@mat:example.org", deviceId: "MATDEVICE", membershipTs: 1234 },
         ]);
 
         expect(mockClient.encryptAndSendToDevice).toHaveBeenCalledTimes(1);
@@ -100,6 +88,7 @@ describe("ToDeviceKeyTransport", () => {
                     call_id: "",
                     scope: "m.room",
                 },
+                sent_ts: expect.any(Number),
             },
         );
 
@@ -153,11 +142,7 @@ describe("ToDeviceKeyTransport", () => {
         const keyBase64Encoded = "ABCDEDF";
         const keyIndex = 2;
         await transport.sendKey(keyBase64Encoded, keyIndex, [
-            mockCallMembership(
-                Object.assign({}, membershipTemplate, { device_id: "MYDEVICE" }),
-                roomId,
-                "@alice:example.org",
-            ),
+            { userId: "@alice:example.org", deviceId: "MYDEVICE", membershipTs: 1234 },
         ]);
 
         transport.start();