reduce sendToDevice payload (#522)

instead of sending one huge request split them up.

src/crypto/algorithms/megolm.js

@@ -74,6 +74,14 @@ OutboundSessionInfo.prototype.needsRotation = function(
     return false;
 };
 
+OutboundSessionInfo.prototype.markSharedWithDevice = function(
+    userId, deviceId, chainIndex,
+) {
+    if (!this.sharedWithDevices[userId]) {
+        this.sharedWithDevices[userId] = {};
+    }
+    this.sharedWithDevices[userId][deviceId] = chainIndex;
+};
 
 /**
  * Determine if this session has been shared with devices which it shouldn't
@@ -262,39 +270,27 @@ MegolmEncryption.prototype._prepareNewSession = async function() {
  *
  * @param {module:crypto/algorithms/megolm.OutboundSessionInfo} session
  *
+ * @param {number} chainIndex current chain index
+ *
+ * @param {object<userId, deviceId>} devicemap
+ *   mapping from userId to deviceId to {@link module:crypto~OlmSessionResult}
+ *
  * @param {object<string, module:crypto/deviceinfo[]>} devicesByUser
  *    map from userid to list of devices
  *
- * @return {module:client.Promise} Promise which resolves once the key sharing
- *     message has been sent.
+ * @return {array<object<userid, deviceInfo>>}
  */
-MegolmEncryption.prototype._shareKeyWithDevices = function(session, devicesByUser) {
-    const self = this;
+MegolmEncryption.prototype._splitUserDeviceMap = function(
+    session, chainIndex, devicemap, devicesByUser,
+) {
+    const maxToDeviceMessagesPerRequest = 20;
 
-    const key = this._olmDevice.getOutboundGroupSessionKey(session.sessionId);
-    const payload = {
-        type: "m.room_key",
-        content: {
-            algorithm: olmlib.MEGOLM_ALGORITHM,
-            room_id: this._roomId,
-            session_id: session.sessionId,
-            session_key: key.key,
-            chain_index: key.chain_index,
-        },
-    };
-
-    const contentMap = {};
-
-    return olmlib.ensureOlmSessionsForDevices(
-        this._olmDevice, this._baseApis, devicesByUser,
-    ).then(function(devicemap) {
-        const promises = [];
-
-        for (const userId in devicesByUser) {
-            if (!devicesByUser.hasOwnProperty(userId)) {
-                continue;
-            }
+    // use an array where the slices of a content map gets stored
+    const mapSlices = [];
+    let currentSliceId = 0; // start inserting in the first slice
+    let entriesInCurrentSlice = 0;
 
+    for (const userId of Object.keys(devicesByUser)) {
         const devicesToShareWith = devicesByUser[userId];
         const sessionResults = devicemap[userId];
 
@@ -312,75 +308,147 @@ MegolmEncryption.prototype._shareKeyWithDevices = function(session, devicesByUse
                 // message because of the lack of keys, but there's not
                 // much point in that really; it will mostly serve to clog
                 // up to_device inboxes.
-                    //
+
+                // mark this device as "handled" because we don't want to try
+                // to claim a one-time-key for dead devices on every message.
+                session.markSharedWithDevice(userId, deviceId, chainIndex);
+
                 // ensureOlmSessionsForUsers has already done the logging,
                 // so just skip it.
                 continue;
             }
 
             console.log(
-                    "sharing keys with device " + userId + ":" + deviceId,
+                "share keys with device " + userId + ":" + deviceId,
             );
 
+            if (entriesInCurrentSlice > maxToDeviceMessagesPerRequest) {
+                // the current slice is filled up. Start inserting into the next slice
+                entriesInCurrentSlice = 0;
+                currentSliceId++;
+            }
+            if (!mapSlices[currentSliceId]) {
+                mapSlices[currentSliceId] = [];
+            }
+
+            mapSlices[currentSliceId].push({
+                userId: userId,
+                deviceInfo: deviceInfo,
+            });
+
+            entriesInCurrentSlice++;
+        }
+    }
+    return mapSlices;
+};
+
+/**
+ * @private
+ *
+ * @param {module:crypto/algorithms/megolm.OutboundSessionInfo} session
+ *
+ * @param {number} chainIndex current chain index
+ *
+ * @param {object<userId, deviceInfo>} userDeviceMap
+ *   mapping from userId to deviceInfo
+ *
+ * @param {object} payload fields to include in the encrypted payload
+ *
+ * @return {module:client.Promise} Promise which resolves once the key sharing
+ *     for the given userDeviceMap is generated and has been sent.
+ */
+MegolmEncryption.prototype._encryptAndSendKeysToDevices = function(
+    session, chainIndex, userDeviceMap, payload,
+) {
     const encryptedContent = {
         algorithm: olmlib.OLM_ALGORITHM,
-                    sender_key: self._olmDevice.deviceCurve25519Key,
+        sender_key: this._olmDevice.deviceCurve25519Key,
         ciphertext: {},
     };
+    const contentMap = {};
+
+    const promises = [];
+    for (let i = 0; i < userDeviceMap.length; i++) {
+        const val = userDeviceMap[i];
+        const userId = val.userId;
+        const deviceInfo = val.deviceInfo;
+        const deviceId = deviceInfo.deviceId;
 
         if (!contentMap[userId]) {
             contentMap[userId] = {};
         }
-
         contentMap[userId][deviceId] = encryptedContent;
 
         promises.push(
             olmlib.encryptMessageForDevice(
                 encryptedContent.ciphertext,
-                        self._userId,
-                        self._deviceId,
-                        self._olmDevice,
+                this._userId,
+                this._deviceId,
+                this._olmDevice,
                 userId,
                 deviceInfo,
                 payload,
             ),
         );
     }
-        }
-
-        if (promises.length === 0) {
-            // no devices to send to
-            return Promise.resolve();
-        }
 
     return Promise.all(promises).then(() => {
-            // TODO: retries
-            return self._baseApis.sendToDevice("m.room.encrypted", contentMap);
+        return this._baseApis.sendToDevice("m.room.encrypted", contentMap).then(() => {
+            // store that we successfully uploaded the keys of the current slice
+            for (const userId of Object.keys(contentMap)) {
+                for (const deviceId of Object.keys(contentMap[userId])) {
+                    session.markSharedWithDevice(
+                        userId, deviceId, chainIndex,
+                    );
+                }
+            }
         });
-    }).then(function() {
-        console.log(`Completed megolm keyshare in ${self._roomId}`);
+    });
+};
 
-        // Add the devices we have shared with to session.sharedWithDevices.
-        //
-        // we deliberately iterate over devicesByUser (ie, the devices we
-        // attempted to share with) rather than the contentMap (those we did
-        // share with), because we don't want to try to claim a one-time-key
-        // for dead devices on every message.
-        for (const userId in devicesByUser) {
-            if (!devicesByUser.hasOwnProperty(userId)) {
-                continue;
-            }
-            if (!session.sharedWithDevices[userId]) {
-                session.sharedWithDevices[userId] = {};
-            }
-            const devicesToShareWith = devicesByUser[userId];
-            for (let i = 0; i < devicesToShareWith.length; i++) {
-                const deviceInfo = devicesToShareWith[i];
-                session.sharedWithDevices[userId][deviceInfo.deviceId] =
-                    key.chain_index;
+/**
+ * @private
+ *
+ * @param {module:crypto/algorithms/megolm.OutboundSessionInfo} session
+ *
+ * @param {object<string, module:crypto/deviceinfo[]>} devicesByUser
+ *    map from userid to list of devices
+ */
+MegolmEncryption.prototype._shareKeyWithDevices = async function(session, devicesByUser) {
+    const key = this._olmDevice.getOutboundGroupSessionKey(session.sessionId);
+    const payload = {
+        type: "m.room_key",
+        content: {
+            algorithm: olmlib.MEGOLM_ALGORITHM,
+            room_id: this._roomId,
+            session_id: session.sessionId,
+            session_key: key.key,
+            chain_index: key.chain_index,
+        },
+    };
+
+    const devicemap = await olmlib.ensureOlmSessionsForDevices(
+        this._olmDevice, this._baseApis, devicesByUser,
+    );
+
+    const userDeviceMaps = this._splitUserDeviceMap(
+        session, key.chain_index, devicemap, devicesByUser,
+    );
+
+    for (let i = 0; i < userDeviceMaps.length; i++) {
+        try {
+            await this._encryptAndSendKeysToDevices(
+                session, key.chain_index, userDeviceMaps[i], payload,
+            );
+            console.log(`Completed megolm keyshare in ${this._roomId} `
+                + `(slice ${i + 1}/${userDeviceMaps.length})`);
+        } catch (e) {
+            console.log(`megolm keyshare in ${this._roomId} `
+                + `(slice ${i + 1}/${userDeviceMaps.length}) failed`);
+
+            throw e;
         }
     }
-    });
 };
 
 /**