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
This commit is contained in:
kaichao 2026-04-08 11:07:23 +08:00 committed by GitHub
parent c44c52b127
commit 8e8b0c0dd7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 70 additions and 53 deletions

View File

@ -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,
store: T,
inbox: Inbox<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 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::<T>::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<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();
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<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 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<T: ChatStore> Context<T> {
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");
}

View File

@ -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 {
pub fn new(ident: Rc<Identity>) -> Self {
impl<S: EphemeralKeyStore> Inbox<S> {
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();
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::<ChatStorage>::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(),