crypto: Add a public method to request and re-request keys.

master
Damir Jelić 2021-04-19 15:00:21 +02:00
parent 8c007510cd
commit 78b7dcac61
3 changed files with 158 additions and 35 deletions

View File

@ -141,6 +141,8 @@ pub(crate) struct KeyRequestMachine {
/// A struct describing an outgoing key request. /// A struct describing an outgoing key request.
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OutgoingKeyRequest { pub struct OutgoingKeyRequest {
/// The user we requested the key from
pub request_recipient: UserId,
/// The unique id of the key request. /// The unique id of the key request.
pub request_id: Uuid, pub request_id: Uuid,
/// The info of the requested key. /// The info of the requested key.
@ -150,11 +152,7 @@ pub struct OutgoingKeyRequest {
} }
impl OutgoingKeyRequest { impl OutgoingKeyRequest {
fn to_request( fn to_request(&self, own_device_id: &DeviceId) -> Result<OutgoingRequest, serde_json::Error> {
&self,
recipient: &UserId,
own_device_id: &DeviceId,
) -> Result<OutgoingRequest, serde_json::Error> {
let content = RoomKeyRequestToDeviceEventContent { let content = RoomKeyRequestToDeviceEventContent {
action: Action::Request, action: Action::Request,
request_id: self.request_id.to_string(), request_id: self.request_id.to_string(),
@ -162,7 +160,22 @@ impl OutgoingKeyRequest {
body: Some(self.info.clone()), body: Some(self.info.clone()),
}; };
wrap_key_request_content(recipient.to_owned(), self.request_id, &content) wrap_key_request_content(self.request_recipient.clone(), self.request_id, &content)
}
fn to_cancelation(
&self,
own_device_id: &DeviceId,
) -> Result<OutgoingRequest, serde_json::Error> {
let content = RoomKeyRequestToDeviceEventContent {
action: Action::CancelRequest,
request_id: self.request_id.to_string(),
requesting_device_id: own_device_id.to_owned(),
body: None,
};
let id = Uuid::new_v4();
wrap_key_request_content(self.request_recipient.clone(), id, &content)
} }
} }
@ -229,7 +242,7 @@ impl KeyRequestMachine {
.into_iter() .into_iter()
.filter(|i| !i.sent_out) .filter(|i| !i.sent_out)
.map(|info| { .map(|info| {
info.to_request(self.user_id(), self.device_id()) info.to_request(self.device_id())
.map_err(CryptoStoreError::from) .map_err(CryptoStoreError::from)
}) })
.collect() .collect()
@ -541,6 +554,69 @@ impl KeyRequestMachine {
} }
} }
/// Create a new outgoing key request for the key with the given session id.
///
/// This will queue up a new to-device request and store the key info so
/// once we receive a forwarded room key we can check that it matches the
/// key we requested.
///
/// This method will return a cancel request and a new key request if the
/// key was already requested, otherwise it will return just the key
/// request.
///
/// # Arguments
///
/// * `room_id` - The id of the room where the key is used in.
///
/// * `sender_key` - The curve25519 key of the sender that owns the key.
///
/// * `session_id` - The id that uniquely identifies the session.
pub async fn request_key(
&self,
room_id: &RoomId,
sender_key: &str,
session_id: &str,
) -> Result<(Option<OutgoingRequest>, OutgoingRequest), CryptoStoreError> {
let key_info = RequestedKeyInfo {
algorithm: EventEncryptionAlgorithm::MegolmV1AesSha2,
room_id: room_id.to_owned(),
sender_key: sender_key.to_owned(),
session_id: session_id.to_owned(),
};
let request = self.store.get_key_request_by_info(&key_info).await?;
if let Some(request) = request {
let cancel = request.to_cancelation(self.device_id())?;
let request = request.to_request(self.device_id())?;
Ok((Some(cancel), request))
} else {
let request = self.request_key_helper(key_info).await?;
Ok((None, request))
}
}
async fn request_key_helper(
&self,
key_info: RequestedKeyInfo,
) -> Result<OutgoingRequest, CryptoStoreError> {
info!("Creating new outgoing room key request {:#?}", key_info);
let request = OutgoingKeyRequest {
request_recipient: self.user_id().to_owned(),
request_id: Uuid::new_v4(),
info: key_info,
sent_out: false,
};
let outgoing_request = request.to_request(self.device_id())?;
self.save_outgoing_key_info(request).await?;
Ok(outgoing_request)
}
/// Create a new outgoing key request for the key with the given session id. /// Create a new outgoing key request for the key with the given session id.
/// ///
/// This will queue up a new to-device request and store the key info so /// This will queue up a new to-device request and store the key info so
@ -570,23 +646,10 @@ impl KeyRequestMachine {
let request = self.store.get_key_request_by_info(&key_info).await?; let request = self.store.get_key_request_by_info(&key_info).await?;
if request.is_some() { if request.is_none() {
// We already sent out a request for this key, nothing to do. self.request_key_helper(key_info).await?;
return Ok(());
} }
info!("Creating new outgoing room key request {:#?}", key_info);
let id = Uuid::new_v4();
let info = OutgoingKeyRequest {
request_id: id,
info: key_info,
sent_out: false,
};
self.save_outgoing_key_info(info).await?;
Ok(()) Ok(())
} }
@ -655,18 +718,9 @@ impl KeyRequestMachine {
// can delete it in one transaction. // can delete it in one transaction.
self.delete_key_info(&key_info).await?; self.delete_key_info(&key_info).await?;
let content = RoomKeyRequestToDeviceEventContent { let request = key_info.to_cancelation(self.device_id())?;
action: Action::CancelRequest, self.outgoing_to_device_requests
request_id: key_info.request_id.to_string(), .insert(request.request_id, request);
requesting_device_id: (&*self.device_id).clone(),
body: None,
};
let id = Uuid::new_v4();
let request = wrap_key_request_content(self.user_id().clone(), id, &content)?;
self.outgoing_to_device_requests.insert(id, request);
Ok(()) Ok(())
} }
@ -840,6 +894,41 @@ mod test {
let machine = get_machine().await; let machine = get_machine().await;
let account = account(); let account = account();
let (_, session) = account
.create_group_session_pair_with_defaults(&room_id())
.await
.unwrap();
assert!(machine
.outgoing_to_device_requests()
.await
.unwrap()
.is_empty());
let (cancel, request) = machine
.request_key(session.room_id(), &session.sender_key, session.session_id())
.await
.unwrap();
assert!(cancel.is_none());
machine
.mark_outgoing_request_as_sent(request.request_id)
.await
.unwrap();
let (cancel, _) = machine
.request_key(session.room_id(), &session.sender_key, session.session_id())
.await
.unwrap();
assert!(cancel.is_some());
}
#[async_test]
async fn re_request_keys() {
let machine = get_machine().await;
let account = account();
let (_, session) = account let (_, session) = account
.create_group_session_pair_with_defaults(&room_id()) .create_group_session_pair_with_defaults(&room_id())
.await .await

View File

@ -918,6 +918,38 @@ impl OlmMachine {
Ok(ToDevice { events }) Ok(ToDevice { events })
} }
/// Request a room key from our devices.
///
/// This method will return a request cancelation and a new key request if
/// the key was already requested, otherwise it will return just the key
/// request.
///
/// The request cancelation *must* be sent out before the request is sent
/// out, otherwise devices will ignore the key request.
///
/// # Arguments
///
/// * `room_id` - The id of the room where the key is used in.
///
/// * `sender_key` - The curve25519 key of the sender that owns the key.
///
/// * `session_id` - The id that uniquely identifies the session.
pub async fn request_room_key(
&self,
event: &SyncMessageEvent<EncryptedEventContent>,
room_id: &RoomId,
) -> MegolmResult<(Option<OutgoingRequest>, OutgoingRequest)> {
let content = match &event.content {
EncryptedEventContent::MegolmV1AesSha2(c) => c,
_ => return Err(EventError::UnsupportedAlgorithm.into()),
};
Ok(self
.key_request_machine
.request_key(room_id, &content.sender_key, &content.session_id)
.await?)
}
/// Decrypt an event from a room timeline. /// Decrypt an event from a room timeline.
/// ///
/// # Arguments /// # Arguments

View File

@ -1346,7 +1346,7 @@ mod test {
#[async_test] #[async_test]
async fn key_request_saving() { async fn key_request_saving() {
let (_, store, _dir) = get_loaded_store().await; let (account, store, _dir) = get_loaded_store().await;
let id = Uuid::new_v4(); let id = Uuid::new_v4();
let info = RequestedKeyInfo { let info = RequestedKeyInfo {
@ -1357,6 +1357,7 @@ mod test {
}; };
let request = OutgoingKeyRequest { let request = OutgoingKeyRequest {
request_recipient: account.user_id().to_owned(),
request_id: id, request_id: id,
info: info.clone(), info: info.clone(),
sent_out: false, sent_out: false,
@ -1378,6 +1379,7 @@ mod test {
assert!(!store.get_unsent_key_requests().await.unwrap().is_empty()); assert!(!store.get_unsent_key_requests().await.unwrap().is_empty());
let request = OutgoingKeyRequest { let request = OutgoingKeyRequest {
request_recipient: account.user_id().to_owned(),
request_id: id, request_id: id,
info: info.clone(), info: info.clone(),
sent_out: true, sent_out: true,