From 8e8b0c0dd7fc6f5e32c3977a14b0a34d7d6b507f Mon Sep 17 00:00:00 2001 From: kaichao Date: Wed, 8 Apr 2026 11:07:23 +0800 Subject: [PATCH] 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 | 66 +++++++++++++------------ core/conversations/src/inbox/handler.rs | 57 ++++++++++++--------- 2 files changed, 70 insertions(+), 53 deletions(-) diff --git a/core/conversations/src/context.rs b/core/conversations/src/context.rs index 36b9187..3415c88 100644 --- a/core/conversations/src/context.rs +++ b/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 { _identity: Rc, - inbox: Inbox, - store: T, + inbox: Inbox, + store: Rc>, } impl Context { @@ -29,20 +29,21 @@ impl Context { /// /// 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, mut store: T) -> Result { + pub fn new_from_store(name: impl Into, store: T) -> Result { 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 Context { /// 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, mut chat_store: T) -> Self { + pub fn new_with_name(name: impl Into, 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 Context { .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::::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 Context { } pub fn list_conversations(&self) -> Result, 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 Context { 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::(&mut *self.store.borrow_mut())?; Ok(payloads .into_iter() @@ -133,7 +136,7 @@ impl Context { 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 Context { &mut self, enc_payload: EncryptedPayload, ) -> Result, ChatError> { - // Look up the ephemeral key from storage - let key_hex = Inbox::extract_ephemeral_key_hex(&enc_payload)?; - 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)?; + let public_key_hex = Inbox::::extract_ephemeral_key_hex(&enc_payload)?; + let (convo, content) = self.inbox.handle_frame(enc_payload, &public_key_hex)?; 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 Context { 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, ChatError> { - let (intro, public_key_hex, private_key) = self.inbox.create_intro_bundle(); - self.store - .save_ephemeral_key(&public_key_hex, &private_key)?; + let intro = self.inbox.create_intro_bundle()?; Ok(intro.into()) } @@ -188,13 +185,20 @@ impl Context { fn load_convo(&self, convo_id: ConversationId) -> Result { 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 skipped_keys = self.store.load_skipped_keys(&record.local_convo_id)?; + let dr_record = self + .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 Context { remote_convo_id: convo.remote_id(), kind: convo.convo_type(), }; - self.store.save_conversation(&convo_info)?; - convo.save_ratchet_state(&mut self.store)?; + self.store.borrow_mut().save_conversation(&convo_info)?; + 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"); } diff --git a/core/conversations/src/inbox/handler.rs b/core/conversations/src/inbox/handler.rs index b9e1126..395f9e3 100644 --- a/core/conversations/src/inbox/handler.rs +++ b/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 { ident: Rc, local_convo_id: String, + store: Rc>, } -impl std::fmt::Debug for Inbox { +impl std::fmt::Debug for Inbox { 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 { - pub fn new(ident: Rc) -> Self { +impl Inbox { + pub fn new(ident: Rc, store: Rc>) -> 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 { 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), 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 Id for Inbox { 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(); - storage.save_ephemeral_key(&key_hex, &private_key).unwrap(); + let bundle = raya_inbox.create_intro_bundle().unwrap(); let (_, mut payloads) = saro_inbox .invite_to_private_convo(&bundle, "hello".as_bytes()) .unwrap(); let payload = payloads.remove(0); + let key_hex = Inbox::::extract_ephemeral_key_hex(&payload.data).unwrap(); - // Look up ephemeral key from storage - 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); + let result = raya_inbox.handle_frame(payload.data, &key_hex); assert!( result.is_ok(),