From 38a6be5465afa40264a20322ad3a0090af1f8ead Mon Sep 17 00:00:00 2001
From: Hugh Nimmo-Smith <hughns@element.io>
Date: Mon, 28 Oct 2024 16:09:55 +0000
Subject: [PATCH] Use new CryptoApi.encryptToDeviceMessages() to send encrypted
 to-device messages from widgets

---
 src/stores/widgets/StopGapWidgetDriver.ts     | 54 +++++++++++----
 .../widgets/StopGapWidgetDriver-test.ts       | 68 ++++++++++++++++++-
 2 files changed, 106 insertions(+), 16 deletions(-)

diff --git a/src/stores/widgets/StopGapWidgetDriver.ts b/src/stores/widgets/StopGapWidgetDriver.ts
index b5609dc5fb9..07fca154e84 100644
--- a/src/stores/widgets/StopGapWidgetDriver.ts
+++ b/src/stores/widgets/StopGapWidgetDriver.ts
@@ -416,26 +416,54 @@ export class StopGapWidgetDriver extends WidgetDriver {
 
     /**
      * Implements {@link WidgetDriver#sendToDevice}
-     * Encrypted to-device events are not supported.
      */
     public async sendToDevice(
         eventType: string,
         encrypted: boolean,
         contentMap: { [userId: string]: { [deviceId: string]: object } },
     ): Promise<void> {
-        if (encrypted) throw new Error("Encrypted to-device events are not supported");
-
         const client = MatrixClientPeg.safeGet();
-        await client.queueToDevice({
-            eventType,
-            batch: Object.entries(contentMap).flatMap(([userId, userContentMap]) =>
-                Object.entries(userContentMap).map(([deviceId, content]) => ({
-                    userId,
-                    deviceId,
-                    payload: content,
-                })),
-            ),
-        });
+
+        if (encrypted) {
+            const crypto = client.getCrypto();
+            if (!crypto) throw new Error("E2EE not enabled");
+
+            // attempt to re-batch these up into a single request
+            const invertedContentMap: { [content: string]: { userId: string; deviceId: string }[] } = {};
+
+            for (const userId of Object.keys(contentMap)) {
+                const userContentMap = contentMap[userId];
+                for (const deviceId of Object.keys(userContentMap)) {
+                    const content = userContentMap[deviceId];
+                    const stringifiedContent = JSON.stringify(content);
+                    invertedContentMap[stringifiedContent] = invertedContentMap[stringifiedContent] || [];
+                    invertedContentMap[stringifiedContent].push({ userId, deviceId });
+                }
+            }
+
+            await Promise.all(
+                Object.entries(invertedContentMap).map(async ([stringifiedContent, recipients]) => {
+                    const batch = await crypto.encryptToDeviceMessages(
+                        eventType,
+                        recipients,
+                        JSON.parse(stringifiedContent),
+                    );
+
+                    await client.queueToDevice(batch);
+                }),
+            );
+        } else {
+            await client.queueToDevice({
+                eventType,
+                batch: Object.entries(contentMap).flatMap(([userId, userContentMap]) =>
+                    Object.entries(userContentMap).map(([deviceId, content]) => ({
+                        userId,
+                        deviceId,
+                        payload: content,
+                    })),
+                ),
+            });
+        }
     }
 
     private pickRooms(roomIds?: (string | Symbols.AnyRoom)[]): Room[] {
diff --git a/test/unit-tests/stores/widgets/StopGapWidgetDriver-test.ts b/test/unit-tests/stores/widgets/StopGapWidgetDriver-test.ts
index 961b702452d..f5c0c5fe230 100644
--- a/test/unit-tests/stores/widgets/StopGapWidgetDriver-test.ts
+++ b/test/unit-tests/stores/widgets/StopGapWidgetDriver-test.ts
@@ -201,10 +201,72 @@ describe("StopGapWidgetDriver", () => {
             });
         });
 
-        it("raises an error if encrypted", async () => {
-            await expect(driver.sendToDevice("org.example.foo", true, contentMap)).rejects.toThrow(
-                "Encrypted to-device events are not supported",
+        it("sends encrypted messages", async () => {
+            const encryptToDeviceMessages = jest
+                .fn()
+                .mockImplementation(
+                    (eventType, recipients: { userId: string; deviceId: string }[], content: object) => ({
+                        eventType: "m.room.encrypted",
+                        batch: recipients.map(({ userId, deviceId }) => ({
+                            userId,
+                            deviceId,
+                            payload: {
+                                eventType,
+                                content,
+                            },
+                        })),
+                    }),
+                );
+
+            MatrixClientPeg.safeGet().getCrypto()!.encryptToDeviceMessages = encryptToDeviceMessages;
+
+            await driver.sendToDevice("org.example.foo", true, {
+                "@alice:example.org": {
+                    aliceMobile: {
+                        hello: "alice",
+                    },
+                },
+                "@bob:example.org": {
+                    bobDesktop: {
+                        hello: "bob",
+                    },
+                },
+            });
+
+            expect(encryptToDeviceMessages).toHaveBeenCalledWith(
+                "org.example.foo",
+                [{ deviceId: "aliceMobile", userId: "@alice:example.org" }],
+                {
+                    hello: "alice",
+                },
             );
+            expect(encryptToDeviceMessages).toHaveBeenCalledWith(
+                "org.example.foo",
+                [{ deviceId: "bobDesktop", userId: "@bob:example.org" }],
+                {
+                    hello: "bob",
+                },
+            );
+            expect(client.queueToDevice).toHaveBeenCalledWith({
+                eventType: "m.room.encrypted",
+                batch: expect.arrayContaining([
+                    {
+                        deviceId: "aliceMobile",
+                        payload: { content: { hello: "alice" }, eventType: "org.example.foo" },
+                        userId: "@alice:example.org",
+                    },
+                ]),
+            });
+            expect(client.queueToDevice).toHaveBeenCalledWith({
+                eventType: "m.room.encrypted",
+                batch: expect.arrayContaining([
+                    {
+                        deviceId: "bobDesktop",
+                        payload: { content: { hello: "bob" }, eventType: "org.example.foo" },
+                        userId: "@bob:example.org",
+                    },
+                ]),
+            });
         });
     });