MatrixRTC key distribution using to-device messaging

hughns · hughns · commit 5897988a4fc0 · 2024-10-30T10:02:00.000Z
diff --git a/src/matrixrtc/CallMembership.ts b/src/matrixrtc/CallMembership.ts
@@ -22,6 +22,7 @@ import { Focus } from "./focus.ts";
 import { isLivekitFocusActive } from "./LivekitFocus.ts";
 
 type CallScope = "m.room" | "m.user";
+
 // Represents an entry in the memberships section of an m.call.member event as it is on the wire
 
 // There are two different data interfaces. One for the Legacy types and one compliant with MSC4143
@@ -39,6 +40,8 @@ export type SessionMembershipData = {
 
     // Application specific data
     scope?: CallScope;
+
+    key_distribution?: KeyDistributionMechanism;
 };
 
 export const isSessionMembershipData = (data: CallMembershipData): data is SessionMembershipData =>
@@ -69,6 +72,7 @@ export type CallMembershipDataLegacy = {
     membershipID: string;
     created_ts?: number;
     foci_active?: Focus[];
+    key_distribution?: KeyDistributionMechanism;
 } & EitherAnd<{ expires: number }, { expires_ts: number }>;
 
 export const isLegacyCallMembershipData = (data: CallMembershipData): data is CallMembershipDataLegacy =>
@@ -103,6 +107,8 @@ const checkCallMembershipDataLegacy = (data: any, errors: string[]): data is Cal
 
 export type CallMembershipData = CallMembershipDataLegacy | SessionMembershipData;
 
+type KeyDistributionMechanism = "room_event" | "to_device";
+
 export class CallMembership {
     public static equal(a: CallMembership, b: CallMembership): boolean {
         return deepCompare(a.membershipData, b.membershipData);
@@ -244,4 +250,8 @@ export class CallMembership {
             }
         }
     }
+
+    public get keyDistributionMethod(): KeyDistributionMechanism {
+        return this.membershipData.key_distribution ?? "room_event";
+    }
 }
diff --git a/src/matrixrtc/MatrixRTCSession.ts b/src/matrixrtc/MatrixRTCSession.ts
@@ -18,7 +18,7 @@ import { logger as rootLogger } from "../logger.ts";
 import { TypedEventEmitter } from "../models/typed-event-emitter.ts";
 import { EventTimeline } from "../models/event-timeline.ts";
 import { Room } from "../models/room.ts";
