Use inbox to handle save and load ephemeral key from store (#81)

* feat: move inbox store out from context

* chore: clear boundary between context and inbox to store dependencies

core/conversations/src/context.rs

@@ -1,5 +1,5 @@
-use std::rc::Rc;
 use std::sync::Arc;
+use std::{cell::RefCell, rc::Rc};
 
 use crypto::Identity;
 use double_ratchets::{RatchetState, restore_ratchet_state};
@@ -20,8 +20,8 @@ pub use crate::inbox::Introduction;
 // Ctx manages lifetimes of objects to process and generate payloads.
 pub struct Context<T: ChatStore> {
     _identity: Rc<Identity>,
-    inbox: Inbox,
+    inbox: Inbox<T>,
-    store: T,
+    store: Rc<RefCell<T>>,
 }
 
 impl<T: ChatStore> Context<T> {
@@ -29,20 +29,21 @@ impl<T: ChatStore> Context<T> {
     ///
     /// If an identity exists in storage, it will be restored.
     /// Otherwise, a new identity will be created with the given name and saved.
-    pub fn new_from_store(name: impl Into<String>, mut store: T) -> Result<Self, ChatError> {
+    pub fn new_from_store(name: impl Into<String>, store: T) -> Result<Self, ChatError> {
         let name = name.into();
+        let store = Rc::new(RefCell::new(store));
 
         // Load or create identity
-        let identity = if let Some(identity) = store.load_identity()? {
+        let identity = if let Some(identity) = store.borrow().load_identity()? {
             identity
         } else {
             let identity = Identity::new(&name);
-            store.save_identity(&identity)?;
+            store.borrow_mut().save_identity(&identity)?;
             identity
         };
 
         let identity = Rc::new(identity);
-        let inbox = Inbox::new(Rc::clone(&identity));
+        let inbox = Inbox::new(Rc::clone(&identity), Rc::clone(&store));
 
         Ok(Self {
             _identity: identity,
@@ -54,15 +55,17 @@ impl<T: ChatStore> Context<T> {
     /// Creates a new in-memory Context (for testing).
     ///
     /// Uses in-memory SQLite database. Each call creates a new isolated database.
-    pub fn new_with_name(name: impl Into<String>, mut chat_store: T) -> Self {
+    pub fn new_with_name(name: impl Into<String>, chat_store: T) -> Self {
         let name = name.into();
         let identity = Identity::new(&name);
+        let chat_store = Rc::new(RefCell::new(chat_store));
         chat_store
+            .borrow_mut()
             .save_identity(&identity)
             .expect("in-memory storage should not fail");
 
         let identity = Rc::new(identity);
-        let inbox = Inbox::new(Rc::clone(&identity));
+        let inbox = Inbox::new(Rc::clone(&identity), Rc::clone(&chat_store));
 
         Self {
             _identity: identity,
@@ -85,7 +88,7 @@ impl<T: ChatStore> Context<T> {
             .invite_to_private_convo(remote_bundle, content)
             .unwrap_or_else(|_| todo!("Log/Surface Error"));
 
-        let remote_id = Inbox::inbox_identifier_for_key(*remote_bundle.installation_key());
+        let remote_id = Inbox::<T>::inbox_identifier_for_key(*remote_bundle.installation_key());
         let payload_bytes = payloads
             .into_iter()
             .map(|p| p.into_envelope(remote_id.clone()))
@@ -96,7 +99,7 @@ impl<T: ChatStore> Context<T> {
     }
 
     pub fn list_conversations(&self) -> Result<Vec<ConversationIdOwned>, ChatError> {
-        let records = self.store.load_conversations()?;
+        let records = self.store.borrow().load_conversations()?;
         Ok(records
             .into_iter()
             .map(|r| Arc::from(r.local_convo_id.as_str()))
@@ -114,7 +117,7 @@ impl<T: ChatStore> Context<T> {
             Conversation::Private(mut convo) => {
                 let payloads = convo.send_message(content)?;
                 let remote_id = convo.remote_id();
-                convo.save_ratchet_state(&mut self.store)?;
+                convo.save_ratchet_state::<T>(&mut *self.store.borrow_mut())?;
 
                 Ok(payloads
                     .into_iter()
@@ -133,7 +136,7 @@ impl<T: ChatStore> Context<T> {
         let enc = EncryptedPayload::decode(env.payload)?;
         match convo_id {
             c if c == self.inbox.id() => self.dispatch_to_inbox(enc),
-            c if self.store.has_conversation(&c)? => self.dispatch_to_convo(&c, enc),
+            c if self.store.borrow().has_conversation(&c)? => self.dispatch_to_convo(&c, enc),
             _ => Ok(None),
         }
     }
@@ -143,20 +146,16 @@ impl<T: ChatStore> Context<T> {
         &mut self,
         enc_payload: EncryptedPayload,
     ) -> Result<Option<ContentData>, ChatError> {
-        // Look up the ephemeral key from storage
+        let public_key_hex = Inbox::<T>::extract_ephemeral_key_hex(&enc_payload)?;
-        let key_hex = Inbox::extract_ephemeral_key_hex(&enc_payload)?;
+        let (convo, content) = self.inbox.handle_frame(enc_payload, &public_key_hex)?;
-        let ephemeral_key = self
-            .store
-            .load_ephemeral_key(&key_hex)?
-            .ok_or(ChatError::UnknownEphemeralKey())?;
-
-        let (convo, content) = self.inbox.handle_frame(&ephemeral_key, enc_payload)?;
 
         match convo {
             Conversation::Private(convo) => self.persist_convo(&convo)?,
         };
 
-        self.store.remove_ephemeral_key(&key_hex)?;
+        self.store
+            .borrow_mut()
+            .remove_ephemeral_key(&public_key_hex)?;
         Ok(content)
     }
 
@@ -171,16 +170,14 @@ impl<T: ChatStore> Context<T> {
         match convo {
             Conversation::Private(mut convo) => {
                 let result = convo.handle_frame(enc_payload)?;
-                convo.save_ratchet_state(&mut self.store)?;
+                convo.save_ratchet_state(&mut *self.store.borrow_mut())?;
                 Ok(result)
             }
         }
     }
 
     pub fn create_intro_bundle(&mut self) -> Result<Vec<u8>, ChatError> {
-        let (intro, public_key_hex, private_key) = self.inbox.create_intro_bundle();
+        let intro = self.inbox.create_intro_bundle()?;
-        self.store
-            .save_ephemeral_key(&public_key_hex, &private_key)?;
         Ok(intro.into())
     }
 
@@ -188,13 +185,20 @@ impl<T: ChatStore> Context<T> {
     fn load_convo(&self, convo_id: ConversationId) -> Result<Conversation, ChatError> {
         let record = self
             .store
+            .borrow()
             .load_conversation(convo_id)?
             .ok_or_else(|| ChatError::NoConvo(convo_id.into()))?;
 
         match record.kind {
             ConversationKind::PrivateV1 => {
-                let dr_record = self.store.load_ratchet_state(&record.local_convo_id)?;
+                let dr_record = self
-                let skipped_keys = self.store.load_skipped_keys(&record.local_convo_id)?;
+                    .store
+                    .borrow()
+                    .load_ratchet_state(&record.local_convo_id)?;
+                let skipped_keys = self
+                    .store
+                    .borrow()
+                    .load_skipped_keys(&record.local_convo_id)?;
                 let dr_state: RatchetState = restore_ratchet_state(dr_record, skipped_keys);
 
                 Ok(Conversation::Private(PrivateV1Convo::new(
@@ -217,8 +221,8 @@ impl<T: ChatStore> Context<T> {
             remote_convo_id: convo.remote_id(),
             kind: convo.convo_type(),
         };
-        self.store.save_conversation(&convo_info)?;
+        self.store.borrow_mut().save_conversation(&convo_info)?;
-        convo.save_ratchet_state(&mut self.store)?;
+        convo.save_ratchet_state(&mut *self.store.borrow_mut())?;
         Ok(Arc::from(convo.id()))
     }
 }
@@ -325,7 +329,7 @@ mod tests {
         let content = alice.handle_payload(&payload.data).unwrap().unwrap();
         assert!(content.is_new_convo);
 
-        let convos = alice.store.load_conversations().unwrap();
+        let convos = alice.store.borrow().load_conversations().unwrap();
         assert_eq!(convos.len(), 1);
         assert_eq!(convos[0].kind.as_str(), "private_v1");
     }

core/conversations/src/inbox/handler.rs

@@ -3,7 +3,9 @@ use chat_proto::logoschat::encryption::EncryptedPayload;
 use prost::Message;
 use prost::bytes::Bytes;
 use rand_core::OsRng;
+use std::cell::RefCell;
 use std::rc::Rc;
+use storage::EphemeralKeyStore;
 
 use crypto::{PrekeyBundle, SymmetricKey32};
 
@@ -21,12 +23,13 @@ fn delivery_address_for_installation(_: PublicKey) -> String {
     "delivery_address".into()
 }
 
-pub struct Inbox {
+pub struct Inbox<S: EphemeralKeyStore> {
     ident: Rc<Identity>,
     local_convo_id: String,
+    store: Rc<RefCell<S>>,
 }
 
-impl std::fmt::Debug for Inbox {
+impl<S: EphemeralKeyStore> std::fmt::Debug for Inbox<S> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("Inbox")
             .field("ident", &self.ident)
@@ -35,25 +38,30 @@ impl std::fmt::Debug for Inbox {
     }
 }
 
-impl Inbox {
+impl<S: EphemeralKeyStore> Inbox<S> {
-    pub fn new(ident: Rc<Identity>) -> Self {
+    pub fn new(ident: Rc<Identity>, store: Rc<RefCell<S>>) -> Self {
         let local_convo_id = Self::inbox_identifier_for_key(ident.public_key());
         Self {
             ident,
             local_convo_id,
+            store,
         }
     }
 
     /// Creates an intro bundle and returns the Introduction along with the
     /// generated ephemeral key pair (public_key_hex, private_key) for the caller to persist.
-    pub fn create_intro_bundle(&self) -> (Introduction, String, PrivateKey) {
+    pub fn create_intro_bundle(&self) -> Result<Introduction, ChatError> {
         let ephemeral = PrivateKey::random();
 
         let ephemeral_key: PublicKey = (&ephemeral).into();
         let public_key_hex = hex::encode(ephemeral_key.as_bytes());
 
+        self.store
+            .borrow_mut()
+            .save_ephemeral_key(&public_key_hex, &ephemeral)?;
+
         let intro = Introduction::new(self.ident.secret(), ephemeral_key, OsRng);
-        (intro, public_key_hex, ephemeral)
+        Ok(intro)
     }
 
     pub fn invite_to_private_convo(
@@ -113,9 +121,15 @@ impl Inbox {
     /// looked up from storage. Returns the created conversation and optional content data.
     pub fn handle_frame(
         &self,
-        ephemeral_key: &PrivateKey,
         enc_payload: EncryptedPayload,
+        public_key_hex: &str,
     ) -> Result<(Conversation, Option<ContentData>), ChatError> {
+        let ephemeral_key = self
+            .store
+            .borrow()
+            .load_ephemeral_key(public_key_hex)?
+            .ok_or(ChatError::UnknownEphemeralKey())?;
+
         let handshake = Self::extract_payload(enc_payload)?;
 
         let header = handshake
@@ -123,11 +137,12 @@ impl Inbox {
             .ok_or(ChatError::UnexpectedPayload("InboxV1Header".into()))?;
 
         // Perform handshake and decrypt frame
-        let (seed_key, frame) = self.perform_handshake(ephemeral_key, header, handshake.payload)?;
+        let (seed_key, frame) =
+            self.perform_handshake(&ephemeral_key, header, handshake.payload)?;
 
         match frame.frame_type.unwrap() {
             proto::inbox_v1_frame::FrameType::InvitePrivateV1(_invite_private_v1) => {
-                let mut convo = PrivateV1Convo::new_responder(seed_key, ephemeral_key);
+                let mut convo = PrivateV1Convo::new_responder(seed_key, &ephemeral_key);
 
                 let Some(enc_payload) = _invite_private_v1.initial_message else {
                     return Err(ChatError::Protocol("missing initial encpayload".into()));
@@ -230,7 +245,7 @@ impl Inbox {
     }
 }
 
-impl Id for Inbox {
+impl<S: EphemeralKeyStore> Id for Inbox<S> {
     fn id(&self) -> ConversationId<'_> {
         &self.local_convo_id
     }
@@ -238,35 +253,33 @@ impl Id for Inbox {
 
 #[cfg(test)]
 mod tests {
+    use std::cell::RefCell;
+
     use super::*;
     use sqlite::{ChatStorage, StorageConfig};
-    use storage::EphemeralKeyStore;
 
     #[test]
     fn test_invite_privatev1_roundtrip() {
-        let mut storage = ChatStorage::new(StorageConfig::InMemory).unwrap();
+        let storage = Rc::new(RefCell::new(
+            ChatStorage::new(StorageConfig::InMemory).unwrap(),
+        ));
 
         let saro_ident = Identity::new("saro");
-        let saro_inbox = Inbox::new(saro_ident.into());
+        let saro_inbox = Inbox::new(saro_ident.into(), Rc::clone(&storage));
 
         let raya_ident = Identity::new("raya");
-        let raya_inbox = Inbox::new(raya_ident.into());
+        let raya_inbox = Inbox::new(raya_ident.into(), Rc::clone(&storage));
 
-        let (bundle, key_hex, private_key) = raya_inbox.create_intro_bundle();
+        let bundle = raya_inbox.create_intro_bundle().unwrap();
-        storage.save_ephemeral_key(&key_hex, &private_key).unwrap();
 
         let (_, mut payloads) = saro_inbox
             .invite_to_private_convo(&bundle, "hello".as_bytes())
             .unwrap();
 
         let payload = payloads.remove(0);
+        let key_hex = Inbox::<ChatStorage>::extract_ephemeral_key_hex(&payload.data).unwrap();
 
-        // Look up ephemeral key from storage
+        let result = raya_inbox.handle_frame(payload.data, &key_hex);
-        let key_hex = Inbox::extract_ephemeral_key_hex(&payload.data).unwrap();
-        let ephemeral_key = storage.load_ephemeral_key(&key_hex).unwrap().unwrap();
-
-        // Test handle_frame with valid payload
-        let result = raya_inbox.handle_frame(&ephemeral_key, payload.data);
 
         assert!(
             result.is_ok(),