feat: move private store out of context

This commit is contained in:
kaichaosun 2026-04-08 15:49:09 +08:00
parent d68c0cb275
commit 9c6d79d575
No known key found for this signature in database
GPG Key ID: 223E0F992F4F03BF
5 changed files with 73 additions and 32 deletions

2
.gitignore vendored
View File

@ -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

View File

@ -85,7 +85,7 @@ impl<T: ChatStore> Context<T> {
) -> Result<(ConversationIdOwned, Vec<AddressedEnvelope>), 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::<T>::inbox_identifier_for_key(*remote_bundle.installation_key());
@ -147,7 +147,9 @@ impl<T: ChatStore> Context<T> {
enc_payload: EncryptedPayload,
) -> Result<Option<ContentData>, ChatError> {
let public_key_hex = Inbox::<T>::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<T: ChatStore> Context<T> {
}
/// Loads a conversation from DB by constructing it from metadata + ratchet state.
fn load_convo(&self, convo_id: ConversationId) -> Result<Conversation, ChatError> {
fn load_convo(&self, convo_id: ConversationId) -> Result<Conversation<T>, ChatError> {
let record = self
.store
.borrow()
@ -205,6 +207,7 @@ impl<T: ChatStore> Context<T> {
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<T: ChatStore> Context<T> {
}
/// Persists a conversation's metadata and ratchet state to DB.
fn persist_convo(&mut self, convo: &PrivateV1Convo) -> Result<ConversationIdOwned, ChatError> {
fn persist_convo(
&mut self,
convo: &PrivateV1Convo<T>,
) -> Result<ConversationIdOwned, ChatError> {
let convo_info = ConversationMeta {
local_convo_id: convo.id().to_string(),
remote_convo_id: convo.remote_id(),

View File

@ -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<S: ConversationStore + RatchetStore> {
Private(PrivateV1Convo<S>),
}

View File

@ -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<S: ConversationStore + RatchetStore> {
local_convo_id: String,
remote_convo_id: String,
dr_state: RatchetState,
store: Rc<RefCell<S>>,
}
impl PrivateV1Convo {
impl<S: ConversationStore + RatchetStore> PrivateV1Convo<S> {
/// 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<RefCell<S>>,
) -> 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<RefCell<S>>,
) -> 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<RefCell<S>>,
) -> 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<S: ConversationStore + RatchetStore> Id for PrivateV1Convo<S> {
fn id(&self) -> ConversationId<'_> {
&self.local_convo_id
}
}
impl Convo for PrivateV1Convo {
impl<S: ConversationStore + RatchetStore> Convo for PrivateV1Convo<S> {
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<S: ConversationStore + RatchetStore> Debug for PrivateV1Convo<S> {
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(),

View File

@ -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<S: EphemeralKeyStore> Inbox<S> {
Ok(intro)
}
pub fn invite_to_private_convo(
pub fn invite_to_private_convo<PS: ConversationStore + RatchetStore>(
&self,
remote_bundle: &Introduction,
initial_message: &[u8],
) -> Result<(PrivateV1Convo, Vec<AddressedEncryptedPayload>), ChatError> {
private_store: Rc<RefCell<PS>>,
) -> Result<(PrivateV1Convo<PS>, Vec<AddressedEncryptedPayload>), ChatError> {
let mut rng = OsRng;
let pkb = PrekeyBundle {
@ -81,7 +82,8 @@ impl<S: EphemeralKeyStore> Inbox<S> {
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<S: EphemeralKeyStore> Inbox<S> {
/// 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<PS: ConversationStore + RatchetStore>(
&self,
enc_payload: EncryptedPayload,
public_key_hex: &str,
) -> Result<(Conversation, Option<ContentData>), ChatError> {
private_store: Rc<RefCell<PS>>,
) -> Result<(Conversation<PS>, Option<ContentData>), ChatError> {
let ephemeral_key = self
.store
.borrow()
@ -142,7 +145,8 @@ impl<S: EphemeralKeyStore> Inbox<S> {
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::<ChatStorage>::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(),