From 9c13b2b4a43b71311261d820e933bcd995590458 Mon Sep 17 00:00:00 2001
From: Bruno Windels <274386+bwindels@users.noreply.github.com>
Date: Fri, 28 Oct 2022 15:41:18 +0200
Subject: [PATCH] add method to fetch missing sender keys

---
 src/matrix/e2ee/DecryptionResult.ts           |  5 ++--
 src/matrix/e2ee/RoomEncryption.js             | 30 +++++++++++++++++++
 .../megolm/decryption/SessionDecryption.ts    |  2 +-
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/matrix/e2ee/DecryptionResult.ts b/src/matrix/e2ee/DecryptionResult.ts
index 0bb2e926..8846616c 100644
--- a/src/matrix/e2ee/DecryptionResult.ts
+++ b/src/matrix/e2ee/DecryptionResult.ts
@@ -27,6 +27,7 @@ limitations under the License.
  */
 
 import type {DeviceIdentity} from "../storage/idb/stores/DeviceIdentityStore";
+import type {TimelineEvent} from "../storage/types";
 
 type DecryptedEvent = {
     type?: string,
@@ -35,12 +36,12 @@ type DecryptedEvent = {
 
 export class DecryptionResult {
     private device?: DeviceIdentity;
-    private roomTracked: boolean = true;
 
     constructor(
         public readonly event: DecryptedEvent,
         public readonly senderCurve25519Key: string,
-        public readonly claimedEd25519Key: string
+        public readonly claimedEd25519Key: string,
+        public readonly encryptedEvent?: TimelineEvent
     ) {}
 
     setDevice(device: DeviceIdentity): void {
diff --git a/src/matrix/e2ee/RoomEncryption.js b/src/matrix/e2ee/RoomEncryption.js
index dbd8fd43..0e17ad6f 100644
--- a/src/matrix/e2ee/RoomEncryption.js
+++ b/src/matrix/e2ee/RoomEncryption.js
@@ -221,7 +221,30 @@ export class RoomEncryption {
         }));
     }
 
+    /** fetches the devices that are not yet known locally from the homeserver to verify the sender of this message. */
+    _fetchKeyAndVerifyDecryptionResults(results, hsApi, log) {
+        const resultsWithoutDevice = results.filter(r => r.isVerificationUnknown);
+        if (resultsWithoutDevice.length) {
+            return log.wrap("fetch unverified senders", async log => {
+                const sendersWithoutDevice = Array.from(resultsWithoutDevice.reduce((senders, r) => {
+                    return senders.add(r.encryptedEvent.sender);
+                }, new Set()));
+                log.set("senders", sendersWithoutDevice);
+                // fetch the devices, ignore return value,
+                // and just reuse _verifyDecryptionResults method so we only have one impl how to verify
+                await this._deviceTracker.devicesForRoomMembers(this._room.id, sendersWithoutDevice, hsApi, log);
+                // now that we've fetched the missing devices, try verifying the results again
+                const txn = await this._storage.readTxn([this._storage.storeNames.deviceIdentities]);
+                return this._verifyDecryptionResults(resultsWithoutDevice, txn);
+                const resultsWithFoundDevice = resultsWithoutDevice.filter(r => !r.isVerificationUnknown);
+                const resultsToEventIdMap = resultsWithFoundDevice.reduce((map, r) => {
+                    map.set(r.encryptedEvent.event_id, r);
+                    return map;
+                }, new Map());
+                return new BatchDecryptionResult(resultsToEventIdMap, new Map(), this);
+            });
         }
+        return new BatchDecryptionResult(new Map(), new Map(), this);
     }
 
     async _requestMissingSessionFromBackup(senderKey, sessionId, log) {
@@ -553,6 +576,13 @@ class BatchDecryptionResult {
     verifyKnownSenders(txn) {
         return this._roomEncryption._verifyDecryptionResults(Array.from(this.results.values()), txn);
     }
+
+    /** Verify any decryption results for which we could not find a device when
+     *  calling `verifyKnownSenders` prior, by fetching them from the homeserver.
+     *  @returns {Promise<BatchDecryptionResult>} the results for which we found a device */
+    fetchAndVerifyRemainingSenders(hsApi, log) {
+        return this._roomEncryption._fetchKeyAndVerifyDecryptionResults(Array.from(this.results.values()), hsApi, log);
+    }
 }
 
 import {createMockStorage} from "../../mocks/Storage";
diff --git a/src/matrix/e2ee/megolm/decryption/SessionDecryption.ts b/src/matrix/e2ee/megolm/decryption/SessionDecryption.ts
index 57ef9a96..36224fc5 100644
--- a/src/matrix/e2ee/megolm/decryption/SessionDecryption.ts
+++ b/src/matrix/e2ee/megolm/decryption/SessionDecryption.ts
@@ -75,7 +75,7 @@ export class SessionDecryption {
                             {encryptedRoomId: payload.room_id, eventRoomId: this.key.roomId});
                     }
                     replayEntries.push(new ReplayDetectionEntry(this.key.sessionId, decryptionResult!.message_index, event));
-                    const result = new DecryptionResult(payload, this.key.senderKey, this.key.claimedEd25519Key);
+                    const result = new DecryptionResult(payload, this.key.senderKey, this.key.claimedEd25519Key, event);
                     results.set(event.event_id, result);
                 } catch (err) {
                     // ignore AbortError from cancelling decryption requests in dispose method