Download device keys in chunks of 250

Depending on the number of users in the request, the server might
overload. To prevent this, the download is broken into chunks of
250 users each. Additionally, no more than 3 requests are kicked off
at the same time to avoid running into rate limiting. Responses are
processed once all chunks have been downloaded.

Fixes: #1619

Signed-off-by: Johannes Marbach <n0-0ne+github@mailbox.org>

spec/unit/crypto/DeviceList.spec.js

@@ -51,6 +51,36 @@ const signedDeviceList = {
     },
 };
 
+const signedDeviceList2 = {
+    "failures": {},
+    "device_keys": {
+        "@test2:sw1v.org": {
+            "QJVRHWAKGH": {
+                "signatures": {
+                    "@test2:sw1v.org": {
+                        "ed25519:QJVRHWAKGH":
+                            "w1xxdLe1iIqzEFHLRVYQeuiM6t2N2ZRiI8s5nDKxf054BP8" +
+                            "1CPEX/AQXh5BhkKAVMlKnwg4T9zU1/wBALeajk3",
+                    },
+                },
+                "user_id": "@test2:sw1v.org",
+                "keys": {
+                    "ed25519:QJVRHWAKGH":
+                        "Ig0/C6T+bBII1l2By2Wnnvtjp1nm/iXBlLU5/QESFXL",
+                    "curve25519:QJVRHWAKGH":
+                        "YR3eQnUvTQzGlWih4rsmJkKxpDxzgkgIgsBd1DEZIbm",
+                },
+                "algorithms": [
+                    "m.olm.v1.curve25519-aes-sha2",
+                    "m.megolm.v1.aes-sha2",
+                ],
+                "device_id": "QJVRHWAKGH",
+                "unsigned": {},
+            },
+        },
+    },
+};
+
 describe('DeviceList', function() {
     let downloadSpy;
     let cryptoStore;
@@ -69,7 +99,7 @@ describe('DeviceList', function() {
         }
     });
 
-    function createTestDeviceList() {
+    function createTestDeviceList(keyDownloadChunkSize = 250) {
         const baseApis = {
             downloadKeysForUsers: downloadSpy,
             getUserId: () => '@test1:sw1v.org',
@@ -78,7 +108,7 @@ describe('DeviceList', function() {
         const mockOlm = {
             verifySignature: function(key, message, signature) {},
         };
-        const dl = new DeviceList(baseApis, cryptoStore, mockOlm);
+        const dl = new DeviceList(baseApis, cryptoStore, mockOlm, keyDownloadChunkSize);
         deviceLists.push(dl);
         return dl;
     }
@@ -150,4 +180,30 @@ describe('DeviceList', function() {
             expect(Object.keys(storedKeys)).toEqual(['HGKAWHRVJQ']);
         });
     });
+
+    it("should download device keys in batches", function() {
+        const dl = createTestDeviceList(1);
+
+        dl.startTrackingDeviceList('@test1:sw1v.org');
+        dl.startTrackingDeviceList('@test2:sw1v.org');
+
+        const queryDefer1 = utils.defer();
+        downloadSpy.mockReturnValueOnce(queryDefer1.promise);
+        const queryDefer2 = utils.defer();
+        downloadSpy.mockReturnValueOnce(queryDefer2.promise);
+
+        const prom1 = dl.refreshOutdatedDeviceLists();
+        expect(downloadSpy).toBeCalledTimes(2);
+        expect(downloadSpy).toHaveBeenNthCalledWith(1, ['@test1:sw1v.org'], {});
+        expect(downloadSpy).toHaveBeenNthCalledWith(2, ['@test2:sw1v.org'], {});
+        queryDefer1.resolve(utils.deepCopy(signedDeviceList));
+        queryDefer2.resolve(utils.deepCopy(signedDeviceList2));
+
+        return prom1.then(() => {
+            const storedKeys1 = dl.getRawStoredDevicesForUser('@test1:sw1v.org');
+            expect(Object.keys(storedKeys1)).toEqual(['HGKAWHRVJQ']);
+            const storedKeys2 = dl.getRawStoredDevicesForUser('@test2:sw1v.org');
+            expect(Object.keys(storedKeys2)).toEqual(['QJVRHWAKGH']);
+        });
+    });
 });

spec/unit/utils.spec.js

@@ -282,4 +282,30 @@ describe("utils", function() {
             expect(target.nonenumerableProp).toBe(undefined);
         });
     });
