From 3781b6ebfa32536511988e2a8c9e7f65c85844ee Mon Sep 17 00:00:00 2001
From: Hugh Nimmo-Smith <hughns@users.noreply.github.com>
Date: Thu, 28 Nov 2024 12:05:39 +0000
Subject: [PATCH] Re-send MatrixRTC media encryption keys for a new joiner even
 if a rotation is in progress (#4561)

---
 spec/unit/matrixrtc/MatrixRTCSession.spec.ts | 59 +++++++++++++++++++-
 src/matrixrtc/MatrixRTCSession.ts            | 10 +++-
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/spec/unit/matrixrtc/MatrixRTCSession.spec.ts b/spec/unit/matrixrtc/MatrixRTCSession.spec.ts
index fb62a669c7b..26368e0b33f 100644
--- a/spec/unit/matrixrtc/MatrixRTCSession.spec.ts
+++ b/spec/unit/matrixrtc/MatrixRTCSession.spec.ts
@@ -864,7 +864,64 @@ describe("MatrixRTCSession", () => {
                 expect(client.cancelPendingEvent).toHaveBeenCalledWith(eventSentinel);
             });
 
-            it("Re-sends key if a new member joins", async () => {
+            it("re-sends key if a new member joins even if a key rotation is in progress", async () => {
+                jest.useFakeTimers();
+                try {
+                    // session with two members
+                    const member2 = Object.assign({}, membershipTemplate, {
+                        device_id: "BBBBBBB",
+                    });
+                    const mockRoom = makeMockRoom([membershipTemplate, member2]);
+                    sess = MatrixRTCSession.roomSessionForRoom(client, mockRoom);
+
+                    // joining will trigger an initial key send
+                    const keysSentPromise1 = new Promise<EncryptionKeysEventContent>((resolve) => {
+                        sendEventMock.mockImplementation((_roomId, _evType, payload) => resolve(payload));
+                    });
+                    sess.joinRoomSession([mockFocus], mockFocus, {
+                        manageMediaKeys: true,
+                        updateEncryptionKeyThrottle: 1000,
+                        makeKeyDelay: 3000,
+                    });
+                    await keysSentPromise1;
+                    expect(sess!.statistics.counters.roomEventEncryptionKeysSent).toEqual(1);
+
+                    // member2 leaves triggering key rotation
+                    mockRoom.getLiveTimeline().getState = jest
+                        .fn()
+                        .mockReturnValue(makeMockRoomState([membershipTemplate], mockRoom.roomId));
+                    sess.onMembershipUpdate();
+
+                    // member2 re-joins which should trigger an immediate re-send
+                    const keysSentPromise2 = new Promise<EncryptionKeysEventContent>((resolve) => {
+                        sendEventMock.mockImplementation((_roomId, _evType, payload) => resolve(payload));
+                    });
+                    mockRoom.getLiveTimeline().getState = jest
+                        .fn()
+                        .mockReturnValue(makeMockRoomState([membershipTemplate, member2], mockRoom.roomId));
+                    sess.onMembershipUpdate();
+                    // but, that immediate resend is throttled so we need to wait a bit
+                    jest.advanceTimersByTime(1000);
+                    const { keys } = await keysSentPromise2;
+                    expect(sess!.statistics.counters.roomEventEncryptionKeysSent).toEqual(2);
+                    // key index should still be the original: 0
+                    expect(keys[0].index).toEqual(0);
+
+                    // check that the key rotation actually happens
+                    const keysSentPromise3 = new Promise<EncryptionKeysEventContent>((resolve) => {
+                        sendEventMock.mockImplementation((_roomId, _evType, payload) => resolve(payload));
+                    });
+                    jest.advanceTimersByTime(2000);
+                    const { keys: rotatedKeys } = await keysSentPromise3;
+                    expect(sess!.statistics.counters.roomEventEncryptionKeysSent).toEqual(3);
+                    // key index should now be the rotated one: 1
+                    expect(rotatedKeys[0].index).toEqual(1);
+                } finally {
+                    jest.useRealTimers();
+                }
+            });
+
+            it("re-sends key if a new member joins", async () => {
                 jest.useFakeTimers();
                 try {
                     const mockRoom = makeMockRoom([membershipTemplate]);
diff --git a/src/matrixrtc/MatrixRTCSession.ts b/src/matrixrtc/MatrixRTCSession.ts
index 046591ede17..147a6e0a05e 100644
--- a/src/matrixrtc/MatrixRTCSession.ts
+++ b/src/matrixrtc/MatrixRTCSession.ts
@@ -878,7 +878,7 @@ export class MatrixRTCSession extends TypedEventEmitter<MatrixRTCSessionEvent, M
             }
         }
 
-        if (this.manageMediaKeys && this.isJoined() && this.makeNewKeyTimeout === undefined) {
+        if (this.manageMediaKeys && this.isJoined()) {
             const oldMembershipIds = new Set(
                 oldMemberships.filter((m) => !this.isMyMembership(m)).map(getParticipantIdFromMembership),
             );
@@ -896,8 +896,12 @@ export class MatrixRTCSession extends TypedEventEmitter<MatrixRTCSessionEvent, M
             this.storeLastMembershipFingerprints();
 
             if (anyLeft) {
-                logger.debug(`Member(s) have left: queueing sender key rotation`);
-                this.makeNewKeyTimeout = setTimeout(this.onRotateKeyTimeout, this.makeKeyDelay);
+                if (this.makeNewKeyTimeout) {
+                    // existing rotation in progress, so let it complete
+                } else {
+                    logger.debug(`Member(s) have left: queueing sender key rotation`);
+                    this.makeNewKeyTimeout = setTimeout(this.onRotateKeyTimeout, this.makeKeyDelay);
+                }
             } else if (anyJoined) {
                 logger.debug(`New member(s) have joined: re-sending keys`);
                 this.requestSendCurrentKey();