crypto: Chunk out key query requests.

master
Damir Jelić 2021-02-26 15:59:27 +01:00
parent 2a09e588f3
commit 6e168051b6
2 changed files with 22 additions and 19 deletions

View File

@ -43,6 +43,8 @@ pub(crate) struct IdentityManager {
} }
impl IdentityManager { impl IdentityManager {
const MAX_KEY_QUERY_USERS: usize = 250;
pub fn new(user_id: Arc<UserId>, device_id: Arc<DeviceIdBox>, store: Store) -> Self { pub fn new(user_id: Arc<UserId>, device_id: Arc<DeviceIdBox>, store: Store) -> Self {
IdentityManager { IdentityManager {
user_id, user_id,
@ -298,19 +300,19 @@ impl IdentityManager {
/// ///
/// [`OlmMachine`]: struct.OlmMachine.html /// [`OlmMachine`]: struct.OlmMachine.html
/// [`receive_keys_query_response`]: #method.receive_keys_query_response /// [`receive_keys_query_response`]: #method.receive_keys_query_response
pub async fn users_for_key_query(&self) -> Option<KeysQueryRequest> { pub async fn users_for_key_query(&self) -> Vec<KeysQueryRequest> {
let mut users = self.store.users_for_key_query(); let users = self.store.users_for_key_query();
if users.is_empty() { if users.is_empty() {
None Vec::new()
} else { } 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() { users
device_keys.insert(user, Vec::new()); .chunks(Self::MAX_KEY_QUERY_USERS)
} .map(|u| u.iter().map(|u| (u.clone(), Vec::new())).collect())
.map(KeysQueryRequest::new)
Some(KeysQueryRequest::new(device_keys)) .collect()
} }
} }
@ -566,7 +568,7 @@ pub(crate) mod test {
#[async_test] #[async_test]
async fn test_manager_creation() { async fn test_manager_creation() {
let manager = manager(); let manager = manager();
assert!(manager.users_for_key_query().await.is_none()) assert!(manager.users_for_key_query().await.is_empty())
} }
#[async_test] #[async_test]

View File

@ -304,16 +304,17 @@ impl OlmMachine {
requests.push(r); requests.push(r);
} }
if let Some(r) = for request in self
self.identity_manager .identity_manager
.users_for_key_query() .users_for_key_query()
.await .await
.map(|r| OutgoingRequest { .into_iter()
request_id: Uuid::new_v4(), .map(|r| OutgoingRequest {
request: Arc::new(r.into()), request_id: Uuid::new_v4(),
}) request: Arc::new(r.into()),
})
{ {
requests.push(r); requests.push(request);
} }
requests.append(&mut self.outgoing_to_device_requests()); requests.append(&mut self.outgoing_to_device_requests());