+
+    describe("chunkPromises", function() {
+        it("should execute promises in chunks", async function() {
+            let promiseCount = 0;
+
+            function fn1() {
+                return new Promise(async function(resolve, reject) {
+                    await utils.sleep(1);
+                    expect(promiseCount).toEqual(0);
+                    ++promiseCount;
+                    resolve();
+                });
+            }
+
+            function fn2() {
+                return new Promise(function(resolve, reject) {
+                    expect(promiseCount).toEqual(1);
+                    ++promiseCount;
+                    resolve();
+                });
+            }
+
+            await utils.chunkPromises([fn1, fn2], 1);
+            expect(promiseCount).toEqual(2);
+        });
+    });
 });

src/crypto/DeviceList.js

@@ -28,7 +28,7 @@ import {DeviceInfo} from './deviceinfo';
 import {CrossSigningInfo} from './CrossSigning';
 import * as olmlib from './olmlib';
 import {IndexedDBCryptoStore} from './store/indexeddb-crypto-store';
-import {defer, sleep} from '../utils';
+import {chunkPromises, defer, sleep} from '../utils';
 
 
 /* State transition diagram for DeviceList._deviceTrackingStatus
@@ -62,7 +62,7 @@ const TRACKING_STATUS_UP_TO_DATE = 3;
  * @alias module:crypto/DeviceList
  */
 export class DeviceList extends EventEmitter {
-    constructor(baseApis, cryptoStore, olmDevice) {
+    constructor(baseApis, cryptoStore, olmDevice, keyDownloadChunkSize = 250) {
         super();
 
         this._cryptoStore = cryptoStore;
@@ -98,6 +98,9 @@ export class DeviceList extends EventEmitter {
         // userId -> promise
         this._keyDownloadsInProgressByUser = {};
 
+        // Maximum number of user IDs per request to prevent server overload (#1619)
+        this._keyDownloadChunkSize = keyDownloadChunkSize;
+
         // Set whenever changes are made other than setting the sync token
         this._dirty = false;
 
@@ -780,13 +783,17 @@ class DeviceListUpdateSerialiser {
             opts.token = this._syncToken;
         }
 
-        this._baseApis.downloadKeysForUsers(
-            downloadUsers, opts,
-        ).then(async (res) => {
-            const dk = res.device_keys || {};
-            const masterKeys = res.master_keys || {};
-            const ssks = res.self_signing_keys || {};
-            const usks = res.user_signing_keys || {};
+        const factories = [];
+        for (let i = 0; i < downloadUsers.length; i += this._deviceList._keyDownloadChunkSize) {
+            const userSlice = downloadUsers.slice(i, i + this._deviceList._keyDownloadChunkSize);
+            factories.push(() => this._baseApis.downloadKeysForUsers(userSlice, opts));
+        }
+
+        chunkPromises(factories, 3).then(async (responses) => {
+            const dk = Object.assign({}, ...(responses.map(res => res.device_keys || {})));
+            const masterKeys = Object.assign({}, ...(responses.map(res => res.master_keys || {})));
+            const ssks = Object.assign({}, ...(responses.map(res => res.self_signing_keys || {})));
+            const usks = Object.assign({}, ...(responses.map(res => res.user_signing_keys || {})));
 
             // yield to other things that want to execute in between users, to
             // avoid wedging the CPU

src/utils.ts

@@ -745,6 +745,15 @@ export function promiseTry<T>(fn: () => T): Promise<T> {
     return new Promise((resolve) => resolve(fn()));
 }
 
+// Creates and awaits all promises, running no more than `chunkSize` at the same time
+export async function chunkPromises<T>(fns: (() => Promise<T>)[], chunkSize: number): Promise<T[]> {
+    const results: T[] = [];
+    for (let i = 0; i < fns.length; i += chunkSize) {
+        results.push(...(await Promise.all(fns.slice(i, i + chunkSize).map(fn => fn()))));
+    }
+    return results;
+}
+
 // We need to be able to access the Node.js crypto library from within the
 // Matrix SDK without needing to `require("crypto")`, which will fail in
 // browsers.  So `index.ts` will call `setCrypto` to store it, and when we need