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,15 +270,151 @@ 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
+ * @return {array<object<userid, deviceInfo>>}
- *     message has been sent.
  */
-MegolmEncryption.prototype._shareKeyWithDevices = function(session, devicesByUser) {
+MegolmEncryption.prototype._splitUserDeviceMap = function(
-    const self = this;
+    session, chainIndex, devicemap, devicesByUser,
+) {
+    const maxToDeviceMessagesPerRequest = 20;
 
+    // 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];
+
+        for (let i = 0; i < devicesToShareWith.length; i++) {
+            const deviceInfo = devicesToShareWith[i];
+            const deviceId = deviceInfo.deviceId;
+
+            const sessionResult = sessionResults[deviceId];
+            if (!sessionResult.sessionId) {
+                // no session with this device, probably because there
+                // were no one-time keys.
+                //
+                // we could send them a to_device message anyway, as a
+                // signal that they have missed out on the key sharing
+                // 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(
+                "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: 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,
+                this._userId,
+                this._deviceId,
+                this._olmDevice,
+                userId,
+                deviceInfo,
+                payload,
+            ),
+        );
+    }
+
+    return Promise.all(promises).then(() => {
+        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,
+                    );
+                }
+            }
+        });
+    });
+};
+
+/**
+ * @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",
@@ -283,104 +427,28 @@ MegolmEncryption.prototype._shareKeyWithDevices = function(session, devicesByUse
         },
     };
 
-    const contentMap = {};
+    const devicemap = await olmlib.ensureOlmSessionsForDevices(
-
-    return olmlib.ensureOlmSessionsForDevices(
         this._olmDevice, this._baseApis, devicesByUser,
-    ).then(function(devicemap) {
+    );
-        const promises = [];
 
-        for (const userId in devicesByUser) {
+    const userDeviceMaps = this._splitUserDeviceMap(
-            if (!devicesByUser.hasOwnProperty(userId)) {
+        session, key.chain_index, devicemap, devicesByUser,
-                continue;
+    );
-            }
 
-            const devicesToShareWith = devicesByUser[userId];
+    for (let i = 0; i < userDeviceMaps.length; i++) {
-            const sessionResults = devicemap[userId];
+        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`);
 
-            for (let i = 0; i < devicesToShareWith.length; i++) {
+            throw e;
-                const deviceInfo = devicesToShareWith[i];
-                const deviceId = deviceInfo.deviceId;
-
-                const sessionResult = sessionResults[deviceId];
-                if (!sessionResult.sessionId) {
-                    // no session with this device, probably because there
-                    // were no one-time keys.
-                    //
-                    // we could send them a to_device message anyway, as a
-                    // signal that they have missed out on the key sharing
-                    // 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.
-                    //
-                    // ensureOlmSessionsForUsers has already done the logging,
-                    // so just skip it.
-                    continue;
-                }
-
-                console.log(
-                    "sharing keys with device " + userId + ":" + deviceId,
-                );
-
-                const encryptedContent = {
-                    algorithm: olmlib.OLM_ALGORITHM,
-                    sender_key: self._olmDevice.deviceCurve25519Key,
-                    ciphertext: {},
-                };
-
-                if (!contentMap[userId]) {
-                    contentMap[userId] = {};
-                }
-
-                contentMap[userId][deviceId] = encryptedContent;
-
-                promises.push(
-                    olmlib.encryptMessageForDevice(
-                        encryptedContent.ciphertext,
-                        self._userId,
-                        self._deviceId,
-                        self._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);
-        });
-    }).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;
-            }
-        }
-    });
 };
 
 /**