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 {
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]

View File

@ -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());