crypto: Add OlmMachine::query_keys_for_users (#2267)

Sometimes we need our key query results to be as up-to-date as possible. Add a mechanism to allow that.

Closes #2263 .

Author: richvdh

bindings/matrix-sdk-crypto-js/CHANGELOG.md

@@ -1,3 +1,8 @@
+# unreleased
+
+-   Add method `OlmMachine.queryKeysForUsers` to build an out-of-band key
+    request.
+
 # v0.1.3
 
 ## Changes in the Javascript bindings

bindings/matrix-sdk-crypto-js/src/machine.rs

@@ -195,10 +195,7 @@ impl OlmMachine {
     /// Users that were already in the list are unaffected.
     #[wasm_bindgen(js_name = "updateTrackedUsers")]
     pub fn update_tracked_users(&self, users: &Array) -> Result<Promise, JsError> {
-        let users = users
-            .iter()
-            .map(|user| Ok(downcast::<identifiers::UserId>(&user, "UserId")?.inner.clone()))
-            .collect::<Result<Vec<ruma::OwnedUserId>, JsError>>()?;
+        let users = user_ids_to_owned_user_ids(users)?;
 
         let me = self.inner.clone();
 
@@ -529,10 +526,7 @@ impl OlmMachine {
         encryption_settings: &encryption::EncryptionSettings,
     ) -> Result<Promise, JsError> {
         let room_id = room_id.inner.clone();
-        let users = users
-            .iter()
-            .map(|user| Ok(downcast::<identifiers::UserId>(&user, "UserId")?.inner.clone()))
-            .collect::<Result<Vec<ruma::OwnedUserId>, JsError>>()?;
+        let users = user_ids_to_owned_user_ids(users)?;
         let encryption_settings =
             matrix_sdk_crypto::olm::EncryptionSettings::from(encryption_settings);
 
@@ -555,6 +549,26 @@ impl OlmMachine {
         }))
     }
 
+    /// Generate an "out-of-band" key query request for the given set of users.
+    ///
+    /// This can be useful if we need the results from `getIdentity` or
+    /// `getUserDevices` to be as up-to-date as possible.
+    ///
+    /// Returns a `KeysQueryRequest` object. The response of the request should
+    /// be passed to the `OlmMachine` with the `mark_request_as_sent`.
+    #[wasm_bindgen(js_name = "queryKeysForUsers")]
+    pub fn query_keys_for_users(
+        &self,
+        users: &Array,
+    ) -> Result<requests::KeysQueryRequest, JsError> {
+        let users = user_ids_to_owned_user_ids(users)?;
+
+        let (request_id, request) =
+            self.inner.query_keys_for_users(users.iter().map(AsRef::as_ref));
+
+        Ok(requests::KeysQueryRequest::try_from((request_id.to_string(), &request))?)
+    }
+
     /// Get the a key claiming request for the user/device pairs that
     /// we are missing Olm sessions for.
     ///
@@ -582,10 +596,7 @@ impl OlmMachine {
     /// empty iterator when calling this method between sync requests.
     #[wasm_bindgen(js_name = "getMissingSessions")]
     pub fn get_missing_sessions(&self, users: &Array) -> Result<Promise, JsError> {
-        let users = users
-            .iter()
-            .map(|user| Ok(downcast::<identifiers::UserId>(&user, "UserId")?.inner.clone()))
-            .collect::<Result<Vec<ruma::OwnedUserId>, JsError>>()?;
+        let users = user_ids_to_owned_user_ids(users)?;
 
         let me = self.inner.clone();
 
@@ -875,3 +886,12 @@ pub(crate) async fn promise_result_to_future(
         }
     }
 }
+
+/// Helper function to take a Javascript array of `UserId`s and turn it into
+/// a Rust `Vec` of `OwnedUserId`s
+fn user_ids_to_owned_user_ids(users: &Array) -> Result<Vec<ruma::OwnedUserId>, JsError> {
+    users
+        .iter()
+        .map(|user| Ok(downcast::<identifiers::UserId>(&user, "UserId")?.inner.clone()))
+        .collect()
+}

bindings/matrix-sdk-crypto-js/tests/machine.test.js

@@ -274,9 +274,17 @@ describe(OlmMachine.name, () => {
         }
     });
 
