From aa94d5d95c732b7a08287eafbd220424771a6486 Mon Sep 17 00:00:00 2001
From: Jonathan de Jong <jonathan@automatia.nl>
Date: Mon, 6 Jun 2022 16:09:32 +0200
Subject: [PATCH] Assume per-user deviceID uniqueness in
 encryptAndSendKeysToDevices (#2136)

* Segment recorded device info by user ID when tracking key shares.

Fixes #2135.

* address review feedback

* fix userIdDeviceInfo

Co-authored-by: Denis Kasak <dkasak@termina.org.uk>
Co-authored-by: Michael Telatynski <7t3chguy@gmail.com>
---
 src/crypto/algorithms/megolm.ts | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/src/crypto/algorithms/megolm.ts b/src/crypto/algorithms/megolm.ts
index b7bb1165a..8bf5a7c62 100644
--- a/src/crypto/algorithms/megolm.ts
+++ b/src/crypto/algorithms/megolm.ts
@@ -592,7 +592,8 @@ class MegolmEncryption extends EncryptionAlgorithm {
         payload: IPayload,
     ): Promise<void> {
         const contentMap: Record<string, Record<string, IEncryptedContent>> = {};
-        const deviceInfoByDeviceId = new Map<string, DeviceInfo>();
+        // Map from userId to a map of deviceId to deviceInfo
+        const deviceInfoByUserIdAndDeviceId = new Map<string, Map<string, DeviceInfo>>();
 
         const promises: Promise<unknown>[] = [];
         for (let i = 0; i < userDeviceMap.length; i++) {
@@ -605,7 +606,18 @@ class MegolmEncryption extends EncryptionAlgorithm {
             const userId = val.userId;
             const deviceInfo = val.deviceInfo;
             const deviceId = deviceInfo.deviceId;
-            deviceInfoByDeviceId.set(deviceId, deviceInfo);
+
+            // Assign to temp value to make type-checking happy
+            let userIdDeviceInfo = deviceInfoByUserIdAndDeviceId.get(userId);
+
+            if (userIdDeviceInfo === undefined) {
+                userIdDeviceInfo = new Map<string, DeviceInfo>();
+
+                deviceInfoByUserIdAndDeviceId.set(userId, userIdDeviceInfo);
+            }
+
+            // We hold by reference, this updates deviceInfoByUserIdAndDeviceId[userId]
+            userIdDeviceInfo.set(deviceId, deviceInfo);
 
             if (!contentMap[userId]) {
                 contentMap[userId] = {};
@@ -660,7 +672,7 @@ class MegolmEncryption extends EncryptionAlgorithm {
                         session.markSharedWithDevice(
                             userId,
                             deviceId,
-                            deviceInfoByDeviceId.get(deviceId).getIdentityKey(),
+                            deviceInfoByUserIdAndDeviceId.get(userId).get(deviceId).getIdentityKey(),
                             chainIndex,
                         );
                     }