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

Fixes #2135.
matrix-org · Apr 15, 2022 · ea6bd5d · ea6bd5d
1 parent 9f45986
commit ea6bd5d
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/src/crypto/algorithms/megolm.ts b/src/crypto/algorithms/megolm.ts
@@ -592,7 +592,7 @@ class MegolmEncryption extends EncryptionAlgorithm {
         payload: IPayload,
     ): Promise<void> {
         const contentMap: Record<string, Record<string, IEncryptedContent>> = {};
-        const deviceInfoByDeviceId = new Map<string, DeviceInfo>();
+        const deviceInfoByDeviceId = new Map<string, Map<string, DeviceInfo>>();
 
         const promises: Promise<unknown>[] = [];
         for (let i = 0; i < userDeviceMap.length; i++) {
@@ -605,7 +605,10 @@ class MegolmEncryption extends EncryptionAlgorithm {
             const userId = val.userId;
             const deviceInfo = val.deviceInfo;
             const deviceId = deviceInfo.deviceId;
-            deviceInfoByDeviceId.set(deviceId, deviceInfo);
+            if (deviceInfoByDeviceId.get(userId) === undefined) {
+                deviceInfoByDeviceId.set(userId, new Map<string, DeviceInfo>());
+            }
+            deviceInfoByDeviceId.get(userId).set(deviceId, deviceInfo);
 
             if (!contentMap[userId]) {
                 contentMap[userId] = {};
@@ -660,7 +663,7 @@ class MegolmEncryption extends EncryptionAlgorithm {
                         session.markSharedWithDevice(
                             userId,
                             deviceId,
-                            deviceInfoByDeviceId.get(deviceId).getIdentityKey(),
+                            deviceInfoByDeviceId.get(userId).get(deviceId).getIdentityKey(),
                             chainIndex,
                         );
                     }