From aff32afefa9aa6eab33a2272b428b11b1a87f740 Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <richard@matrix.org>
Date: Wed, 26 Jul 2017 08:04:46 +0100
Subject: [PATCH] Make olmlib.encryptMessageForDevice async

---
 spec/unit/crypto/algorithms/megolm.spec.js |  3 +-
 src/crypto/algorithms/megolm.js            | 56 ++++++++++++----------
 src/crypto/algorithms/olm.js               | 14 ++++--
 src/crypto/olmlib.js                       |  5 +-
 4 files changed, 45 insertions(+), 33 deletions(-)

diff --git a/spec/unit/crypto/algorithms/megolm.spec.js b/spec/unit/crypto/algorithms/megolm.spec.js
index ec7e65d41..f83307612 100644
--- a/spec/unit/crypto/algorithms/megolm.spec.js
+++ b/spec/unit/crypto/algorithms/megolm.spec.js
@@ -60,7 +60,8 @@ describe("MegolmDecryption", function() {
         // we stub out the olm encryption bits
         mockOlmLib = {};
         mockOlmLib.ensureOlmSessionsForDevices = expect.createSpy();
-        mockOlmLib.encryptMessageForDevice = expect.createSpy();
+        mockOlmLib.encryptMessageForDevice =
+            expect.createSpy().andReturn(Promise.resolve());
         megolmDecryption.olmlib = mockOlmLib;
     });
 
diff --git a/src/crypto/algorithms/megolm.js b/src/crypto/algorithms/megolm.js
index bb9bf4d25..9f630327f 100644
--- a/src/crypto/algorithms/megolm.js
+++ b/src/crypto/algorithms/megolm.js
@@ -288,7 +288,7 @@ MegolmEncryption.prototype._shareKeyWithDevices = function(session, devicesByUse
     return olmlib.ensureOlmSessionsForDevices(
         this._olmDevice, this._baseApis, devicesByUser,
     ).then(function(devicemap) {
-        let haveTargets = false;
+        const promises = [];
 
         for (const userId in devicesByUser) {
             if (!devicesByUser.hasOwnProperty(userId)) {
@@ -328,31 +328,35 @@ MegolmEncryption.prototype._shareKeyWithDevices = function(session, devicesByUse
                     ciphertext: {},
                 };
 
-                olmlib.encryptMessageForDevice(
-                    encryptedContent.ciphertext,
-                    self._userId,
-                    self._deviceId,
-                    self._olmDevice,
-                    userId,
-                    deviceInfo,
-                    payload,
-                );
-
                 if (!contentMap[userId]) {
                     contentMap[userId] = {};
                 }
 
                 contentMap[userId][deviceId] = encryptedContent;
-                haveTargets = true;
+
+                promises.push(
+                    olmlib.encryptMessageForDevice(
+                        encryptedContent.ciphertext,
+                        self._userId,
+                        self._deviceId,
+                        self._olmDevice,
+                        userId,
+                        deviceInfo,
+                        payload,
+                    ),
+                );
             }
         }
 
-        if (!haveTargets) {
+        if (promises.length === 0) {
+            // no devices to send to
             return Promise.resolve();
         }
 
-        // TODO: retries
-        return self._baseApis.sendToDevice("m.room.encrypted", contentMap);
+        return Promise.all(promises).then(() => {
+            // TODO: retries
+            return self._baseApis.sendToDevice("m.room.encrypted", contentMap);
+        });
     }).then(function() {
         console.log(`Completed megolm keyshare in ${self._roomId}`);
 
@@ -753,7 +757,7 @@ MegolmDecryption.prototype.shareKeysWithDevice = function(keyRequest) {
             //
             // ensureOlmSessionsForUsers has already done the logging,
             // so just skip it.
-            return;
+            return null;
         }
 
         console.log(
@@ -772,7 +776,7 @@ MegolmDecryption.prototype.shareKeysWithDevice = function(keyRequest) {
             ciphertext: {},
         };
 
-        this.olmlib.encryptMessageForDevice(
+        return this.olmlib.encryptMessageForDevice(
             encryptedContent.ciphertext,
             this._userId,
             this._deviceId,
@@ -780,16 +784,16 @@ MegolmDecryption.prototype.shareKeysWithDevice = function(keyRequest) {
             userId,
             deviceInfo,
             payload,
-        );
+        ).then(() => {
+            const contentMap = {
+                [userId]: {
+                    [deviceId]: encryptedContent,
+                },
+            };
 
-        const contentMap = {
-            [userId]: {
-                [deviceId]: encryptedContent,
-            },
-        };
-
-        // TODO: retries
-        return this._baseApis.sendToDevice("m.room.encrypted", contentMap);
+            // TODO: retries
+            return this._baseApis.sendToDevice("m.room.encrypted", contentMap);
+        });
     }).done();
 };
 
diff --git a/src/crypto/algorithms/olm.js b/src/crypto/algorithms/olm.js
index 9cfe49a9b..26f2dd4b3 100644
--- a/src/crypto/algorithms/olm.js
+++ b/src/crypto/algorithms/olm.js
@@ -107,6 +107,8 @@ OlmEncryption.prototype.encryptMessage = function(room, eventType, content) {
             ciphertext: {},
         };
 
+        const promises = [];
+
         for (let i = 0; i < users.length; ++i) {
             const userId = users[i];
             const devices = self._crypto.getStoredDevicesForUser(userId);
@@ -123,15 +125,17 @@ OlmEncryption.prototype.encryptMessage = function(room, eventType, content) {
                     continue;
                 }
 
-                olmlib.encryptMessageForDevice(
-                    encryptedContent.ciphertext,
-                    self._userId, self._deviceId, self._olmDevice,
-                    userId, deviceInfo, payloadFields,
+                promises.push(
+                    olmlib.encryptMessageForDevice(
+                        encryptedContent.ciphertext,
+                        self._userId, self._deviceId, self._olmDevice,
+                        userId, deviceInfo, payloadFields,
+                    ),
                 );
             }
         }
 
-        return encryptedContent;
+        return Promise.all(promises).return(encryptedContent);
     });
 };
 
diff --git a/src/crypto/olmlib.js b/src/crypto/olmlib.js
index 476aef57d..5cc162bcd 100644
--- a/src/crypto/olmlib.js
+++ b/src/crypto/olmlib.js
@@ -48,8 +48,11 @@ module.exports.MEGOLM_ALGORITHM = "m.megolm.v1.aes-sha2";
  * @param {string} recipientUserId
  * @param {module:crypto/deviceinfo} recipientDevice
  * @param {object} payloadFields fields to include in the encrypted payload
+ *
+ * Returns a promise which resolves (to undefined) when the payload
+ *    has been encrypted into `resultsObject`
  */
-module.exports.encryptMessageForDevice = function(
+module.exports.encryptMessageForDevice = async function(
     resultsObject,
     ourUserId, ourDeviceId, olmDevice, recipientUserId, recipientDevice,
     payloadFields,