From 6e168051b62f70f141a4a602b46c51853c31c164 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Damir=20Jeli=C4=87?= <poljar@termina.org.uk>
Date: Fri, 26 Feb 2021 15:59:27 +0100
Subject: [PATCH] crypto: Chunk out key query requests.

---
 matrix_sdk_crypto/src/identities/manager.rs | 22 +++++++++++----------
 matrix_sdk_crypto/src/machine.rs            | 19 +++++++++---------
 2 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/matrix_sdk_crypto/src/identities/manager.rs b/matrix_sdk_crypto/src/identities/manager.rs
index 1333fadd..5351d638 100644
--- a/matrix_sdk_crypto/src/identities/manager.rs
+++ b/matrix_sdk_crypto/src/identities/manager.rs
@@ -43,6 +43,8 @@ pub(crate) struct IdentityManager {
 }
 
 impl IdentityManager {
+    const MAX_KEY_QUERY_USERS: usize = 250;
+
     pub fn new(user_id: Arc<UserId>, device_id: Arc<DeviceIdBox>, store: Store) -> Self {
         IdentityManager {
             user_id,
@@ -298,19 +300,19 @@ impl IdentityManager {
     ///
     /// [`OlmMachine`]: struct.OlmMachine.html
     /// [`receive_keys_query_response`]: #method.receive_keys_query_response
-    pub async fn users_for_key_query(&self) -> Option<KeysQueryRequest> {
-        let mut users = self.store.users_for_key_query();
+    pub async fn users_for_key_query(&self) -> Vec<KeysQueryRequest> {
+        let users = self.store.users_for_key_query();
 
         if users.is_empty() {
-            None
+            Vec::new()
         } else {
-            let mut device_keys: BTreeMap<UserId, Vec<Box<DeviceId>>> = BTreeMap::new();
+            let users: Vec<UserId> = users.into_iter().collect();
 
-            for user in users.drain() {
-                device_keys.insert(user, Vec::new());
-            }
-
-            Some(KeysQueryRequest::new(device_keys))
+            users
+                .chunks(Self::MAX_KEY_QUERY_USERS)
+                .map(|u| u.iter().map(|u| (u.clone(), Vec::new())).collect())
+                .map(KeysQueryRequest::new)
+                .collect()
         }
     }
 
@@ -566,7 +568,7 @@ pub(crate) mod test {
     #[async_test]
     async fn test_manager_creation() {
         let manager = manager();
-        assert!(manager.users_for_key_query().await.is_none())
+        assert!(manager.users_for_key_query().await.is_empty())
     }
 
     #[async_test]
diff --git a/matrix_sdk_crypto/src/machine.rs b/matrix_sdk_crypto/src/machine.rs
index dca6603b..6ace47b5 100644
--- a/matrix_sdk_crypto/src/machine.rs
+++ b/matrix_sdk_crypto/src/machine.rs
@@ -304,16 +304,17 @@ impl OlmMachine {
             requests.push(r);
         }
 
-        if let Some(r) =
-            self.identity_manager
-                .users_for_key_query()
-                .await
-                .map(|r| OutgoingRequest {
-                    request_id: Uuid::new_v4(),
-                    request: Arc::new(r.into()),
-                })
+        for request in self
+            .identity_manager
+            .users_for_key_query()
+            .await
+            .into_iter()
+            .map(|r| OutgoingRequest {
+                request_id: Uuid::new_v4(),
+                request: Arc::new(r.into()),
+            })
         {
-            requests.push(r);
+            requests.push(request);
         }
 
         requests.append(&mut self.outgoing_to_device_requests());