From ea6bd5dcb8be48f3024f9ad3f1e44072ea542b92 Mon Sep 17 00:00:00 2001
From: Denis Kasak <dkasak@termina.org.uk>
Date: Fri, 15 Apr 2022 12:31:48 +0200
Subject: [PATCH] Segment recorded device info by user ID when tracking key
 shares.

Fixes #2135.
---
 src/crypto/algorithms/megolm.ts | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/crypto/algorithms/megolm.ts b/src/crypto/algorithms/megolm.ts
index f960dd4f15e..568f4e0cf31 100644
--- 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,
                         );
                     }