Skip to content

Commit

Permalink
Prepare our /keys/claim response logic to handle multiple one-time keys
Browse files Browse the repository at this point in the history 
This is useful if we ever decide to switch to X3DH for the session
establishment. It also refactors a bit the /keys/claim response handling
method.
  • Loading branch information
@poljar
poljar committed Nov 22, 2023
1 parent 06c9a1a commit 8a0dbdb
Show file tree
Hide file tree
Showing 6 changed files with 287 additions and 184 deletions.
4 changes: 2 additions & 2 deletions crates/matrix-sdk-crypto/src/dehydrated_devices.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -491,15 +491,15 @@ mod tests {
.await
.expect("We should be able to create a request to upload a dehydrated device");

let (_key_id, one_time_key) = request
let (key_id, one_time_key) = request
.one_time_keys
.pop_first()
.expect("The dehydrated device creation request should contain a one-time key");

// Ensure that we know about the public keys of the dehydrated device.
receive_device_keys(&alice, user_id(), &request.device_id, request.device_keys).await;
// Create a 1-to-1 Olm session with the dehydrated device.
create_session(&alice, user_id(), &request.device_id, one_time_key).await;
create_session(&alice, user_id(), &request.device_id, key_id, one_time_key).await;

// Send a room key to the dehydrated device.
let (event, group_session) = send_room_key(&alice, room_id, user_id()).await;
Expand Down
7 changes: 7 additions & 0 deletions crates/matrix-sdk-crypto/src/error.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -242,6 +242,13 @@ pub enum SessionCreationError {
)]
OneTimeKeyNotSigned(OwnedUserId, OwnedDeviceId),

/// The signed one-time key is missing.
#[error(
"Tried to create a new Olm session for {0} {1}, but the signed \
one-time key is missing"
)]
OneTimeKeyMissing(OwnedUserId, OwnedDeviceId),

