Add CryptoApi.encryptToDeviceMessages() and deprecate Crypto.encryptAndSendToDevices()

package.json

@@ -50,7 +50,7 @@
     ],
     "dependencies": {
         "@babel/runtime": "^7.12.5",
-        "@matrix-org/matrix-sdk-crypto-wasm": "^7.0.0",
+        "@matrix-org/matrix-sdk-crypto-wasm": "github:matrix-org/matrix-rust-sdk-crypto-wasm#9a167f7ca220cfb7e192d5312e3159dcce5391d4",
         "@matrix-org/olm": "3.2.15",
         "another-json": "^0.2.0",
         "bs58": "^6.0.0",

src/client.ts

@@ -3209,20 +3209,43 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
      * Encrypts and sends a given object via Olm to-device messages to a given
      * set of devices.
      *
-     * @param userDeviceInfoArr - list of deviceInfo objects representing the devices to send to
+     * @param userDeviceInfoArr - @deprecated list of deviceInfo objects representing the devices to send to
+     * @param devices - list of (userId, deviceId) pairs objects representing the devices to send to
      *
      * @param payload - fields to include in the encrypted payload
      *
-     * @returns Promise which
-     *     resolves once the message has been encrypted and sent to the given
-     *     userDeviceMap, and returns the `{ contentMap, deviceInfoByDeviceId }`
-     *     of the successfully sent messages.
+     * @returns Promise which resolves once the payload has been encrypted and sent to the given devices
      */
-    public encryptAndSendToDevices(userDeviceInfoArr: IOlmDevice<DeviceInfo>[], payload: object): Promise<void> {
-        if (!this.crypto) {
+    public encryptAndSendToDevices(userDeviceInfoArr: IOlmDevice<DeviceInfo>[], payload: object): Promise<void>;
+    public encryptAndSendToDevices(devices: { userId: string; deviceId: string }[], payload: object): Promise<void>;
+    public async encryptAndSendToDevices(
+        devices: { userId: string; deviceId: string }[] | IOlmDevice<DeviceInfo>[],
+        payload: object,
+    ): Promise<void> {
+        if (devices.length == 0) {
+            return;
+        }
+
+        if ("deviceinfo" in devices[0]) {
+            if (!this.crypto) {
+                throw new Error("End-to-End encryption disabled");
+            }
+            return this.crypto.encryptAndSendToDevices(devices as IOlmDevice<DeviceInfo>[], payload);
+        }
+
+        const crypto = this.getCrypto();
+        if (!crypto) {
             throw new Error("End-to-End encryption disabled");
         }
-        return this.crypto.encryptAndSendToDevices(userDeviceInfoArr, payload);
+
+        const { type } = payload as { type: string };
+
+        const batch = await crypto.encryptToDeviceMessages(
+            type,
+            devices as { userId: string; deviceId: string }[],
+            payload,
+        );
+        await this.queueToDevice(batch);
     }
 
     /**

src/crypto-api/index.ts

@@ -16,6 +16,7 @@ limitations under the License.
 
 import type { SecretsBundle } from "@matrix-org/matrix-sdk-crypto-wasm";
 import type { IMegolmSessionData } from "../@types/crypto.ts";
+import type { ToDeviceBatch, ToDevicePayload } from "../models/ToDeviceMessage.ts";
 import { Room } from "../models/room.ts";
 import { DeviceMap } from "../models/device.ts";
 import { UIAuthCallback } from "../interactive-auth.ts";
@@ -550,6 +551,22 @@ export interface CryptoApi {
      * @param secrets - The secrets bundle received from the other device
      */
     importSecretsBundle?(secrets: Awaited<ReturnType<SecretsBundle["to_json"]>>): Promise<void>;
+
+    /**
+     * Encrypts a given payload object via Olm to-device messages to a given
+     * set of devices.
+     *
+     * @param eventType the type of the event to send
+     * @param devices an array of (user ID, device ID) pairs to encrypt the payload for
+     * @param payload the payload to encrypt
+     *
+     * @returns a promise which resolves to the batch of encrypted payloads which can then be sent via {@link MatrixClient#queueToDevice}
+     */
+    encryptToDeviceMessages(
+        eventType: string,
+        devices: { userId: string; deviceId: string }[],
+        payload: ToDevicePayload,
+    ): Promise<ToDeviceBatch>;
 }
 
 /** A reason code for a failure to decrypt an event. */

src/crypto/index.ts

@@ -55,7 +55,7 @@ import { IStore } from "../store/index.ts";
 import { Room, RoomEvent } from "../models/room.ts";
 import { RoomMember, RoomMemberEvent } from "../models/room-member.ts";
 import { EventStatus, IContent, IEvent, MatrixEvent, MatrixEventEvent } from "../models/event.ts";
-import { ToDeviceBatch } from "../models/ToDeviceMessage.ts";
+import { ToDeviceBatch, ToDevicePayload } from "../models/ToDeviceMessage.ts";
 import { ClientEvent, IKeysUploadResponse, ISignedKey, IUploadKeySignaturesResponse, MatrixClient } from "../client.ts";
 import { IRoomEncryption, RoomList } from "./RoomList.ts";
 import { IKeyBackupInfo } from "./keybackup.ts";
@@ -3522,60 +3522,17 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
      *     resolves once the message has been encrypted and sent to the given
      *     userDeviceMap, and returns the `{ contentMap, deviceInfoByDeviceId }`
      *     of the successfully sent messages.
+     *
+     * @deprecated Instead use {@link encryptToDeviceMessages} followed by {@link MatrixClient.queueToDevice}.
      */
     public async encryptAndSendToDevices(userDeviceInfoArr: IOlmDevice<DeviceInfo>[], payload: object): Promise<void> {
-        const toDeviceBatch: ToDeviceBatch = {
-            eventType: EventType.RoomMessageEncrypted,
-            batch: [],
-        };
-
         try {
-            await Promise.all(
-                userDeviceInfoArr.map(async ({ userId, deviceInfo }) => {
-                    const deviceId = deviceInfo.deviceId;
-                    const encryptedContent: IEncryptedContent = {
-                        algorithm: olmlib.OLM_ALGORITHM,
-                        sender_key: this.olmDevice.deviceCurve25519Key!,
-                        ciphertext: {},
-                        [ToDeviceMessageId]: uuidv4(),
-                    };
-
-                    toDeviceBatch.batch.push({
-                        userId,
-                        deviceId,
-                        payload: encryptedContent,
-                    });
-
-                    await olmlib.ensureOlmSessionsForDevices(
-                        this.olmDevice,
-                        this.baseApis,
-                        new Map([[userId, [deviceInfo]]]),
-                    );
-                    await olmlib.encryptMessageForDevice(
-                        encryptedContent.ciphertext,
-                        this.userId,
-                        this.deviceId,
-                        this.olmDevice,
-                        userId,
-                        deviceInfo,
-                        payload,
-                    );
-                }),
+            const toDeviceBatch = await this.prepareToDeviceBatch(
+                EventType.RoomMessageEncrypted,
+                userDeviceInfoArr,
+                payload,
             );
 
-            // prune out any devices that encryptMessageForDevice could not encrypt for,
-            // in which case it will have just not added anything to the ciphertext object.
-            // There's no point sending messages to devices if we couldn't encrypt to them,
-            // since that's effectively a blank message.
-            toDeviceBatch.batch = toDeviceBatch.batch.filter((msg) => {
-                if (Object.keys(msg.payload.ciphertext).length > 0) {
-                    return true;
-                } else {
-                    logger.log(`No ciphertext for device ${msg.userId}:${msg.deviceId}: pruning`);
-                    return false;
-                }
-            });
-
             try {
                 await this.baseApis.queueToDevice(toDeviceBatch);
             } catch (e) {
@@ -4305,6 +4262,93 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
     public async startDehydration(createNewKey?: boolean): Promise<void> {
         throw new Error("Not implemented");
     }
+
+    private async prepareToDeviceBatch(
+        eventType: string,
+        userDeviceInfoArr: IOlmDevice<DeviceInfo>[],
+        payload: object,
+    ): Promise<ToDeviceBatch> {
+        const toDeviceBatch: ToDeviceBatch = {
+            eventType,
+            batch: [],
+        };
+
+        await Promise.all(
+            userDeviceInfoArr.map(async ({ userId, deviceInfo }) => {
+                const deviceId = deviceInfo.deviceId;
+                const encryptedContent: IEncryptedContent = {
+                    algorithm: olmlib.OLM_ALGORITHM,
+                    sender_key: this.olmDevice.deviceCurve25519Key!,
+                    ciphertext: {},
+                    [ToDeviceMessageId]: uuidv4(),
+                };
+
+                toDeviceBatch.batch.push({
+                    userId,
+                    deviceId,
+                    payload: encryptedContent,
+                });
+
+                await olmlib.ensureOlmSessionsForDevices(
+                    this.olmDevice,
+                    this.baseApis,
+                    new Map([[userId, [deviceInfo]]]),
+                );
+                await olmlib.encryptMessageForDevice(
+                    encryptedContent.ciphertext,
+                    this.userId,
+                    this.deviceId,
+                    this.olmDevice,
+                    userId,
+                    deviceInfo,
+                    payload,
+                );
+            }),
+        );
+
+        // prune out any devices that encryptMessageForDevice could not encrypt for,
+        // in which case it will have just not added anything to the ciphertext object.
+        // There's no point sending messages to devices if we couldn't encrypt to them,
+        // since that's effectively a blank message.
+        toDeviceBatch.batch = toDeviceBatch.batch.filter((msg) => {
+            if (Object.keys(msg.payload.ciphertext).length > 0) {
+                return true;
+            } else {
+                logger.log(`No ciphertext for device ${msg.userId}:${msg.deviceId}: pruning`);
+                return false;
+            }
+        });
+
+        return toDeviceBatch;
+    }
+
+    public async encryptToDeviceMessages(
+        eventType: string,
+        devices: { userId: string; deviceId: string }[],
+        payload: ToDevicePayload,
+    ): Promise<ToDeviceBatch> {
+        const userIds = new Set(devices.map(({ userId }) => userId));
+        const deviceInfoMap = await this.downloadKeys(Array.from(userIds), false);
+
+        const userDeviceInfoArr: IOlmDevice<DeviceInfo>[] = [];
+
+        devices.forEach(({ userId, deviceId }) => {
+            const devices = deviceInfoMap.get(userId);
+            if (!devices) {
+                logger.warn(`No devices found for user ${userId}`);
+                return;
+            }
+
+            if (devices.has(deviceId)) {
+                // Send the message to a specific device
+                userDeviceInfoArr.push({ userId, deviceInfo: devices.get(deviceId)! });
+            } else {
+                logger.warn(`No device found for user ${userId} with id ${deviceId}`);
+            }
+        });
+
+        return this.prepareToDeviceBatch(eventType, userDeviceInfoArr, payload);
+    }
 }
 
 /**

src/embedded.ts

@@ -401,14 +401,30 @@ export class RoomWidgetClient extends MatrixClient {
         await this.widgetApi.sendToDevice(eventType, false, recursiveMapToObject(contentMap));
     }
 
-    public async encryptAndSendToDevices(userDeviceInfoArr: IOlmDevice<DeviceInfo>[], payload: object): Promise<void> {
+    public async encryptAndSendToDevices(
+        devices: { userId: string; deviceId: string }[] | IOlmDevice<DeviceInfo>[],
+        payload: object,
+    ): Promise<void> {
         // map: user Id → device Id → payload
         const contentMap: MapWithDefault<string, Map<string, object>> = new MapWithDefault(() => new Map());
-        for (const {
-            userId,
-            deviceInfo: { deviceId },
-        } of userDeviceInfoArr) {
-            contentMap.getOrCreate(userId).set(deviceId, payload);
+
+        if (devices.length == 0) {
+            return;
+        }
+
+        if ("deviceinfo" in devices[0]) {
+            // pre-CryptoApi style:
+            for (const {
+                userId,
+                deviceInfo: { deviceId },
+            } of devices as IOlmDevice<DeviceInfo>[]) {
+                contentMap.getOrCreate(userId).set(deviceId, payload);
+            }
+        } else {
+            // new CryptoApi style:
+            for (const { userId, deviceId } of devices as { userId: string; deviceId: string }[]) {
+                contentMap.getOrCreate(userId).set(deviceId, payload);
+            }
         }
 
         await this.widgetApi.sendToDevice((payload as { type: string }).type, true, recursiveMapToObject(contentMap));

src/rust-crypto/rust-crypto.ts

@@ -21,6 +21,7 @@ import type { IEventDecryptionResult, IMegolmSessionData } from "../@types/crypt
 import { KnownMembership } from "../@types/membership.ts";
 import type { IDeviceLists, IToDeviceEvent } from "../sync-accumulator.ts";
 import type { IEncryptedEventInfo } from "../crypto/api.ts";
+import type { ToDevicePayload, ToDeviceBatch } from "../models/ToDeviceMessage.ts";
 import { MatrixEvent, MatrixEventEvent } from "../models/event.ts";
 import { Room } from "../models/room.ts";
 import { RoomMember } from "../models/room-member.ts";
@@ -1713,6 +1714,37 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, RustCryptoEv
     public async getOwnIdentity(): Promise<RustSdkCryptoJs.OwnUserIdentity | undefined> {
         return await this.olmMachine.getIdentity(new RustSdkCryptoJs.UserId(this.userId));
     }
+
+    public async encryptToDeviceMessages(
+        eventType: string,
+        devices: { userId: string; deviceId: string }[],
+        payload: ToDevicePayload,
+    ): Promise<ToDeviceBatch> {
+        const batch: ToDeviceBatch = {
+            batch: [],
+            eventType,
+        };
+
+        for (const { userId, deviceId } of devices) {
+            const device: RustSdkCryptoJs.Device | undefined = await this.olmMachine.getDevice(
+                new RustSdkCryptoJs.UserId(userId),
+                new RustSdkCryptoJs.DeviceId(deviceId),
+            );
+
+            if (device) {
+                const encryptedPayload = JSON.parse(await device.encryptToDeviceEvent(eventType, payload));
+                batch.batch.push({
+                    deviceId,
+                    userId,
+                    payload: encryptedPayload,
+                });
+            } else {
+                this.logger.warn(`encryptToDeviceMessages: unknown device ${userId}:${deviceId}`);
+            }
+        }
+
+        return batch;
+    }
 }
 
 class EventDecryptor {