From 9c6d79d57500cd45cc2273ac05c902f5d8c8496b Mon Sep 17 00:00:00 2001 From: kaichaosun Date: Wed, 8 Apr 2026 15:49:09 +0800 Subject: [PATCH] feat: move private store out of context --- .gitignore | 2 +- core/conversations/src/context.rs | 14 +++-- core/conversations/src/conversation.rs | 6 +-- .../src/conversation/privatev1.rs | 52 ++++++++++++++----- core/conversations/src/inbox/handler.rs | 31 ++++++----- 5 files changed, 73 insertions(+), 32 deletions(-) diff --git a/.gitignore b/.gitignore index 9d00a2f..c968a3d 100644 --- a/.gitignore +++ b/.gitignore @@ -29,8 +29,8 @@ target # Temporary data folder tmp - .DS_Store +justfile # Generated C headers (produced by `make` in examples/c-ffi; do not commit) crates/client-ffi/client_ffi.h diff --git a/core/conversations/src/context.rs b/core/conversations/src/context.rs index 3415c88..2c14aff 100644 --- a/core/conversations/src/context.rs +++ b/core/conversations/src/context.rs @@ -85,7 +85,7 @@ impl Context { ) -> Result<(ConversationIdOwned, Vec), ChatError> { let (convo, payloads) = self .inbox - .invite_to_private_convo(remote_bundle, content) + .invite_to_private_convo(remote_bundle, content, Rc::clone(&self.store)) .unwrap_or_else(|_| todo!("Log/Surface Error")); let remote_id = Inbox::::inbox_identifier_for_key(*remote_bundle.installation_key()); @@ -147,7 +147,9 @@ impl Context { enc_payload: EncryptedPayload, ) -> Result, ChatError> { let public_key_hex = Inbox::::extract_ephemeral_key_hex(&enc_payload)?; - let (convo, content) = self.inbox.handle_frame(enc_payload, &public_key_hex)?; + let (convo, content) = + self.inbox + .handle_frame(enc_payload, &public_key_hex, Rc::clone(&self.store))?; match convo { Conversation::Private(convo) => self.persist_convo(&convo)?, @@ -182,7 +184,7 @@ impl Context { } /// Loads a conversation from DB by constructing it from metadata + ratchet state. - fn load_convo(&self, convo_id: ConversationId) -> Result { + fn load_convo(&self, convo_id: ConversationId) -> Result, ChatError> { let record = self .store .borrow() @@ -205,6 +207,7 @@ impl Context { record.local_convo_id, record.remote_convo_id, dr_state, + Rc::clone(&self.store), ))) } ConversationKind::Unknown(_) => Err(ChatError::BadBundleValue(format!( @@ -215,7 +218,10 @@ impl Context { } /// Persists a conversation's metadata and ratchet state to DB. - fn persist_convo(&mut self, convo: &PrivateV1Convo) -> Result { + fn persist_convo( + &mut self, + convo: &PrivateV1Convo, + ) -> Result { let convo_info = ConversationMeta { local_convo_id: convo.id().to_string(), remote_convo_id: convo.remote_id(), diff --git a/core/conversations/src/conversation.rs b/core/conversations/src/conversation.rs index 2c058dd..1580d78 100644 --- a/core/conversations/src/conversation.rs +++ b/core/conversations/src/conversation.rs @@ -4,7 +4,7 @@ use crate::types::{AddressedEncryptedPayload, ContentData}; use chat_proto::logoschat::encryption::EncryptedPayload; use std::fmt::Debug; use std::sync::Arc; -use storage::ConversationKind; +use storage::{ConversationKind, ConversationStore, RatchetStore}; pub use crate::errors::ChatError; pub use privatev1::PrivateV1Convo; @@ -36,6 +36,6 @@ pub trait Convo: Id + Debug { fn convo_type(&self) -> ConversationKind; } -pub enum Conversation { - Private(PrivateV1Convo), +pub enum Conversation { + Private(PrivateV1Convo), } diff --git a/core/conversations/src/conversation/privatev1.rs b/core/conversations/src/conversation/privatev1.rs index f2f8a22..85427f0 100644 --- a/core/conversations/src/conversation/privatev1.rs +++ b/core/conversations/src/conversation/privatev1.rs @@ -9,8 +9,8 @@ use chat_proto::logoschat::{ use crypto::{PrivateKey, PublicKey, SymmetricKey32}; use double_ratchets::{Header, InstallationKeyPair, RatchetState}; use prost::{Message, bytes::Bytes}; -use std::fmt::Debug; -use storage::ConversationKind; +use std::{cell::RefCell, fmt::Debug, rc::Rc}; +use storage::{ConversationKind, ConversationStore}; use crate::{ conversation::{ChatError, ConversationId, Convo, Id}, @@ -55,23 +55,34 @@ impl BaseConvoId { } } -pub struct PrivateV1Convo { +pub struct PrivateV1Convo { local_convo_id: String, remote_convo_id: String, dr_state: RatchetState, + store: Rc>, } -impl PrivateV1Convo { +impl PrivateV1Convo { /// Reconstructs a PrivateV1Convo from persisted metadata and ratchet state. - pub fn new(local_convo_id: String, remote_convo_id: String, dr_state: RatchetState) -> Self { + pub fn new( + local_convo_id: String, + remote_convo_id: String, + dr_state: RatchetState, + store: Rc>, + ) -> Self { Self { local_convo_id, remote_convo_id, dr_state, + store, } } - pub fn new_initiator(seed_key: SymmetricKey32, remote: PublicKey) -> Self { + pub fn new_initiator( + seed_key: SymmetricKey32, + remote: PublicKey, + store: Rc>, + ) -> Self { let base_convo_id = BaseConvoId::new(&seed_key); let local_convo_id = base_convo_id.id_for_participant(Role::Initiator); let remote_convo_id = base_convo_id.id_for_participant(Role::Responder); @@ -86,10 +97,15 @@ impl PrivateV1Convo { local_convo_id, remote_convo_id, dr_state, + store, } } - pub fn new_responder(seed_key: SymmetricKey32, dh_self: &PrivateKey) -> Self { + pub fn new_responder( + seed_key: SymmetricKey32, + dh_self: &PrivateKey, + store: Rc>, + ) -> Self { let base_convo_id = BaseConvoId::new(&seed_key); let local_convo_id = base_convo_id.id_for_participant(Role::Responder); let remote_convo_id = base_convo_id.id_for_participant(Role::Initiator); @@ -105,6 +121,7 @@ impl PrivateV1Convo { local_convo_id, remote_convo_id, dr_state, + store, } } @@ -177,13 +194,13 @@ impl PrivateV1Convo { } } -impl Id for PrivateV1Convo { +impl Id for PrivateV1Convo { fn id(&self) -> ConversationId<'_> { &self.local_convo_id } } -impl Convo for PrivateV1Convo { +impl Convo for PrivateV1Convo { fn send_message( &mut self, content: &[u8], @@ -216,6 +233,8 @@ impl Convo for PrivateV1Convo { return Err(ChatError::ProtocolExpectation("None", "Some".into())); }; + self.save_ratchet_state(&mut *self.store.borrow_mut())?; + // Handle FrameTypes let output = match frame_type { FrameType::Content(bytes) => self.handle_content(bytes.into()), @@ -234,7 +253,7 @@ impl Convo for PrivateV1Convo { } } -impl Debug for PrivateV1Convo { +impl Debug for PrivateV1Convo { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("PrivateV1Convo") .field("dr_state", &"******") @@ -245,6 +264,7 @@ impl Debug for PrivateV1Convo { #[cfg(test)] mod tests { use crypto::PrivateKey; + use sqlite::{ChatStorage, StorageConfig}; use super::*; @@ -253,14 +273,22 @@ mod tests { let saro = PrivateKey::random(); let raya = PrivateKey::random(); + let saro_storage = Rc::new(RefCell::new( + ChatStorage::new(StorageConfig::InMemory).unwrap(), + )); + + let raya_storage = Rc::new(RefCell::new( + ChatStorage::new(StorageConfig::InMemory).unwrap(), + )); + let pub_raya = PublicKey::from(&raya); let seed_key = saro.diffie_hellman(&pub_raya).DANGER_to_bytes(); let seed_key_saro = SymmetricKey32::from(seed_key); let seed_key_raya = SymmetricKey32::from(seed_key); let send_content_bytes = vec![0, 2, 4, 6, 8]; - let mut sr_convo = PrivateV1Convo::new_initiator(seed_key_saro, pub_raya); - let mut rs_convo = PrivateV1Convo::new_responder(seed_key_raya, &raya); + let mut sr_convo = PrivateV1Convo::new_initiator(seed_key_saro, pub_raya, saro_storage); + let mut rs_convo = PrivateV1Convo::new_responder(seed_key_raya, &raya, raya_storage); let send_frame = PrivateV1Frame { conversation_id: "_".into(), diff --git a/core/conversations/src/inbox/handler.rs b/core/conversations/src/inbox/handler.rs index 395f9e3..8e8df24 100644 --- a/core/conversations/src/inbox/handler.rs +++ b/core/conversations/src/inbox/handler.rs @@ -5,7 +5,7 @@ use prost::bytes::Bytes; use rand_core::OsRng; use std::cell::RefCell; use std::rc::Rc; -use storage::EphemeralKeyStore; +use storage::{ConversationStore, EphemeralKeyStore, RatchetStore}; use crypto::{PrekeyBundle, SymmetricKey32}; @@ -64,11 +64,12 @@ impl Inbox { Ok(intro) } - pub fn invite_to_private_convo( + pub fn invite_to_private_convo( &self, remote_bundle: &Introduction, initial_message: &[u8], - ) -> Result<(PrivateV1Convo, Vec), ChatError> { + private_store: Rc>, + ) -> Result<(PrivateV1Convo, Vec), ChatError> { let mut rng = OsRng; let pkb = PrekeyBundle { @@ -81,7 +82,8 @@ impl Inbox { let (seed_key, ephemeral_pub) = InboxHandshake::perform_as_initiator(self.ident.secret(), &pkb, &mut rng); - let mut convo = PrivateV1Convo::new_initiator(seed_key, *remote_bundle.ephemeral_key()); + let mut convo = + PrivateV1Convo::new_initiator(seed_key, *remote_bundle.ephemeral_key(), private_store); let mut payloads = convo.send_message(initial_message)?; @@ -119,11 +121,12 @@ impl Inbox { /// Handles an incoming inbox frame. The caller must provide the ephemeral private key /// looked up from storage. Returns the created conversation and optional content data. - pub fn handle_frame( + pub fn handle_frame( &self, enc_payload: EncryptedPayload, public_key_hex: &str, - ) -> Result<(Conversation, Option), ChatError> { + private_store: Rc>, + ) -> Result<(Conversation, Option), ChatError> { let ephemeral_key = self .store .borrow() @@ -142,7 +145,8 @@ impl Inbox { 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, private_store); let Some(enc_payload) = _invite_private_v1.initial_message else { return Err(ChatError::Protocol("missing initial encpayload".into())); @@ -260,26 +264,29 @@ mod tests { #[test] fn test_invite_privatev1_roundtrip() { - let storage = Rc::new(RefCell::new( + let saro_storage = Rc::new(RefCell::new( + ChatStorage::new(StorageConfig::InMemory).unwrap(), + )); + let raya_storage = Rc::new(RefCell::new( ChatStorage::new(StorageConfig::InMemory).unwrap(), )); let saro_ident = Identity::new("saro"); - let saro_inbox = Inbox::new(saro_ident.into(), Rc::clone(&storage)); + let saro_inbox = Inbox::new(saro_ident.into(), Rc::clone(&saro_storage)); let raya_ident = Identity::new("raya"); - let raya_inbox = Inbox::new(raya_ident.into(), Rc::clone(&storage)); + let raya_inbox = Inbox::new(raya_ident.into(), Rc::clone(&raya_storage)); let bundle = raya_inbox.create_intro_bundle().unwrap(); let (_, mut payloads) = saro_inbox - .invite_to_private_convo(&bundle, "hello".as_bytes()) + .invite_to_private_convo(&bundle, "hello".as_bytes(), saro_storage) .unwrap(); let payload = payloads.remove(0); let key_hex = Inbox::::extract_ephemeral_key_hex(&payload.data).unwrap(); - let result = raya_inbox.handle_frame(payload.data, &key_hex); + let result = raya_inbox.handle_frame(payload.data, &key_hex, raya_storage); assert!( result.is_ok(),