From 8663fd402b202b775eea347ae2e091f48a1f18e6 Mon Sep 17 00:00:00 2001 From: Johannes Marbach Date: Wed, 10 Mar 2021 20:05:42 +0100 Subject: [PATCH] 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 --- spec/unit/crypto/DeviceList.spec.js | 60 ++++++++++++++++++++++++++++- spec/unit/utils.spec.js | 26 +++++++++++++ src/crypto/DeviceList.js | 25 +++++++----- src/utils.ts | 9 +++++ 4 files changed, 109 insertions(+), 11 deletions(-) diff --git a/spec/unit/crypto/DeviceList.spec.js b/spec/unit/crypto/DeviceList.spec.js index 3ba25ae1f..a9ff83054 100644 --- a/spec/unit/crypto/DeviceList.spec.js +++ b/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']); + }); + }); }); diff --git a/spec/unit/utils.spec.js b/spec/unit/utils.spec.js index 0b9f8ab93..686912913 100644 --- a/spec/unit/utils.spec.js +++ b/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); + }); + }); }); diff --git a/src/crypto/DeviceList.js b/src/crypto/DeviceList.js index 797f4b860..97836b5a4 100644 --- a/src/crypto/DeviceList.js +++ b/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 diff --git a/src/utils.ts b/src/utils.ts index 95b3c11c5..17b3924cd 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -745,6 +745,15 @@ export function promiseTry(fn: () => T): Promise { return new Promise((resolve) => resolve(fn())); } +// Creates and awaits all promises, running no more than `chunkSize` at the same time +export async function chunkPromises(fns: (() => Promise)[], chunkSize: number): Promise { + 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