From d4f0799e6c424ad8e629df0354d4aa8d362f9ab0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Damir=20Jeli=C4=87?= Date: Tue, 24 Mar 2020 17:25:01 +0100 Subject: [PATCH] crypto: Create inbound sessions from room key events. --- src/crypto/machine.rs | 41 ++++++++++++++++-------- src/crypto/memory_stores.rs | 62 +++++++++++++++++++++++++++++++++++++ src/crypto/mod.rs | 1 + 3 files changed, 91 insertions(+), 13 deletions(-) create mode 100644 src/crypto/memory_stores.rs diff --git a/src/crypto/machine.rs b/src/crypto/machine.rs index 8eb015fc..afb1727a 100644 --- a/src/crypto/machine.rs +++ b/src/crypto/machine.rs @@ -20,6 +20,7 @@ use std::result::Result as StdResult; use std::sync::Arc; use super::error::{OlmError, Result, SignatureError, VerificationResult}; +use super::memory_stores::GroupSessionStore; use super::olm::Account; #[cfg(feature = "sqlite-cryptostore")] use super::store::sqlite::SqliteStore; @@ -30,10 +31,10 @@ use crate::api; use api::r0::keys; use cjson; -use olm_rs::session::OlmMessage; -use olm_rs::utility::OlmUtility; -use serde_json::json; -use serde_json::Value; +use olm_rs::{ + inbound_group_session::OlmInboundGroupSession, session::OlmMessage, utility::OlmUtility, +}; +use serde_json::{json, Value}; use tokio::sync::Mutex; use tracing::{debug, info, instrument, trace, warn}; @@ -70,6 +71,8 @@ pub struct OlmMachine { /// Persists all the encrytpion keys so a client can resume the session /// without the need to create new keys. store: Box, + /// A cache of all the inbound group sessions we know about. + inbound_group_sessions: GroupSessionStore, } impl OlmMachine { @@ -86,6 +89,7 @@ impl OlmMachine { account: Arc::new(Mutex::new(Account::new())), uploaded_signed_key_count: None, store: Box::new(MemoryStore::new()), + inbound_group_sessions: GroupSessionStore::new(), }) } @@ -118,6 +122,7 @@ impl OlmMachine { account: Arc::new(Mutex::new(account)), uploaded_signed_key_count: None, store: Box::new(store), + inbound_group_sessions: GroupSessionStore::new(), }) } @@ -490,18 +495,27 @@ impl OlmMachine { OlmMessage::from_type_and_ciphertext(message_type.into(), ciphertext.body.clone()) .map_err(|_| OlmError::UnsupportedOlmType)?; - Ok(self + let decrypted_event = self .decrypt_olm_message(&event.sender.to_string(), &content.sender_key, message) - .await?) + .await?; + debug!("Decrypted a to-device event {:?}", decrypted_event); + self.handle_decrypted_to_device_event(&content.sender_key, &decrypted_event); + + Ok(decrypted_event) } else { warn!("Olm event doesn't contain a ciphertext for our key"); Err(OlmError::MissingCiphertext) } } - fn add_room_key(&self, event: &ToDeviceRoomKey) { + fn add_room_key(&mut self, sender_key: &str, event: &ToDeviceRoomKey) { match event.content.algorithm { - Algorithm::MegolmV1AesSha2 => {} + Algorithm::MegolmV1AesSha2 => { + // TODO check for all the valid fields. + let session = OlmInboundGroupSession::new(&event.content.session_key).unwrap(); + self.inbound_group_sessions + .add(&event.sender.to_string(), sender_key, session); + } _ => warn!( "Received room key with unsupported key algorithm {}", event.content.algorithm @@ -513,7 +527,11 @@ impl OlmMachine { // TODO } - fn handle_decrypted_to_device_event(&self, event: &EventResult) { + fn handle_decrypted_to_device_event( + &mut self, + sender_key: &str, + event: &EventResult, + ) { let event = if let EventResult::Ok(e) = event { e } else { @@ -522,7 +540,7 @@ impl OlmMachine { }; match event { - ToDeviceEvent::RoomKey(e) => self.add_room_key(e), + ToDeviceEvent::RoomKey(e) => self.add_room_key(sender_key, e), ToDeviceEvent::ForwardedRoomKey(e) => self.add_forwarded_room_key(e), _ => warn!("Received a unexpected encrypted to-device event"), } @@ -572,9 +590,6 @@ impl OlmMachine { continue; } }; - - debug!("Decrypted a to-device event {:?}", decrypted_event); - self.handle_decrypted_to_device_event(&decrypted_event); } ToDeviceEvent::RoomKeyRequest(e) => self.handle_room_key_request(e), ToDeviceEvent::KeyVerificationAccept(..) diff --git a/src/crypto/memory_stores.rs b/src/crypto/memory_stores.rs new file mode 100644 index 00000000..8ae5e94f --- /dev/null +++ b/src/crypto/memory_stores.rs @@ -0,0 +1,62 @@ +// 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 olm_rs::inbound_group_session::OlmInboundGroupSession; +use std::collections::HashMap; + +#[derive(Debug)] +pub struct GroupSessionStore { + entries: HashMap>>, +} + +impl GroupSessionStore { + pub fn new() -> Self { + GroupSessionStore { + entries: HashMap::new(), + } + } + + pub fn add( + &mut self, + room_id: &str, + sender_key: &str, + session: OlmInboundGroupSession, + ) -> bool { + if !self.entries.contains_key(room_id) { + self.entries.insert(room_id.to_owned(), HashMap::new()); + } + + let mut room_map = self.entries.get_mut(room_id).unwrap(); + + if !room_map.contains_key(sender_key) { + room_map.insert(sender_key.to_owned(), HashMap::new()); + } + + let mut sender_map = room_map.get_mut(sender_key).unwrap(); + let ret = sender_map.insert(session.session_id(), session); + + ret.is_some() + } + + pub fn get( + &self, + room_id: &str, + sender_key: &str, + session_id: &str, + ) -> Option<&OlmInboundGroupSession> { + self.entries + .get(room_id) + .and_then(|m| m.get(sender_key).and_then(|m| m.get(session_id))) + } +} diff --git a/src/crypto/mod.rs b/src/crypto/mod.rs index c1d8ed10..a16a4205 100644 --- a/src/crypto/mod.rs +++ b/src/crypto/mod.rs @@ -15,6 +15,7 @@ mod error; // TODO remove this. mod machine; +mod memory_stores; #[allow(dead_code)] mod olm; mod store;