crypto: Create inbound sessions from room key events.

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::{
-use olm_rs::utility::OlmUtility;
+    inbound_group_session::OlmInboundGroupSession, session::OlmMessage, utility::OlmUtility,
-use serde_json::json;
+};
-use serde_json::Value;
+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<dyn CryptoStore>,
+    /// 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<ToDeviceEvent>) {
+    fn handle_decrypted_to_device_event(
+        &mut self,
+        sender_key: &str,
+        event: &EventResult<ToDeviceEvent>,
+    ) {
         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(..)

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<String, HashMap<String, HashMap<String, OlmInboundGroupSession>>>,
+}
+
+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)))
+    }
+}

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;