/// The one-time key algorithm is unsupported.
#[error(
"Tried to create a new Olm session for {0} {1}, but the one-time \
Expand Down
39 changes: 27 additions & 12 deletions crates/matrix-sdk-crypto/src/machine.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2078,7 +2078,9 @@ pub(crate) mod tests {
use ruma::{
api::{
client::{
keys::{get_keys, get_keys::v3::Response as KeyQueryResponse, upload_keys},
keys::{
claim_keys, get_keys, get_keys::v3::Response as KeyQueryResponse, upload_keys,
},
sync::sync_events::DeviceLists,
to_device::send_event_to_device::v3::Response as ToDeviceResponse,
},
Expand Down Expand Up @@ -2235,15 +2237,20 @@ pub(crate) mod tests {
bob: &UserId,
use_fallback_key: bool,
) -> (OlmMachine, OlmMachine) {
let (alice, bob, one_time_keys) = get_machine_pair(alice, bob, use_fallback_key).await;
let (alice, bob, mut one_time_keys) = get_machine_pair(alice, bob, use_fallback_key).await;

let one_time_key = one_time_keys.values().next().unwrap();
let (device_key_id, one_time_key) = one_time_keys.pop_first().unwrap();

let one_time_keys = BTreeMap::from([(
(bob.user_id().to_owned(), bob.device_id().to_owned()),
one_time_key,
bob.user_id().to_owned(),
BTreeMap::from([(
bob.device_id().to_owned(),
BTreeMap::from([(device_key_id, one_time_key)]),
)]),
)]);
alice.inner.session_manager.create_sessions(&one_time_keys).await.unwrap();

let response = claim_keys::v3::Response::new(one_time_keys);
alice.inner.session_manager.create_sessions(&response).await.unwrap();

(alice, bob)
}
Expand Down Expand Up @@ -2540,23 +2547,29 @@ pub(crate) mod tests {
machine: &OlmMachine,
user_id: &UserId,
device_id: &DeviceId,
key_id: OwnedDeviceKeyId,
one_time_key: Raw<OneTimeKey>,
) {
let one_time_keys =
BTreeMap::from([((user_id.to_owned(), device_id.to_owned()), &one_time_key)]);
machine.inner.session_manager.create_sessions(&one_time_keys).await.unwrap();
let one_time_keys = BTreeMap::from([(
user_id.to_owned(),
BTreeMap::from([(device_id.to_owned(), BTreeMap::from([(key_id, one_time_key)]))]),
)]);

let response = claim_keys::v3::Response::new(one_time_keys);
machine.inner.session_manager.create_sessions(&response).await.unwrap();
}

#[async_test]
async fn test_session_creation() {
let (alice_machine, bob_machine, mut one_time_keys) =
get_machine_pair(alice_id(), user_id(), false).await;
let (_key_id, one_time_key) = one_time_keys.pop_first().unwrap();
let (key_id, one_time_key) = one_time_keys.pop_first().unwrap();

create_session(
&alice_machine,
bob_machine.user_id(),
bob_machine.device_id(),
key_id,
one_time_key,
)
.await;
Expand All @@ -2577,7 +2590,7 @@ pub(crate) mod tests {
async fn test_getting_most_recent_session() {
let (alice_machine, bob_machine, mut one_time_keys) =
get_machine_pair(alice_id(), user_id(), false).await;
let (_key_id, one_time_key) = one_time_keys.pop_first().unwrap();
let (key_id, one_time_key) = one_time_keys.pop_first().unwrap();

let device = alice_machine
.get_device(bob_machine.user_id(), bob_machine.device_id(), None)
Expand All @@ -2591,17 +2604,19 @@ pub(crate) mod tests {
&alice_machine,
bob_machine.user_id(),
bob_machine.device_id(),
key_id,
one_time_key.to_owned(),
)
.await;

for _ in 0..10 {
let (_key_id, one_time_key) = one_time_keys.pop_first().unwrap();
let (key_id, one_time_key) = one_time_keys.pop_first().unwrap();

create_session(
&alice_machine,
bob_machine.user_id(),
bob_machine.device_id(),
key_id,
one_time_key.to_owned(),
)
.await;
Expand Down
112 changes: 73 additions & 39 deletions crates/matrix-sdk-crypto/src/olm/account.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,11 @@ use crate::{
OlmError, SignatureError,
};

#[derive(Debug)]
enum PrekeyBundle {
Olm3DH { key: SignedKey },
}

#[derive(Debug, Clone)]
pub(crate) enum SessionType {
New(Session),
Expand Down Expand Up @@ -835,6 +840,43 @@ impl Account {
}
}

#[instrument(
skip_all,
fields(
user_id = %device.user_id(),
device_id = %device.device_id(),
algorithms = ?device.algorithms()
)
)]
fn find_pre_key_bundle(
device: &ReadOnlyDevice,
key_map: &BTreeMap<OwnedDeviceKeyId, Raw<ruma::encryption::OneTimeKey>>,
) -> Result<PrekeyBundle, SessionCreationError> {
let mut keys = key_map.iter();

let first_key = keys.next().ok_or_else(|| {
SessionCreationError::OneTimeKeyMissing(
device.user_id().to_owned(),
device.device_id().into(),
)
})?;

let first_key_id = first_key.0.to_owned();
let first_key = OneTimeKey::deserialize(first_key_id.algorithm(), first_key.1)?;

let ret = match first_key {
OneTimeKey::SignedKey(key) => Ok(PrekeyBundle::Olm3DH { key }),
_ => Err(SessionCreationError::OneTimeKeyUnknown(
device.user_id().to_owned(),
device.device_id().into(),
)),
};

trace!(result = ?ret, "Finished searching for a valid pre-key bundle");

ret
}

/// Create a new session with another account given a one-time key and a
/// device.
///
Expand All @@ -850,45 +892,39 @@ impl Account {
pub fn create_outbound_session(
&self,
device: &ReadOnlyDevice,
one_time_key: &Raw<ruma::encryption::OneTimeKey>,
key_map: &BTreeMap<OwnedDeviceKeyId, Raw<ruma::encryption::OneTimeKey>>,
) -> Result<Session, SessionCreationError> {
let one_time_key: SignedKey = match one_time_key.deserialize_as() {
Ok(OneTimeKey::SignedKey(k)) => k,
Ok(OneTimeKey::Key(_)) => {
return Err(SessionCreationError::OneTimeKeyNotSigned(
device.user_id().to_owned(),
device.device_id().into(),
));
}
Ok(_) => {
return Err(SessionCreationError::OneTimeKeyUnknown(
device.user_id().to_owned(),
device.device_id().into(),
));
}
Err(e) => return Err(SessionCreationError::InvalidJson(e)),
};
let pre_key_bundle = Self::find_pre_key_bundle(device, key_map)?;

match pre_key_bundle {
PrekeyBundle::Olm3DH { key } => {
device.verify_one_time_key(&key).map_err(|error| {
SessionCreationError::InvalidSignature {
signing_key: device.ed25519_key(),
one_time_key: key.clone(),
error,
}
})?;

device.verify_one_time_key(&one_time_key).map_err(|error| {
SessionCreationError::InvalidSignature {
signing_key: device.ed25519_key(),
one_time_key: one_time_key.clone(),
error,
let identity_key = device.curve25519_key().ok_or_else(|| {
SessionCreationError::DeviceMissingCurveKey(
device.user_id().to_owned(),
device.device_id().into(),
)
})?;

let is_fallback = key.fallback();
let one_time_key = key.key();
let config = device.olm_session_config();

Ok(self.create_outbound_session_helper(
config,
identity_key,
one_time_key,
is_fallback,
))
}
})?;

let identity_key = device.curve25519_key().ok_or_else(|| {
SessionCreationError::DeviceMissingCurveKey(
device.user_id().to_owned(),
device.device_id().into(),
)
})?;

let is_fallback = one_time_key.fallback();
let one_time_key = one_time_key.key();
let config = device.olm_session_config();

Ok(self.create_outbound_session_helper(config, identity_key, one_time_key, is_fallback))
}
}

/// Create a new session with another account given a pre-key Olm message.
Expand Down Expand Up @@ -949,11 +985,9 @@ impl Account {

other.generate_one_time_keys_helper(1);
let one_time_map = other.signed_one_time_keys();
let one_time = one_time_map.values().next().unwrap();

let device = ReadOnlyDevice::from_account(other);

let mut our_session = self.create_outbound_session(&device, one_time).unwrap();
let mut our_session = self.create_outbound_session(&device, &one_time_map).unwrap();

other.mark_keys_as_published();

Expand Down
Loading

0 comments on commit 8a0dbdb

Please sign in to comment.