+    test("Can build a key query request", async () => {
+        const m = await machine();
+        const request = m.queryKeysForUsers([new UserId("@alice:example.org")]);
+        expect(request).toBeInstanceOf(KeysQueryRequest);
+        const body = JSON.parse(request.body);
+        expect(Object.keys(body.device_keys)).toContain("@alice:example.org");
+    });
+
     describe("setup workflow to mark requests as sent", () => {
         let m;
-        let ougoingRequests;
+        let outgoingRequests;
 
         beforeAll(async () => {
             m = await machine(new UserId("@alice:example.org"), new DeviceId("DEVICEID"));

crates/matrix-sdk-crypto/CHANGELOG.md

@@ -43,3 +43,6 @@
 
 - Fix parsing error for `POST /_matrix/client/v3/keys/signatures/upload`
   responses generated by Synapse.
+
+- Add new API `OlmMachine::query_keys_for_users` for generating out-of-band key
+  queries.

crates/matrix-sdk-crypto/src/identities/manager.rs

@@ -105,7 +105,8 @@ impl IdentityManager {
     ///
     /// # Arguments
     ///
-    /// * `request_id` - The request_id returned by users_for_key_query
+    /// * `request_id` - The request_id returned by `users_for_key_query` or
+    ///   `build_key_query_for_users`
     /// * `response` - The keys query response of the request that the client
     /// performed.
     pub async fn receive_keys_query_response(
@@ -619,6 +620,38 @@ impl IdentityManager {
         Ok((changes, changed_identity))
     }
 
+    /// Generate an "out-of-band" key query request for the given set of users.
+    ///
+    /// Unlike the regular key query requests returned by `users_for_key_query`,
+    /// there can be several of these in flight at once. This can be useful
+    /// if we need results to be as up-to-date as possible.
+    ///
+    /// Once the request has been made, the response can be fed back into the
+    /// IdentityManager and store by calling `receive_keys_query_response`.
+    ///
+    /// # Arguments
+    ///
+    /// * `users` - list of users whose keys should be queried
+    ///
+    /// # Returns
+    ///
+    /// A tuple containing the request ID for the request, and the request
+    /// itself.
+    pub(crate) fn build_key_query_for_users<'a>(
+        &self,
+        users: impl IntoIterator<Item = &'a UserId>,
+    ) -> (OwnedTransactionId, KeysQueryRequest) {
+        // Since this is an "out-of-band" request, we just make up a transaction ID and
+        // do not store the details in `self.keys_query_request_details`.
+        //
+        // `receive_keys_query_response` will process the response as normal, except
+        // that it will not mark the users as "up-to-date".
+
+        // We assume that there aren't too many users here; if we find a usecase that
+        // requires lots of users to be up-to-date we may need to rethink this.
+        (TransactionId::new(), KeysQueryRequest::new(users.into_iter().map(|u| u.to_owned())))
+    }
+
     /// Get a list of key query requests needed.
     ///
     /// # Returns
@@ -969,6 +1002,7 @@ pub(crate) mod tests {
     use serde_json::json;
 
     use super::testing::{device_id, key_query, manager, other_key_query, other_user_id, user_id};
+    use crate::identities::manager::testing::own_key_query;
 
     fn key_query_with_failures() -> KeysQueryResponse {
         let response = json!({
@@ -1161,4 +1195,24 @@ pub(crate) mod tests {
             .iter()
             .any(|(_, r)| r.device_keys.contains_key(alice)));
     }
+
+    #[async_test]
+    async fn test_out_of_band_key_query() {
+        // build the request
+        let manager = manager().await;
+        let (reqid, req) = manager.build_key_query_for_users(vec![user_id()]);
+        assert!(req.device_keys.contains_key(user_id()));
+
+        // make up a response and check it is processed
+        let (device_changes, identity_changes) =
+            manager.receive_keys_query_response(&reqid, &own_key_query()).await.unwrap();
+        assert_eq!(device_changes.new.len(), 1);
+        assert_eq!(device_changes.new[0].device_id(), "LVWOVGOXME");
+        assert_eq!(identity_changes.new.len(), 1);
+        assert_eq!(identity_changes.new[0].user_id(), user_id());
+
+        let devices = manager.store.get_user_devices(user_id()).await.unwrap();
+        assert_eq!(devices.devices().count(), 1);
+        assert_eq!(devices.devices().next().unwrap().device_id(), "LVWOVGOXME");
+    }
 }

crates/matrix-sdk-crypto/src/machine.rs

@@ -86,8 +86,8 @@ use crate::{
         Signatures,
     },
     verification::{Verification, VerificationMachine, VerificationRequest},
-    CrossSigningKeyExport, CryptoStoreError, LocalTrust, ReadOnlyDevice, RoomKeyImportResult,
-    SignatureError, ToDeviceRequest,
+    CrossSigningKeyExport, CryptoStoreError, KeysQueryRequest, LocalTrust, ReadOnlyDevice,
+    RoomKeyImportResult, SignatureError, ToDeviceRequest,
 };
 
 /// State machine implementation of the Olm/Megolm encryption protocol used for
@@ -402,6 +402,30 @@ impl OlmMachine {
         Ok(requests)
     }
 
+    /// Generate an "out-of-band" key query request for the given set of users.
+    ///
+    /// This can be useful if we need the results from [`get_identity`] or
+    /// [`get_user_devices`] to be as up-to-date as possible.
+    ///
+    /// # Arguments
+    ///
+    /// * `users` - list of users whose keys should be queried
+    ///
+    /// # Returns
+    ///
+    /// A request to be sent out to the server. Once sent, the response should
+    /// be passed back to the state machine using [`mark_request_as_sent`].
+    ///
+    /// [`mark_request_as_sent`]: OlmMachine::mark_request_as_sent
+    /// [`get_identity`]: OlmMachine::get_identity
+    /// [`get_user_devices`]: OlmMachine::get_user_devices
+    pub fn query_keys_for_users<'a>(
+        &self,
+        users: impl IntoIterator<Item = &'a UserId>,
+    ) -> (OwnedTransactionId, KeysQueryRequest) {
+        self.inner.identity_manager.build_key_query_for_users(users)
+    }
+
     /// Mark the request with the given request id as sent.
     ///
     /// # Arguments
@@ -472,7 +496,7 @@ impl OlmMachine {
     ///
     /// These requests may require user interactive auth.
     ///
-    /// [`mark_request_as_sent`]: #method.mark_request_as_sent`mark_request_
+    /// [`mark_request_as_sent`]: #method.mark_request_as_sent
     pub async fn bootstrap_cross_signing(
         &self,
         reset: bool,
@@ -2238,6 +2262,14 @@ pub(crate) mod tests {
         assert_eq!(device.device_id(), alice_device_id);
     }
 
+    #[async_test]
+    async fn test_query_keys_for_users() {
+        let (machine, _) = get_prepared_machine(false).await;
+        let alice_id = user_id!("@alice:example.org");
+        let (_, request) = machine.query_keys_for_users(vec![alice_id]);
+        assert!(request.device_keys.contains_key(alice_id));
+    }
+
     #[async_test]
     async fn test_missing_sessions_calculation() {
         let (machine, _) = get_machine_after_query().await;