Use new CryptoApi.encryptToDeviceMessages() to send encrypted to-devi…

…ce messages from widgets (#28315)
element-hq · Oct 30, 2024 · c23c9df · c23c9df
1 parent 5c45ca5
commit c23c9df
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 16 deletions.
diff --git 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
@@ -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",
+                    },
+                ]),
+            });
         });
     });