627 lines
22 KiB
Rust
627 lines
22 KiB
Rust
// Copyright 2020 The Matrix.org Foundation C.I.C.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
use std::{
|
|
collections::{BTreeMap, BTreeSet, HashMap, HashSet},
|
|
sync::Arc,
|
|
};
|
|
|
|
use dashmap::DashMap;
|
|
use futures::future::join_all;
|
|
use matrix_sdk_common::{executor::spawn, uuid::Uuid};
|
|
use ruma::{
|
|
events::{
|
|
room::{encrypted::EncryptedEventContent, history_visibility::HistoryVisibility},
|
|
AnyMessageEventContent, AnyToDeviceEventContent, EventType,
|
|
},
|
|
serde::Raw,
|
|
to_device::DeviceIdOrAllDevices,
|
|
DeviceId, DeviceIdBox, RoomId, UserId,
|
|
};
|
|
use tracing::{debug, info, trace};
|
|
|
|
use crate::{
|
|
error::{EventError, MegolmResult, OlmResult},
|
|
olm::{Account, InboundGroupSession, OutboundGroupSession, Session, ShareInfo, ShareState},
|
|
store::{Changes, Result as StoreResult, Store},
|
|
Device, EncryptionSettings, OlmError, ToDeviceRequest,
|
|
};
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub(crate) struct GroupSessionCache {
|
|
store: Store,
|
|
sessions: Arc<DashMap<RoomId, OutboundGroupSession>>,
|
|
/// A map from the request id to the group session that the request belongs
|
|
/// to. Used to mark requests belonging to the session as shared.
|
|
sessions_being_shared: Arc<DashMap<Uuid, OutboundGroupSession>>,
|
|
}
|
|
|
|
impl GroupSessionCache {
|
|
pub(crate) fn new(store: Store) -> Self {
|
|
Self { store, sessions: Default::default(), sessions_being_shared: Default::default() }
|
|
}
|
|
|
|
pub(crate) fn insert(&self, session: OutboundGroupSession) {
|
|
self.sessions.insert(session.room_id().to_owned(), session);
|
|
}
|
|
|
|
/// Either get a session for the given room from the cache or load it from
|
|
/// the store.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `room_id` - The id of the room this session is used for.
|
|
pub async fn get_or_load(&self, room_id: &RoomId) -> StoreResult<Option<OutboundGroupSession>> {
|
|
// Get the cached session, if there isn't one load one from the store
|
|
// and put it in the cache.
|
|
if let Some(s) = self.sessions.get(room_id) {
|
|
Ok(Some(s.clone()))
|
|
} else if let Some(s) = self.store.get_outbound_group_sessions(room_id).await? {
|
|
for request_id in s.pending_request_ids() {
|
|
self.sessions_being_shared.insert(request_id, s.clone());
|
|
}
|
|
|
|
self.sessions.insert(room_id.clone(), s.clone());
|
|
|
|
Ok(Some(s))
|
|
} else {
|
|
Ok(None)
|
|
}
|
|
}
|
|
|
|
/// Get an outbound group session for a room, if one exists.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `room_id` - The id of the room for which we should get the outbound
|
|
/// group session.
|
|
fn get(&self, room_id: &RoomId) -> Option<OutboundGroupSession> {
|
|
self.sessions.get(room_id).map(|s| s.clone())
|
|
}
|
|
|
|
/// Get or load the session for the given room with the given session id.
|
|
///
|
|
/// This is the same as [get_or_load()](#method.get_or_load) but it will
|
|
/// filter out the session if it doesn't match the given session id.
|
|
pub async fn get_with_id(
|
|
&self,
|
|
room_id: &RoomId,
|
|
session_id: &str,
|
|
) -> StoreResult<Option<OutboundGroupSession>> {
|
|
Ok(self.get_or_load(room_id).await?.filter(|o| session_id == o.session_id()))
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct GroupSessionManager {
|
|
account: Account,
|
|
/// Store for the encryption keys.
|
|
/// Persists all the encryption keys so a client can resume the session
|
|
/// without the need to create new keys.
|
|
store: Store,
|
|
/// The currently active outbound group sessions.
|
|
sessions: GroupSessionCache,
|
|
}
|
|
|
|
impl GroupSessionManager {
|
|
const MAX_TO_DEVICE_MESSAGES: usize = 250;
|
|
|
|
pub(crate) fn new(account: Account, store: Store) -> Self {
|
|
Self { account, store: store.clone(), sessions: GroupSessionCache::new(store) }
|
|
}
|
|
|
|
pub async fn invalidate_group_session(&self, room_id: &RoomId) -> StoreResult<bool> {
|
|
if let Some(s) = self.sessions.get(room_id) {
|
|
s.invalidate_session();
|
|
|
|
let mut changes = Changes::default();
|
|
changes.outbound_group_sessions.push(s.clone());
|
|
self.store.save_changes(changes).await?;
|
|
|
|
Ok(true)
|
|
} else {
|
|
Ok(false)
|
|
}
|
|
}
|
|
|
|
pub async fn mark_request_as_sent(&self, request_id: &Uuid) -> StoreResult<()> {
|
|
if let Some((_, s)) = self.sessions.sessions_being_shared.remove(request_id) {
|
|
s.mark_request_as_sent(request_id);
|
|
|
|
let mut changes = Changes::default();
|
|
changes.outbound_group_sessions.push(s.clone());
|
|
self.store.save_changes(changes).await?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub fn get_outbound_group_session(&self, room_id: &RoomId) -> Option<OutboundGroupSession> {
|
|
self.sessions.get(room_id)
|
|
}
|
|
|
|
pub async fn encrypt(
|
|
&self,
|
|
room_id: &RoomId,
|
|
content: AnyMessageEventContent,
|
|
) -> MegolmResult<EncryptedEventContent> {
|
|
let session = if let Some(s) = self.sessions.get(room_id) {
|
|
s
|
|
} else {
|
|
panic!("Session wasn't created nor shared");
|
|
};
|
|
|
|
if session.expired() {
|
|
panic!("Session expired");
|
|
}
|
|
|
|
let content = session.encrypt(content).await;
|
|
|
|
let mut changes = Changes::default();
|
|
changes.outbound_group_sessions.push(session);
|
|
self.store.save_changes(changes).await?;
|
|
|
|
Ok(content)
|
|
}
|
|
|
|
/// Create a new outbound group session.
|
|
///
|
|
/// This also creates a matching inbound group session and saves that one in
|
|
/// the store.
|
|
pub async fn create_outbound_group_session(
|
|
&self,
|
|
room_id: &RoomId,
|
|
settings: EncryptionSettings,
|
|
) -> OlmResult<(OutboundGroupSession, InboundGroupSession)> {
|
|
let (outbound, inbound) = self
|
|
.account
|
|
.create_group_session_pair(room_id, settings)
|
|
.await
|
|
.map_err(|_| EventError::UnsupportedAlgorithm)?;
|
|
|
|
self.sessions.insert(outbound.clone());
|
|
Ok((outbound, inbound))
|
|
}
|
|
|
|
pub async fn get_or_create_outbound_session(
|
|
&self,
|
|
room_id: &RoomId,
|
|
settings: EncryptionSettings,
|
|
) -> OlmResult<(OutboundGroupSession, Option<InboundGroupSession>)> {
|
|
let outbound_session = self.sessions.get_or_load(room_id).await?;
|
|
|
|
// If there is no session or the session has expired or is invalid,
|
|
// create a new one.
|
|
if let Some(s) = outbound_session {
|
|
if s.expired() || s.invalidated() {
|
|
self.create_outbound_group_session(room_id, settings)
|
|
.await
|
|
.map(|(o, i)| (o, i.into()))
|
|
} else {
|
|
Ok((s, None))
|
|
}
|
|
} else {
|
|
self.create_outbound_group_session(room_id, settings).await.map(|(o, i)| (o, i.into()))
|
|
}
|
|
}
|
|
|
|
/// Encrypt the given content for the given devices and create a to-device
|
|
/// requests that sends the encrypted content to them.
|
|
async fn encrypt_session_for(
|
|
content: AnyToDeviceEventContent,
|
|
devices: Vec<Device>,
|
|
message_index: u32,
|
|
) -> OlmResult<(
|
|
Uuid,
|
|
ToDeviceRequest,
|
|
BTreeMap<UserId, BTreeMap<DeviceIdBox, ShareInfo>>,
|
|
Vec<Session>,
|
|
)> {
|
|
let mut messages = BTreeMap::new();
|
|
let mut changed_sessions = Vec::new();
|
|
let mut share_infos = BTreeMap::new();
|
|
|
|
let encrypt = |device: Device, content: AnyToDeviceEventContent| async move {
|
|
let mut message = BTreeMap::new();
|
|
let mut share_infos = BTreeMap::new();
|
|
|
|
let encrypted = device.encrypt(content.clone()).await;
|
|
|
|
let used_session = match encrypted {
|
|
Ok((session, encrypted)) => {
|
|
message
|
|
.entry(device.user_id().to_owned())
|
|
.or_insert_with(BTreeMap::new)
|
|
.insert(
|
|
DeviceIdOrAllDevices::DeviceId(device.device_id().into()),
|
|
Raw::from(AnyToDeviceEventContent::RoomEncrypted(encrypted)),
|
|
);
|
|
share_infos
|
|
.entry(device.user_id().to_owned())
|
|
.or_insert_with(BTreeMap::new)
|
|
.insert(
|
|
device.device_id().to_owned(),
|
|
ShareInfo {
|
|
sender_key: session.sender_key().to_owned(),
|
|
message_index,
|
|
},
|
|
);
|
|
|
|
Some(session)
|
|
}
|
|
// TODO we'll want to create m.room_key.withheld here.
|
|
Err(OlmError::MissingSession)
|
|
| Err(OlmError::EventError(EventError::MissingSenderKey)) => None,
|
|
Err(e) => return Err(e),
|
|
};
|
|
|
|
Ok((used_session, share_infos, message))
|
|
};
|
|
|
|
let tasks: Vec<_> =
|
|
devices.iter().map(|d| spawn(encrypt(d.clone(), content.clone()))).collect();
|
|
|
|
let results = join_all(tasks).await;
|
|
|
|
for result in results {
|
|
let (used_session, infos, message) = result.expect("Encryption task panicked")?;
|
|
|
|
if let Some(session) = used_session {
|
|
changed_sessions.push(session);
|
|
}
|
|
|
|
for (user, device_messages) in message.into_iter() {
|
|
messages.entry(user).or_insert_with(BTreeMap::new).extend(device_messages);
|
|
}
|
|
|
|
for (user, infos) in infos.into_iter() {
|
|
share_infos.entry(user).or_insert_with(BTreeMap::new).extend(infos);
|
|
}
|
|
}
|
|
|
|
let id = Uuid::new_v4();
|
|
|
|
let request =
|
|
ToDeviceRequest { event_type: EventType::RoomEncrypted, txn_id: id, messages };
|
|
|
|
trace!(
|
|
recipient_count = request.message_count(),
|
|
transaction_id = ?id,
|
|
"Created a to-device request carrying a room_key"
|
|
);
|
|
|
|
Ok((id, request, share_infos, changed_sessions))
|
|
}
|
|
|
|
/// Given a list of user and an outbound session, return the list of users
|
|
/// and their devices that this session should be shared with.
|
|
///
|
|
/// Returns a boolean indicating whether the session needs to be rotated and
|
|
/// the list of users/devices that should receive the session.
|
|
pub async fn collect_session_recipients(
|
|
&self,
|
|
users: impl Iterator<Item = &UserId>,
|
|
history_visibility: HistoryVisibility,
|
|
outbound: &OutboundGroupSession,
|
|
) -> OlmResult<(bool, HashMap<UserId, Vec<Device>>)> {
|
|
let users: HashSet<&UserId> = users.collect();
|
|
let mut devices: HashMap<UserId, Vec<Device>> = HashMap::new();
|
|
|
|
debug!(
|
|
users = ?users,
|
|
history_visibility = ?history_visibility,
|
|
session_id = outbound.session_id(),
|
|
"Calculating group session recipients"
|
|
);
|
|
|
|
let users_shared_with: HashSet<UserId> =
|
|
outbound.shared_with_set.iter().map(|k| k.key().clone()).collect();
|
|
|
|
let users_shared_with: HashSet<&UserId> = users_shared_with.iter().collect();
|
|
|
|
// A user left if a user is missing from the set of users that should
|
|
// get the session but is in the set of users that received the session.
|
|
let user_left = !users_shared_with.difference(&users).collect::<HashSet<_>>().is_empty();
|
|
|
|
let visibility_changed = outbound.settings().history_visibility != history_visibility;
|
|
|
|
// To protect the room history we need to rotate the session if either:
|
|
//
|
|
// 1. Any user left the room.
|
|
// 2. Any of the users' devices got deleted or blacklisted.
|
|
// 3. The history visibility changed.
|
|
//
|
|
// This is calculated in the following code and stored in this variable.
|
|
let mut should_rotate = user_left || visibility_changed;
|
|
|
|
for user_id in users {
|
|
let user_devices = self.store.get_user_devices(user_id).await?;
|
|
let non_blacklisted_devices: Vec<Device> =
|
|
user_devices.devices().filter(|d| !d.is_blacklisted()).collect();
|
|
|
|
// If we haven't already concluded that the session should be
|
|
// rotated for other reasons, we also need to check whether any
|
|
// of the devices in the session got deleted or blacklisted in the
|
|
// meantime. If so, we should also rotate the session.
|
|
if !should_rotate {
|
|
// Device IDs that should receive this session
|
|
let non_blacklisted_device_ids: HashSet<&DeviceId> =
|
|
non_blacklisted_devices.iter().map(|d| d.device_id()).collect();
|
|
|
|
if let Some(shared) = outbound.shared_with_set.get(user_id) {
|
|
#[allow(clippy::map_clone)]
|
|
// Devices that received this session
|
|
let shared: HashSet<DeviceIdBox> =
|
|
shared.iter().map(|d| d.key().clone()).collect();
|
|
let shared: HashSet<&DeviceId> = shared.iter().map(|d| d.as_ref()).collect();
|
|
|
|
// The set difference between
|
|
//
|
|
// 1. Devices that had previously received the session, and
|
|
// 2. Devices that would now receive the session
|
|
//
|
|
// represents newly deleted or blacklisted devices. If this
|
|
// set is non-empty, we must rotate.
|
|
let newly_deleted_or_blacklisted =
|
|
shared.difference(&non_blacklisted_device_ids).collect::<HashSet<_>>();
|
|
|
|
if !newly_deleted_or_blacklisted.is_empty() {
|
|
should_rotate = true;
|
|
}
|
|
};
|
|
}
|
|
|
|
devices.entry(user_id.clone()).or_insert_with(Vec::new).extend(non_blacklisted_devices);
|
|
}
|
|
|
|
debug!(
|
|
should_rotate = should_rotate,
|
|
session_id = outbound.session_id(),
|
|
"Done calculating group session recipients"
|
|
);
|
|
|
|
Ok((should_rotate, devices))
|
|
}
|
|
|
|
pub async fn encrypt_request(
|
|
chunk: Vec<Device>,
|
|
content: AnyToDeviceEventContent,
|
|
outbound: OutboundGroupSession,
|
|
message_index: u32,
|
|
being_shared: Arc<DashMap<Uuid, OutboundGroupSession>>,
|
|
) -> OlmResult<Vec<Session>> {
|
|
let (id, request, share_infos, used_sessions) =
|
|
Self::encrypt_session_for(content.clone(), chunk, message_index).await?;
|
|
|
|
if !request.messages.is_empty() {
|
|
outbound.add_request(id, request.into(), share_infos);
|
|
being_shared.insert(id, outbound.clone());
|
|
}
|
|
|
|
Ok(used_sessions)
|
|
}
|
|
|
|
pub(crate) fn session_cache(&self) -> GroupSessionCache {
|
|
self.sessions.clone()
|
|
}
|
|
|
|
/// Get to-device requests to share a group session with users in a room.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// `room_id` - The room id of the room where the group session will be
|
|
/// used.
|
|
///
|
|
/// `users` - The list of users that should receive the group session.
|
|
///
|
|
/// `encryption_settings` - The settings that should be used for the group
|
|
/// session.
|
|
pub async fn share_group_session(
|
|
&self,
|
|
room_id: &RoomId,
|
|
users: impl Iterator<Item = &UserId>,
|
|
encryption_settings: impl Into<EncryptionSettings>,
|
|
) -> OlmResult<Vec<Arc<ToDeviceRequest>>> {
|
|
debug!(room_id = room_id.as_str(), "Checking if a room key needs to be shared",);
|
|
|
|
let encryption_settings = encryption_settings.into();
|
|
let history_visibility = encryption_settings.history_visibility.clone();
|
|
let mut changes = Changes::default();
|
|
|
|
let (outbound, inbound) =
|
|
self.get_or_create_outbound_session(room_id, encryption_settings.clone()).await?;
|
|
|
|
if let Some(inbound) = inbound {
|
|
changes.outbound_group_sessions.push(outbound.clone());
|
|
changes.inbound_group_sessions.push(inbound);
|
|
}
|
|
|
|
let (should_rotate, devices) =
|
|
self.collect_session_recipients(users, history_visibility, &outbound).await?;
|
|
|
|
let outbound = if should_rotate {
|
|
let old_session_id = outbound.session_id();
|
|
|
|
let (outbound, inbound) =
|
|
self.create_outbound_group_session(room_id, encryption_settings).await?;
|
|
changes.outbound_group_sessions.push(outbound.clone());
|
|
changes.inbound_group_sessions.push(inbound);
|
|
|
|
debug!(
|
|
room_id = room_id.as_str(),
|
|
old_session_id = old_session_id,
|
|
session_id = outbound.session_id(),
|
|
"A user/device has left the group since we last sent a message, \
|
|
rotating the outbound session.",
|
|
);
|
|
|
|
outbound
|
|
} else {
|
|
outbound
|
|
};
|
|
|
|
let devices: Vec<Device> = devices
|
|
.into_iter()
|
|
.map(|(_, d)| {
|
|
d.into_iter()
|
|
.filter(|d| matches!(outbound.is_shared_with(d), ShareState::NotShared))
|
|
})
|
|
.flatten()
|
|
.collect();
|
|
|
|
let key_content = outbound.as_content().await;
|
|
let message_index = outbound.message_index().await;
|
|
|
|
if !devices.is_empty() {
|
|
let users = devices.iter().fold(BTreeMap::new(), |mut acc, d| {
|
|
acc.entry(d.user_id()).or_insert_with(BTreeSet::new).insert(d.device_id());
|
|
acc
|
|
});
|
|
|
|
info!(
|
|
index = message_index,
|
|
users = ?users,
|
|
room_id = room_id.as_str(),
|
|
"Sharing an outbound group session",
|
|
);
|
|
}
|
|
|
|
let tasks: Vec<_> = devices
|
|
.chunks(Self::MAX_TO_DEVICE_MESSAGES)
|
|
.map(|chunk| {
|
|
spawn(Self::encrypt_request(
|
|
chunk.to_vec(),
|
|
key_content.clone(),
|
|
outbound.clone(),
|
|
message_index,
|
|
self.sessions.sessions_being_shared.clone(),
|
|
))
|
|
})
|
|
.collect();
|
|
|
|
for result in join_all(tasks).await {
|
|
let used_sessions: OlmResult<Vec<Session>> = result.expect("Encryption task panicked");
|
|
|
|
changes.sessions.extend(used_sessions?);
|
|
}
|
|
|
|
let requests = outbound.pending_requests();
|
|
|
|
debug!(
|
|
room_id = room_id.as_str(),
|
|
session_id = outbound.session_id(),
|
|
request_count = requests.len(),
|
|
"Done generating to-device requests for a room key share"
|
|
);
|
|
|
|
if requests.is_empty() {
|
|
debug!(
|
|
room_id = room_id.as_str(),
|
|
session_id = outbound.session_id(),
|
|
"The outbound group session doesn't need to be shared with \
|
|
anyone, marking as shared",
|
|
);
|
|
outbound.mark_as_shared();
|
|
}
|
|
|
|
let session_count = changes.sessions.len();
|
|
|
|
self.store.save_changes(changes).await?;
|
|
|
|
debug!(
|
|
room_id = room_id.as_str(),
|
|
session_id = outbound.session_id(),
|
|
session_count = session_count,
|
|
"Stored the changed sessions after encrypting an room key"
|
|
);
|
|
|
|
Ok(requests)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod test {
|
|
use matrix_sdk_common::uuid::Uuid;
|
|
use matrix_sdk_test::response_from_file;
|
|
use ruma::{
|
|
api::{
|
|
client::r0::keys::{claim_keys, get_keys},
|
|
IncomingResponse,
|
|
},
|
|
room_id, user_id, DeviceIdBox, UserId,
|
|
};
|
|
use serde_json::Value;
|
|
|
|
use crate::{EncryptionSettings, OlmMachine};
|
|
|
|
fn alice_id() -> UserId {
|
|
user_id!("@alice:example.org")
|
|
}
|
|
|
|
fn alice_device_id() -> DeviceIdBox {
|
|
"JLAFKJWSCS".into()
|
|
}
|
|
|
|
fn keys_query_response() -> get_keys::Response {
|
|
let data = include_bytes!("../../benches/keys_query.json");
|
|
let data: Value = serde_json::from_slice(data).unwrap();
|
|
let data = response_from_file(&data);
|
|
get_keys::Response::try_from_http_response(data)
|
|
.expect("Can't parse the keys upload response")
|
|
}
|
|
|
|
fn keys_claim_response() -> claim_keys::Response {
|
|
let data = include_bytes!("../../benches/keys_claim.json");
|
|
let data: Value = serde_json::from_slice(data).unwrap();
|
|
let data = response_from_file(&data);
|
|
claim_keys::Response::try_from_http_response(data)
|
|
.expect("Can't parse the keys upload response")
|
|
}
|
|
|
|
async fn machine() -> OlmMachine {
|
|
let keys_query = keys_query_response();
|
|
let keys_claim = keys_claim_response();
|
|
let uuid = Uuid::new_v4();
|
|
|
|
let machine = OlmMachine::new(&alice_id(), &alice_device_id());
|
|
|
|
machine.mark_request_as_sent(&uuid, &keys_query).await.unwrap();
|
|
machine.mark_request_as_sent(&uuid, &keys_claim).await.unwrap();
|
|
|
|
machine
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_sharing() {
|
|
let machine = machine().await;
|
|
let room_id = room_id!("!test:localhost");
|
|
let keys_claim = keys_claim_response();
|
|
|
|
let users: Vec<_> = keys_claim.one_time_keys.keys().collect();
|
|
|
|
let requests = machine
|
|
.share_group_session(&room_id, users.clone().into_iter(), EncryptionSettings::default())
|
|
.await
|
|
.unwrap();
|
|
|
|
let event_count: usize = requests.iter().map(|r| r.message_count()).sum();
|
|
|
|
// The keys claim response has a couple of one-time keys with invalid
|
|
// signatures, thus only 148 sessions are actually created, we check
|
|
// that all 148 valid sessions get an room key.
|
|
assert_eq!(event_count, 148);
|
|
}
|
|
}
|