-import { MatrixClient } from "../client.ts";
+import { MatrixClient, SendToDeviceContentMap } from "../client.ts";
 import { EventType } from "../@types/event.ts";
 import { UpdateDelayedEventAction } from "../@types/requests.ts";
 import {
@@ -31,14 +31,15 @@ import {
 import { RoomStateEvent } from "../models/room-state.ts";
 import { Focus } from "./focus.ts";
 import { randomString, secureRandomBase64Url } from "../randomstring.ts";
-import { EncryptionKeysEventContent } from "./types.ts";
+import { EncryptionKeysEventContent, EncryptionKeysToDeviceContent } from "./types.ts";
 import { decodeBase64, encodeUnpaddedBase64 } from "../base64.ts";
 import { KnownMembership } from "../@types/membership.ts";
 import { MatrixError } from "../http-api/errors.ts";
 import { MatrixEvent } from "../models/event.ts";
 import { isLivekitFocusActive } from "./LivekitFocus.ts";
 import { ExperimentalGroupCallRoomMemberState } from "../webrtc/groupCall.ts";
 import { sleep } from "../utils.ts";
+import type { RoomWidgetClient } from "../embedded.ts";
 
 const logger = rootLogger.getChild("MatrixRTCSession");
 
@@ -162,8 +163,21 @@ export class MatrixRTCSession extends TypedEventEmitter<MatrixRTCSessionEvent, M
              * The number of times we have received a room event containing encryption keys.
              */
             roomEventEncryptionKeysReceived: 0,
+            /**
+             * The number of times we have sent a to-device event containing encryption keys.
+             */
+            toDeviceEncryptionKeysSent: 0,
+            /**
+             * The number of times we have received a to-device event containing encryption keys.
+             */
+            toDeviceEncryptionKeysReceived: 0,
         },
         totals: {
+            /**
+             * The total age (in milliseconds) of all to-device events containing encryption keys that we have received.
+             * We track the total age so that we can later calculate the average age of all keys received.
+             */
+            toDeviceEncryptionKeysReceivedTotalAge: 0,
             /**
              * The total age (in milliseconds) of all room events containing encryption keys that we have received.
              * We track the total age so that we can later calculate the average age of all keys received.
@@ -546,7 +560,7 @@ export class MatrixRTCSession extends TypedEventEmitter<MatrixRTCSessionEvent, M
     }
 
     /**
-     * Requests that we resend our current keys to the room. May send a keys event immediately
+     * Requests that we (re)-send our current keys to the room. May send a keys event immediately
      * or queue for alter if one has already been sent recently.
      */
     private requestSendCurrentKey(): void {
@@ -602,21 +616,10 @@ export class MatrixRTCSession extends TypedEventEmitter<MatrixRTCSessionEvent, M
         const keyToSend = myKeys[keyIndexToSend];
 
         try {
-            const content: EncryptionKeysEventContent = {
-                keys: [
-                    {
-                        index: keyIndexToSend,
-                        key: encodeUnpaddedBase64(keyToSend),
-                    },
-                ],
-                device_id: deviceId,
-                call_id: "",
-                sent_ts: Date.now(),
-            };
-
-            this.statistics.counters.roomEventEncryptionKeysSent += 1;
-
-            await this.client.sendEvent(this.room.roomId, EventType.CallEncryptionKeysPrefix, content);
+            await Promise.all([
+                this.sendKeysViaRoomEvent(deviceId, keyToSend, keyIndexToSend),
+                this.sendKeysViaToDevice(deviceId, keyToSend, keyIndexToSend),
+            ]);
 
             logger.debug(
                 `Embedded-E2EE-LOG updateEncryptionKeyEvent participantId=${userId}:${deviceId} numKeys=${myKeys.length} currentKeyIndex=${this.currentEncryptionKeyIndex} keyIndexToSend=${keyIndexToSend}`,
@@ -639,6 +642,96 @@ export class MatrixRTCSession extends TypedEventEmitter<MatrixRTCSessionEvent, M
         }
     };
 
+    private async sendKeysViaRoomEvent(deviceId: string, key: Uint8Array, index: number): Promise<void> {
+        const membersRequiringRoomEvent = this.memberships.filter(
+            (m) => !this.isMyMembership(m) && m.keyDistributionMethod === "room_event",
+        );
+
+        if (membersRequiringRoomEvent.length === 0) {
+            logger.info("No members require keys via room event");
+            return;
+        }
+
+        logger.info(
+            `Sending encryption keys event for: ${membersRequiringRoomEvent.map((m) => `${m.sender}:${m.deviceId}`).join(", ")}`,
+        );
+
+        const content: EncryptionKeysEventContent = {
+            keys: [
+                {
+                    index,
+                    key: encodeUnpaddedBase64(key),
+                },
+            ],
+            device_id: deviceId,
+            call_id: "",
+            sent_ts: Date.now(),
+        };
+
+        this.statistics.counters.roomEventEncryptionKeysSent += 1;
+
+        await this.client.sendEvent(this.room.roomId, EventType.CallEncryptionKeysPrefix, content);
+    }
+
+    private async sendKeysViaToDevice(deviceId: string, key: Uint8Array, index: number): Promise<void> {
+        const membershipsRequiringToDevice = this.memberships.filter(
+            (m) => !this.isMyMembership(m) && m.sender && m.keyDistributionMethod === "to_device",
+        );
+
+        if (membershipsRequiringToDevice.length === 0) {
+            logger.info("No members require keys via to-device event");
+            return;
+        }
+
+        const content: EncryptionKeysToDeviceContent = {
+            keys: [{ index, key: encodeUnpaddedBase64(key) }],
+            device_id: deviceId,
+            call_id: "",
+            room_id: this.room.roomId,
+            sent_ts: Date.now(),
+        };
+
+        logger.info(
+            `Sending encryption keys to-device batch for: ${membershipsRequiringToDevice.map(({ sender, deviceId }) => `${sender}:${deviceId}`).join(", ")}`,
+        );
+
+        this.statistics.counters.toDeviceEncryptionKeysSent += membershipsRequiringToDevice.length;
+
+        // we don't do an instanceof due to circular dependency issues
+        if ("widgetApi" in this.client) {
+            logger.info("Sending keys via widgetApi");
+            // embedded mode, getCrypto() returns null and so we make some assumptions about the underlying implementation
+
+            const contentMap: SendToDeviceContentMap = new Map();
+
+            membershipsRequiringToDevice.forEach(({ sender, deviceId }) => {
+                if (!contentMap.has(sender!)) {
+                    contentMap.set(sender!, new Map());
+                }
+
+                contentMap.get(sender!)!.set(deviceId, content);
+            });
+
+            await (this.client as unknown as RoomWidgetClient).sendToDeviceViaWidgetApi(
+                EventType.CallEncryptionKeysPrefix,
+                true,
+                contentMap,
+            );
+        } else {
+            const crypto = this.client.getCrypto();
+            if (!crypto) {
+                logger.error("No crypto instance available to send keys via to-device event");
+                return;
+            }
+
+            const devices = membershipsRequiringToDevice.map(({ deviceId, sender }) => ({ userId: sender!, deviceId }));
+
+            const batch = await crypto.encryptToDeviceMessages(EventType.CallEncryptionKeysPrefix, devices, content);
+
+            await this.client.queueToDevice(batch);
+        }
+    }
+
     /**
      * Sets a timer for the soonest membership expiry
      */
@@ -714,9 +807,17 @@ export class MatrixRTCSession extends TypedEventEmitter<MatrixRTCSessionEvent, M
             return;
         }
 
-        this.statistics.counters.roomEventEncryptionKeysReceived += 1;
-        const age = Date.now() - (typeof content.sent_ts === "number" ? content.sent_ts : event.getTs());
-        this.statistics.totals.roomEventEncryptionKeysReceivedTotalAge += age;
+        let age: number;
+
+        if (event.getRoomId()) {
+            this.statistics.counters.roomEventEncryptionKeysReceived += 1;
+            age = Date.now() - (typeof content.sent_ts === "number" ? content.sent_ts : event.getTs());
+            this.statistics.totals.roomEventEncryptionKeysReceivedTotalAge += age;
+        } else {
+            this.statistics.counters.toDeviceEncryptionKeysReceived += 1;
+            age = Date.now() - (content as EncryptionKeysToDeviceContent).sent_ts;
+            this.statistics.totals.toDeviceEncryptionKeysReceivedTotalAge += age;
+        }
 
         for (const key of content.keys) {
             if (!key) {
@@ -795,8 +896,8 @@ export class MatrixRTCSession extends TypedEventEmitter<MatrixRTCSessionEvent, M
                 logger.debug(`Member(s) have left: queueing sender key rotation`);
                 this.makeNewKeyTimeout = setTimeout(this.onRotateKeyTimeout, MAKE_KEY_DELAY);
             } else if (anyJoined) {
-                logger.debug(`New member(s) have joined: re-sending keys`);
-                this.requestSendCurrentKey();
+                logger.debug(`New member(s) have joined: rotating keys`);
+                this.makeNewKeyTimeout = setTimeout(this.onRotateKeyTimeout, MAKE_KEY_DELAY);
             } else if (oldFingerprints) {
                 // does it look like any of the members have updated their memberships?
                 const newFingerprints = this.lastMembershipFingerprints!;
@@ -849,6 +950,7 @@ export class MatrixRTCSession extends TypedEventEmitter<MatrixRTCSessionEvent, M
             foci_active: this.ownFociPreferred,
             membershipID: this.membershipId,
             ...(createdTs ? { created_ts: createdTs } : {}),
+            key_distribution: "to_device",
         };
     }
     /**
@@ -862,6 +964,7 @@ export class MatrixRTCSession extends TypedEventEmitter<MatrixRTCSessionEvent, M
             device_id: deviceId,
             focus_active: { type: "livekit", focus_selection: "oldest_membership" },
             foci_preferred: this.ownFociPreferred ?? [],
+            key_distribution: "to_device",
         };
     }
 
diff --git a/src/matrixrtc/MatrixRTCSessionManager.ts b/src/matrixrtc/MatrixRTCSessionManager.ts
@@ -22,6 +22,7 @@ import { RoomState, RoomStateEvent } from "../models/room-state.ts";
 import { MatrixEvent } from "../models/event.ts";
 import { MatrixRTCSession } from "./MatrixRTCSession.ts";
 import { EventType } from "../@types/event.ts";
+import { EncryptionKeysToDeviceContent } from "./types.ts";
 
 const logger = rootLogger.getChild("MatrixRTCSessionManager");
 
@@ -56,7 +57,7 @@ export class MatrixRTCSessionManager extends TypedEventEmitter<MatrixRTCSessionM
 
     public start(): void {
         // We shouldn't need to null-check here, but matrix-client.spec.ts mocks getRooms
-        // returing nothing, and breaks tests if you change it to return an empty array :'(
+        // returning nothing, and breaks tests if you change it to return an empty array :'(
         for (const room of this.client.getRooms() ?? []) {
             const session = MatrixRTCSession.roomSessionForRoom(this.client, room);
             if (session.memberships.length > 0) {
@@ -67,6 +68,7 @@ export class MatrixRTCSessionManager extends TypedEventEmitter<MatrixRTCSessionM
         this.client.on(ClientEvent.Room, this.onRoom);
         this.client.on(RoomEvent.Timeline, this.onTimeline);
         this.client.on(RoomStateEvent.Events, this.onRoomState);
+        this.client.on(ClientEvent.ToDeviceEvent, this.onToDeviceEvent);
     }
 
     public stop(): void {
@@ -78,6 +80,7 @@ export class MatrixRTCSessionManager extends TypedEventEmitter<MatrixRTCSessionM
         this.client.off(ClientEvent.Room, this.onRoom);
         this.client.off(RoomEvent.Timeline, this.onTimeline);
         this.client.off(RoomStateEvent.Events, this.onRoomState);
+        this.client.off(ClientEvent.ToDeviceEvent, this.onToDeviceEvent);
     }
 
     /**
@@ -100,15 +103,40 @@ export class MatrixRTCSessionManager extends TypedEventEmitter<MatrixRTCSessionM
         return this.roomSessions.get(room.roomId)!;
     }
 
-    private async consumeCallEncryptionEvent(event: MatrixEvent, isRetry = false): Promise<void> {
+    private onTimeline = (event: MatrixEvent): void => {
+        this.consumeCallEncryptionEvent(event, (event) => event.getRoomId(), false);
+    };
+
+    private onToDeviceEvent = (event: MatrixEvent): void => {
+        if (!event.isEncrypted()) {
+            logger.warn("Ignoring unencrypted to-device call encryption event", event);
+            return;
+        }
+        this.consumeCallEncryptionEvent(
+            event,
+            (event) => event.getContent<EncryptionKeysToDeviceContent>().room_id,
+            false,
+        );
+    };
+
+    /**
+     * @param event - the event to consume
+     * @param roomIdExtractor - the function to extract the room id from the event
+     * @param isRetry - whether this is a retry. If false we will retry decryption failures once
+     */
+    private consumeCallEncryptionEvent = async (
+        event: MatrixEvent,
+        roomIdExtractor: (event: MatrixEvent) => string | undefined,
+        isRetry: boolean,
+    ): Promise<void> => {
         await this.client.decryptEventIfNeeded(event);
         if (event.isDecryptionFailure()) {
             if (!isRetry) {
                 logger.warn(
                     `Decryption failed for event ${event.getId()}: ${event.decryptionFailureReason} will retry once only`,
                 );
                 // retry after 1 second. After this we give up.
-                setTimeout(() => this.consumeCallEncryptionEvent(event, true), 1000);
+                setTimeout(() => this.consumeCallEncryptionEvent(event, roomIdExtractor, true), 1000);
             } else {
                 logger.warn(`Decryption failed for event ${event.getId()}: ${event.decryptionFailureReason}`);
             }
@@ -117,18 +145,20 @@ export class MatrixRTCSessionManager extends TypedEventEmitter<MatrixRTCSessionM
             logger.info(`Decryption succeeded for event ${event.getId()} after retry`);
         }
 
-        if (event.getType() !== EventType.CallEncryptionKeysPrefix) return Promise.resolve();
+        if (event.getType() !== EventType.CallEncryptionKeysPrefix) return;
+        const roomId = roomIdExtractor(event);
+        if (!roomId) {
+            logger.error("Received call encryption event with no room_id!");
+            return;
+        }
+
+        const room = this.client.getRoom(roomId);
 
-        const room = this.client.getRoom(event.getRoomId());
         if (!room) {
-            logger.error(`Got room state event for unknown room ${event.getRoomId()}!`);
-            return Promise.resolve();
+            logger.error(`Got encryption event for unknown room ${roomId}!`);
+            return;
         }
-
         this.getRoomSession(room).onCallEncryption(event);
-    }
-    private onTimeline = (event: MatrixEvent): void => {
-        this.consumeCallEncryptionEvent(event);
     };
 
     private onRoom = (room: Room): void => {
diff --git a/src/matrixrtc/types.ts b/src/matrixrtc/types.ts
@@ -26,6 +26,11 @@ export interface EncryptionKeysEventContent {
     sent_ts?: number;
 }
 
+export interface EncryptionKeysToDeviceContent extends EncryptionKeysEventContent {
+    room_id?: string;
+    sent_ts: number;
+}
+
 export type CallNotifyType = "ring" | "notify";
 
 export interface ICallNotifyContent {
