From b0c1dbca33add864ec8c1a901820fc3f9f920563 Mon Sep 17 00:00:00 2001 From: kaichaosun Date: Thu, 5 Feb 2026 15:25:26 +0800 Subject: [PATCH] chore: ratchet session own storage --- conversations/src/chat.rs | 51 ++-- conversations/src/common.rs | 4 +- conversations/src/dm/privatev1.rs | 92 +++++--- conversations/src/errors.rs | 2 + conversations/src/inbox/inbox.rs | 15 +- double-ratchets/examples/out_of_order_demo.rs | 45 ++-- double-ratchets/examples/storage_demo.rs | 99 +++----- double-ratchets/src/storage/session.rs | 220 +++++++++--------- 8 files changed, 270 insertions(+), 258 deletions(-) diff --git a/conversations/src/chat.rs b/conversations/src/chat.rs index 295f3ba..b8727ce 100644 --- a/conversations/src/chat.rs +++ b/conversations/src/chat.rs @@ -71,8 +71,9 @@ pub struct ChatManager { inbox: Inbox, /// Storage for chat metadata (identity, inbox keys, chat records). storage: ChatStorage, - /// Storage for ratchet state (delegated to double-ratchets crate). - ratchet_storage: RatchetStorage, + /// Storage config for creating ratchet storage instances. + /// Each PrivateV1Convo gets its own storage instance (with RatchetSession). + storage_config: StorageConfig, } impl ChatManager { @@ -85,9 +86,6 @@ impl ChatManager { pub fn open(config: StorageConfig) -> Result { let mut storage = ChatStorage::new(config.clone())?; - // Initialize ratchet storage (delegated to double-ratchets crate) - let ratchet_storage = RatchetStorage::with_config(config)?; - // Load or create identity let identity = if let Some(identity) = storage.load_identity()? { identity @@ -108,7 +106,7 @@ impl ChatManager { chats: HashMap::new(), inbox, storage, - ratchet_storage, + storage_config: config, }) } @@ -117,6 +115,11 @@ impl ChatManager { Self::open(StorageConfig::InMemory) } + /// Creates a new RatchetStorage instance using the stored config. + fn create_ratchet_storage(&self) -> Result { + Ok(RatchetStorage::with_config(self.storage_config.clone())?) + } + /// Get the local identity's public address. /// /// This address can be shared with others so they can identify you. @@ -144,15 +147,20 @@ impl ChatManager { /// Start a new private conversation with someone using their introduction bundle. /// /// Returns the chat ID and envelopes that must be delivered to the remote party. - /// The chat state is automatically persisted. + /// The chat state is automatically persisted (via RatchetSession). pub fn start_private_chat( &mut self, remote_bundle: &Introduction, initial_message: &str, ) -> Result<(String, Vec), ChatManagerError> { - let (convo, payloads) = self - .inbox - .invite_to_private_convo(remote_bundle, initial_message.to_string())?; + // Create new storage for this conversation's RatchetSession + let ratchet_storage = self.create_ratchet_storage()?; + + let (convo, payloads) = self.inbox.invite_to_private_convo( + ratchet_storage, + remote_bundle, + initial_message.to_string(), + )?; let chat_id = convo.id().to_string(); @@ -169,8 +177,7 @@ impl ChatManager { ); self.storage.save_chat(&chat_record)?; - // Persist ratchet state (delegated to double-ratchets storage) - self.ratchet_storage.save(&chat_id, convo.ratchet_state())?; + // Ratchet state is automatically persisted by RatchetSession // Store in memory cache self.chats.insert(chat_id.clone(), convo); @@ -196,8 +203,7 @@ impl ChatManager { let payloads = chat.send_message(content)?; - // Persist updated ratchet state (delegated to double-ratchets storage) - self.ratchet_storage.save(chat_id, chat.ratchet_state())?; + // Ratchet state is automatically persisted by RatchetSession let remote_id = chat.remote_id(); Ok(payloads @@ -212,10 +218,10 @@ impl ChatManager { return Ok(()); } - // Try to load ratchet state from double-ratchets storage - if self.ratchet_storage.exists(chat_id)? { - let dr_state = self.ratchet_storage.load(chat_id)?; - let convo = PrivateV1Convo::from_state(chat_id.to_string(), dr_state); + // Try to load conversation from storage via RatchetSession + let ratchet_storage = self.create_ratchet_storage()?; + if ratchet_storage.exists(chat_id)? { + let convo = PrivateV1Convo::open(ratchet_storage, chat_id.to_string())?; self.chats.insert(chat_id.to_string(), convo); Ok(()) } else if self.storage.chat_exists(chat_id)? { @@ -237,8 +243,11 @@ impl ChatManager { /// Returns the decrypted content if successful. /// Any new chats or state changes are automatically persisted. pub fn handle_incoming(&mut self, payload: &[u8]) -> Result { + // Create storage for potential new conversation + let ratchet_storage = self.create_ratchet_storage()?; + // Try to handle as inbox message (new chat invitation) - match self.inbox.handle_frame(payload) { + match self.inbox.handle_frame(ratchet_storage, payload) { Ok((chat, content_data)) => { let chat_id = chat.id().to_string(); @@ -308,7 +317,9 @@ impl ChatManager { self.chats.remove(chat_id); self.storage.delete_chat(chat_id)?; // Also delete ratchet state from double-ratchets storage - let _ = self.ratchet_storage.delete(chat_id); + if let Ok(mut ratchet_storage) = self.create_ratchet_storage() { + let _ = ratchet_storage.delete(chat_id); + } Ok(()) } } diff --git a/conversations/src/common.rs b/conversations/src/common.rs index e8492a8..9714a68 100644 --- a/conversations/src/common.rs +++ b/conversations/src/common.rs @@ -2,6 +2,7 @@ use std::fmt::Debug; pub use crate::errors::ChatError; use crate::types::{AddressedEncryptedPayload, ContentData}; +use double_ratchets::storage::RatchetStorage; pub type ChatId<'a> = &'a str; @@ -12,13 +13,14 @@ pub trait HasChatId: Debug { pub trait InboundMessageHandler { fn handle_frame( &mut self, + storage: RatchetStorage, encoded_payload: &[u8], ) -> Result<(Box, Vec), ChatError>; } pub trait Chat: HasChatId + Debug { fn send_message(&mut self, content: &[u8]) - -> Result, ChatError>; + -> Result, ChatError>; fn remote_id(&self) -> String; } diff --git a/conversations/src/dm/privatev1.rs b/conversations/src/dm/privatev1.rs index 89905e6..4c22918 100644 --- a/conversations/src/dm/privatev1.rs +++ b/conversations/src/dm/privatev1.rs @@ -3,7 +3,10 @@ use chat_proto::logoschat::{ encryption::{Doubleratchet, EncryptedPayload, encrypted_payload::Encryption}, }; use crypto::SecretKey; -use double_ratchets::{Header, InstallationKeyPair, RatchetState}; +use double_ratchets::{ + Header, InstallationKeyPair, + storage::{RatchetSession, RatchetStorage}, +}; use prost::{Message, bytes::Bytes}; use std::fmt::Debug; use x25519_dalek::PublicKey; @@ -18,48 +21,61 @@ use crate::{ pub struct PrivateV1Convo { chat_id: String, - dr_state: RatchetState, + session: RatchetSession, } impl PrivateV1Convo { - pub fn new_initiator(chat_id: String, seed_key: SecretKey, remote: PublicKey) -> Self { + /// Create a new conversation as the initiator (sender of first message). + /// + /// The session will be persisted to the provided storage. + pub fn new_initiator( + storage: RatchetStorage, + chat_id: String, + seed_key: SecretKey, + remote: PublicKey, + ) -> Result { // TODO: Danger - Fix double-ratchets types to Accept SecretKey - // perhaps update the DH to work with cryptocrate. - // init_sender doesn't take ownership of the key so a reference can be used. + // perhaps update the DH to work with crypto crate. let shared_secret: [u8; 32] = seed_key.as_bytes().to_vec().try_into().unwrap(); - Self { - chat_id, - dr_state: RatchetState::init_sender(shared_secret, remote), - } + let session = RatchetSession::create_sender_session(storage, &chat_id, shared_secret, remote)?; + + Ok(Self { chat_id, session }) } + /// Create a new conversation as the responder (receiver of first message). + /// + /// The session will be persisted to the provided storage. pub fn new_responder( + storage: RatchetStorage, chat_id: String, seed_key: SecretKey, dh_self: InstallationKeyPair, - ) -> Self { - Self { - chat_id, - // TODO: Danger - Fix double-ratchets types to Accept SecretKey - dr_state: RatchetState::init_receiver(seed_key.as_bytes().to_owned(), dh_self), - } + ) -> Result { + // TODO: Danger - Fix double-ratchets types to Accept SecretKey + let shared_secret: [u8; 32] = seed_key.as_bytes().to_owned(); + let session = RatchetSession::create_receiver_session(storage, &chat_id, shared_secret, dh_self)?; + + Ok(Self { chat_id, session }) } - /// Restore a conversation from a loaded RatchetState. - pub fn from_state(chat_id: String, dr_state: RatchetState) -> Self { - Self { chat_id, dr_state } + /// Open an existing conversation from storage. + pub fn open(storage: RatchetStorage, chat_id: String) -> Result { + let session = RatchetSession::open(storage, &chat_id)?; + + Ok(Self { chat_id, session }) } - /// Get a reference to the ratchet state for storage. - pub fn ratchet_state(&self) -> &RatchetState { - &self.dr_state + /// Consumes the conversation and returns the underlying storage. + /// Useful when you need to reuse the storage for another conversation. + pub fn into_storage(self) -> RatchetStorage { + self.session.into_storage() } - fn encrypt(&mut self, frame: PrivateV1Frame) -> EncryptedPayload { + fn encrypt(&mut self, frame: PrivateV1Frame) -> Result { let encoded_bytes = frame.encode_to_vec(); - let (cipher_text, header) = self.dr_state.encrypt_message(&encoded_bytes); + let (cipher_text, header) = self.session.encrypt_message(&encoded_bytes)?; - EncryptedPayload { + Ok(EncryptedPayload { encryption: Some(Encryption::Doubleratchet(Doubleratchet { dh: Bytes::from(Vec::from(header.dh_pub.to_bytes())), msg_num: header.msg_num, @@ -67,7 +83,7 @@ impl PrivateV1Convo { ciphertext: Bytes::from(cipher_text), aux: "".into(), })), - } + }) } fn decrypt(&mut self, payload: EncryptedPayload) -> Result { @@ -101,7 +117,7 @@ impl PrivateV1Convo { // Decrypt into Frame let content_bytes = self - .dr_state + .session .decrypt_message(&dr_header.ciphertext, header) .map_err(|e| EncryptionError::Decryption(e.to_string()))?; Ok(PrivateV1Frame::decode(content_bytes.as_slice()).unwrap()) @@ -126,7 +142,7 @@ impl Chat for PrivateV1Convo { frame_type: Some(FrameType::Content(content.to_vec().into())), }; - let data = self.encrypt(frame); + let data = self.encrypt(frame)?; Ok(vec![AddressedEncryptedPayload { delivery_address: "delivery_address".into(), @@ -143,13 +159,14 @@ impl Chat for PrivateV1Convo { impl Debug for PrivateV1Convo { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("PrivateV1Convo") - .field("dr_state", &"******") + .field("session", &"******") .finish() } } #[cfg(test)] mod tests { + use double_ratchets::storage::RatchetStorage; use x25519_dalek::StaticSecret; use super::*; @@ -163,18 +180,27 @@ mod tests { let seed_key = saro.diffie_hellman(&pub_raya); let send_content_bytes = vec![0, 2, 4, 6, 8]; + + // Create in-memory storage for both parties + let storage_sender = RatchetStorage::in_memory().unwrap(); + let storage_receiver = RatchetStorage::in_memory().unwrap(); + let mut sr_convo = PrivateV1Convo::new_initiator( - "test_chat".to_string(), + storage_sender, + "test_chat_sender".to_string(), SecretKey::from(seed_key.to_bytes()), pub_raya, - ); + ) + .unwrap(); let installation_key_pair = InstallationKeyPair::from(raya); let mut rs_convo = PrivateV1Convo::new_responder( - "test_chat".to_string(), + storage_receiver, + "test_chat_receiver".to_string(), SecretKey::from(seed_key.to_bytes()), installation_key_pair, - ); + ) + .unwrap(); let send_frame = PrivateV1Frame { conversation_id: "_".into(), @@ -182,7 +208,7 @@ mod tests { timestamp: timestamp_millis(), frame_type: Some(FrameType::Content(Bytes::from(send_content_bytes.clone()))), }; - let payload = sr_convo.encrypt(send_frame.clone()); + let payload = sr_convo.encrypt(send_frame.clone()).unwrap(); let recv_frame = rs_convo.decrypt(payload).unwrap(); assert!( diff --git a/conversations/src/errors.rs b/conversations/src/errors.rs index 27cc64a..ec2e766 100644 --- a/conversations/src/errors.rs +++ b/conversations/src/errors.rs @@ -22,6 +22,8 @@ pub enum ChatError { NoConvo(u32), #[error("chat with id '{0}' was not found")] NoChatId(String), + #[error("session error: {0}")] + Session(#[from] double_ratchets::SessionError), } #[derive(Error, Debug)] diff --git a/conversations/src/inbox/inbox.rs b/conversations/src/inbox/inbox.rs index 588c171..89a6e6e 100644 --- a/conversations/src/inbox/inbox.rs +++ b/conversations/src/inbox/inbox.rs @@ -6,6 +6,7 @@ use std::collections::HashMap; use std::rc::Rc; use crypto::{PrekeyBundle, SecretKey}; +use double_ratchets::storage::RatchetStorage; use crate::common::{Chat, ChatId, HasChatId, InboundMessageHandler}; use crate::dm::privatev1::PrivateV1Convo; @@ -92,6 +93,7 @@ impl Inbox { pub fn invite_to_private_convo( &self, + storage: RatchetStorage, remote_bundle: &Introduction, initial_message: String, ) -> Result<(PrivateV1Convo, Vec), ChatError> { @@ -111,7 +113,7 @@ impl Inbox { // Generate unique chat ID let chat_id = generate_chat_id(); let mut convo = - PrivateV1Convo::new_initiator(chat_id, seed_key, remote_bundle.ephemeral_key); + PrivateV1Convo::new_initiator(storage, chat_id, seed_key, remote_bundle.ephemeral_key)?; let mut payloads = convo.send_message(initial_message.as_bytes())?; @@ -228,6 +230,7 @@ impl HasChatId for Inbox { impl InboundMessageHandler for Inbox { fn handle_frame( &mut self, + storage: RatchetStorage, message: &[u8], ) -> Result<(Box, Vec), ChatError> { if message.len() == 0 { @@ -253,7 +256,7 @@ impl InboundMessageHandler for Inbox { let chat_id = generate_chat_id(); let installation_keypair = double_ratchets::InstallationKeyPair::from(ephemeral_key.clone()); - let convo = PrivateV1Convo::new_responder(chat_id, seed_key, installation_keypair); + let convo = PrivateV1Convo::new_responder(storage, chat_id, seed_key, installation_keypair)?; // TODO: Update PrivateV1 Constructor with DR, initial_message Ok((Box::new(convo), vec![])) @@ -274,9 +277,13 @@ mod tests { let raya_ident = Identity::new(); let mut raya_inbox = Inbox::new(raya_ident.into()); + // Create in-memory storage for both parties + let storage_sender = RatchetStorage::in_memory().unwrap(); + let storage_receiver = RatchetStorage::in_memory().unwrap(); + let (bundle, _secret) = raya_inbox.create_bundle(); let (_, payloads) = saro_inbox - .invite_to_private_convo(&bundle.into(), "hello".into()) + .invite_to_private_convo(storage_sender, &bundle.into(), "hello".into()) .unwrap(); let payload = payloads @@ -287,7 +294,7 @@ mod tests { payload.data.encode(&mut buf).unwrap(); // Test handle_frame with valid payload - let result = raya_inbox.handle_frame(&buf); + let result = raya_inbox.handle_frame(storage_receiver, &buf); assert!( result.is_ok(), diff --git a/double-ratchets/examples/out_of_order_demo.rs b/double-ratchets/examples/out_of_order_demo.rs index e99cdfd..1e5841e 100644 --- a/double-ratchets/examples/out_of_order_demo.rs +++ b/double-ratchets/examples/out_of_order_demo.rs @@ -9,9 +9,9 @@ fn main() { println!("=== Out-of-Order Message Handling Demo ===\n"); let alice_db_file = NamedTempFile::new().unwrap(); - let alice_db_path = alice_db_file.path().to_str().unwrap(); + let alice_db_path = alice_db_file.path().to_str().unwrap().to_string(); let bob_db_file = NamedTempFile::new().unwrap(); - let bob_db_path = bob_db_file.path().to_str().unwrap(); + let bob_db_path = bob_db_file.path().to_str().unwrap().to_string(); let shared_secret = [0x42u8; 32]; let bob_keypair = InstallationKeyPair::generate(); @@ -25,13 +25,13 @@ fn main() { // Phase 1: Alice sends 5 messages, Bob receives 1, 3, 5 (skipping 2, 4) { - let mut alice_storage = RatchetStorage::new(alice_db_path, encryption_key) + let alice_storage = RatchetStorage::new(&alice_db_path, encryption_key) .expect("Failed to create Alice storage"); - let mut bob_storage = - RatchetStorage::new(bob_db_path, encryption_key).expect("Failed to create Bob storage"); + let bob_storage = RatchetStorage::new(&bob_db_path, encryption_key) + .expect("Failed to create Bob storage"); let mut alice_session: RatchetSession = RatchetSession::create_sender_session( - &mut alice_storage, + alice_storage, conv_id, shared_secret, bob_public, @@ -39,7 +39,7 @@ fn main() { .unwrap(); let mut bob_session: RatchetSession = RatchetSession::create_receiver_session( - &mut bob_storage, + bob_storage, conv_id, shared_secret, bob_keypair, @@ -71,10 +71,10 @@ fn main() { // Phase 2: Simulate app restart by reopening storage println!("\n Simulating app restart..."); { - let mut bob_storage = - RatchetStorage::new(bob_db_path, encryption_key).expect("Failed to reopen Bob storage"); + let bob_storage = RatchetStorage::new(&bob_db_path, encryption_key) + .expect("Failed to reopen Bob storage"); - let bob_session: RatchetSession = RatchetSession::open(&mut bob_storage, conv_id).unwrap(); + let bob_session: RatchetSession = RatchetSession::open(bob_storage, conv_id).unwrap(); println!( " After restart, Bob's skipped_keys: {}", bob_session.state().skipped_keys.len() @@ -85,11 +85,10 @@ fn main() { println!("\nBob receives delayed message 2..."); let (ct4, header4) = messages[3].clone(); // Save for replay test { - let mut bob_storage = - RatchetStorage::new(bob_db_path, encryption_key).expect("Failed to open Bob storage"); + let bob_storage = + RatchetStorage::new(&bob_db_path, encryption_key).expect("Failed to open Bob storage"); - let mut bob_session: RatchetSession = - RatchetSession::open(&mut bob_storage, conv_id).unwrap(); + let mut bob_session: RatchetSession = RatchetSession::open(bob_storage, conv_id).unwrap(); let (ct, header) = &messages[1]; let pt = bob_session.decrypt_message(ct, header.clone()).unwrap(); @@ -102,11 +101,10 @@ fn main() { println!("\nBob receives delayed message 4..."); { - let mut bob_storage = - RatchetStorage::new(bob_db_path, encryption_key).expect("Failed to open Bob storage"); + let bob_storage = + RatchetStorage::new(&bob_db_path, encryption_key).expect("Failed to open Bob storage"); - let mut bob_session: RatchetSession = - RatchetSession::open(&mut bob_storage, conv_id).unwrap(); + let mut bob_session: RatchetSession = RatchetSession::open(bob_storage, conv_id).unwrap(); let pt = bob_session.decrypt_message(&ct4, header4.clone()).unwrap(); println!(" Received: \"{}\"", String::from_utf8_lossy(&pt)); @@ -120,11 +118,10 @@ fn main() { println!("\n--- Replay Protection Demo ---"); println!("Trying to decrypt message 4 again (should fail)..."); { - let mut bob_storage = - RatchetStorage::new(bob_db_path, encryption_key).expect("Failed to open Bob storage"); + let bob_storage = + RatchetStorage::new(&bob_db_path, encryption_key).expect("Failed to open Bob storage"); - let mut bob_session: RatchetSession = - RatchetSession::open(&mut bob_storage, conv_id).unwrap(); + let mut bob_session: RatchetSession = RatchetSession::open(bob_storage, conv_id).unwrap(); match bob_session.decrypt_message(&ct4, header4) { Ok(_) => println!(" ERROR: Replay attack succeeded!"), @@ -133,8 +130,8 @@ fn main() { } // Cleanup - let _ = std::fs::remove_file(alice_db_path); - let _ = std::fs::remove_file(bob_db_path); + let _ = std::fs::remove_file(&alice_db_path); + let _ = std::fs::remove_file(&bob_db_path); println!("\n=== Demo Complete ==="); } diff --git a/double-ratchets/examples/storage_demo.rs b/double-ratchets/examples/storage_demo.rs index 8202995..4196d66 100644 --- a/double-ratchets/examples/storage_demo.rs +++ b/double-ratchets/examples/storage_demo.rs @@ -9,48 +9,47 @@ fn main() { println!("=== Double Ratchet Storage Demo ===\n"); let alice_db_file = NamedTempFile::new().unwrap(); - let alice_db_path = alice_db_file.path().to_str().unwrap(); + let alice_db_path = alice_db_file.path().to_str().unwrap().to_string(); let bob_db_file = NamedTempFile::new().unwrap(); - let bob_db_path = bob_db_file.path().to_str().unwrap(); + let bob_db_path = bob_db_file.path().to_str().unwrap().to_string(); let encryption_key = "super-secret-key-123!"; + let conv_id = "conv1"; // Initial conversation with encryption { - let mut alice_storage = RatchetStorage::new(alice_db_path, encryption_key) + let alice_storage = RatchetStorage::new(&alice_db_path, encryption_key) .expect("Failed to create alice encrypted storage"); - let mut bob_storage = RatchetStorage::new(bob_db_path, encryption_key) + let bob_storage = RatchetStorage::new(&bob_db_path, encryption_key) .expect("Failed to create bob encrypted storage"); println!( " Encrypted database created at: {}, {}", alice_db_path, bob_db_path ); - run_conversation(&mut alice_storage, &mut bob_storage); + run_conversation(alice_storage, bob_storage, conv_id); } // Restart with correct key println!("\n Simulating restart with encryption key..."); { - let mut alice_storage = RatchetStorage::new(alice_db_path, encryption_key) + let alice_storage = RatchetStorage::new(&alice_db_path, encryption_key) .expect("Failed to create alice encrypted storage"); - let mut bob_storage = RatchetStorage::new(bob_db_path, encryption_key) + let bob_storage = RatchetStorage::new(&bob_db_path, encryption_key) .expect("Failed to create bob encrypted storage"); - continue_after_restart(&mut alice_storage, &mut bob_storage); + continue_after_restart(alice_storage, bob_storage, conv_id); } - let _ = std::fs::remove_file(alice_db_path); - let _ = std::fs::remove_file(bob_db_path); + let _ = std::fs::remove_file(&alice_db_path); + let _ = std::fs::remove_file(&bob_db_path); } /// Simulates a conversation between Alice and Bob. /// Each party saves/loads state from storage for each operation. -fn run_conversation(alice_storage: &mut RatchetStorage, bob_storage: &mut RatchetStorage) { +fn run_conversation(alice_storage: RatchetStorage, bob_storage: RatchetStorage, conv_id: &str) { // === Setup: Simulate X3DH key exchange === let shared_secret = [0x42u8; 32]; // In reality, this comes from X3DH let bob_keypair = InstallationKeyPair::generate(); - let conv_id = "conv1"; - let mut alice_session: RatchetSession = RatchetSession::create_sender_session( alice_storage, conv_id, @@ -66,46 +65,31 @@ fn run_conversation(alice_storage: &mut RatchetStorage, bob_storage: &mut Ratche println!(" Sessions created for Alice and Bob"); // === Message 1: Alice -> Bob === - let (ct1, h1) = { - let result = alice_session - .encrypt_message(b"Hello Bob! This is message 1.") - .unwrap(); - println!(" Alice sent: \"Hello Bob! This is message 1.\""); - result - }; + let (ct1, h1) = alice_session + .encrypt_message(b"Hello Bob! This is message 1.") + .unwrap(); + println!(" Alice sent: \"Hello Bob! This is message 1.\""); - { - let pt = bob_session.decrypt_message(&ct1, h1).unwrap(); - println!(" Bob received: \"{}\"", String::from_utf8_lossy(&pt)); - } + let pt = bob_session.decrypt_message(&ct1, h1).unwrap(); + println!(" Bob received: \"{}\"", String::from_utf8_lossy(&pt)); // === Message 2: Bob -> Alice (triggers DH ratchet) === - let (ct2, h2) = { - let result = bob_session - .encrypt_message(b"Hi Alice! Got your message.") - .unwrap(); - println!(" Bob sent: \"Hi Alice! Got your message.\""); - result - }; + let (ct2, h2) = bob_session + .encrypt_message(b"Hi Alice! Got your message.") + .unwrap(); + println!(" Bob sent: \"Hi Alice! Got your message.\""); - { - let pt = alice_session.decrypt_message(&ct2, h2).unwrap(); - println!(" Alice received: \"{}\"", String::from_utf8_lossy(&pt)); - } + let pt = alice_session.decrypt_message(&ct2, h2).unwrap(); + println!(" Alice received: \"{}\"", String::from_utf8_lossy(&pt)); // === Message 3: Alice -> Bob === - let (ct3, h3) = { - let result = alice_session - .encrypt_message(b"Great! Let's keep chatting.") - .unwrap(); - println!(" Alice sent: \"Great! Let's keep chatting.\""); - result - }; + let (ct3, h3) = alice_session + .encrypt_message(b"Great! Let's keep chatting.") + .unwrap(); + println!(" Alice sent: \"Great! Let's keep chatting.\""); - { - let pt = bob_session.decrypt_message(&ct3, h3).unwrap(); - println!(" Bob received: \"{}\"", String::from_utf8_lossy(&pt)); - } + let pt = bob_session.decrypt_message(&ct3, h3).unwrap(); + println!(" Bob received: \"{}\"", String::from_utf8_lossy(&pt)); // Print final state println!( @@ -115,27 +99,20 @@ fn run_conversation(alice_storage: &mut RatchetStorage, bob_storage: &mut Ratche ); } -fn continue_after_restart(alice_storage: &mut RatchetStorage, bob_storage: &mut RatchetStorage) { +fn continue_after_restart(alice_storage: RatchetStorage, bob_storage: RatchetStorage, conv_id: &str) { // Load persisted states - let conv_id = "conv1"; - let mut alice_session: RatchetSession = RatchetSession::open(alice_storage, conv_id).unwrap(); let mut bob_session: RatchetSession = RatchetSession::open(bob_storage, conv_id).unwrap(); - println!(" Sessions restored for Alice and Bob",); + println!(" Sessions restored for Alice and Bob"); // Continue conversation - let (ct, header) = { - let result = alice_session - .encrypt_message(b"Message after restart!") - .unwrap(); - println!(" Alice sent: \"Message after restart!\""); - result - }; + let (ct, header) = alice_session + .encrypt_message(b"Message after restart!") + .unwrap(); + println!(" Alice sent: \"Message after restart!\""); - { - let pt = bob_session.decrypt_message(&ct, header).unwrap(); - println!(" Bob received: \"{}\"", String::from_utf8_lossy(&pt)); - } + let pt = bob_session.decrypt_message(&ct, header).unwrap(); + println!(" Bob received: \"{}\"", String::from_utf8_lossy(&pt)); println!( " Final state: Alice msg_send={}, Bob msg_recv={}", diff --git a/double-ratchets/src/storage/session.rs b/double-ratchets/src/storage/session.rs index e7ad71e..1afa28c 100644 --- a/double-ratchets/src/storage/session.rs +++ b/double-ratchets/src/storage/session.rs @@ -1,5 +1,6 @@ //! Session wrapper for automatic state persistence. +use storage::StorageConfig; use x25519_dalek::PublicKey; use crate::{ @@ -13,16 +14,19 @@ use super::RatchetStorage; /// A session wrapper that automatically persists ratchet state after operations. /// Provides rollback semantics - state is only saved if the operation succeeds. -pub struct RatchetSession<'a, D: HkdfInfo + Clone = DefaultDomain> { - storage: &'a mut RatchetStorage, +/// +/// This struct owns its storage, making it easy to store in other structs +/// and use across multiple operations without lifetime concerns. +pub struct RatchetSession { + storage: RatchetStorage, conversation_id: String, state: RatchetState, } -impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> { +impl RatchetSession { /// Opens an existing session from storage. pub fn open( - storage: &'a mut RatchetStorage, + storage: RatchetStorage, conversation_id: impl Into, ) -> Result { let conversation_id = conversation_id.into(); @@ -34,9 +38,18 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> { }) } + /// Opens an existing session with the given storage configuration. + pub fn open_with_config( + config: StorageConfig, + conversation_id: impl Into, + ) -> Result { + let storage = RatchetStorage::with_config(config)?; + Self::open(storage, conversation_id) + } + /// Creates a new session and persists the initial state. pub fn create( - storage: &'a mut RatchetStorage, + mut storage: RatchetStorage, conversation_id: impl Into, state: RatchetState, ) -> Result { @@ -49,9 +62,19 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> { }) } + /// Creates a new session with the given storage configuration. + pub fn create_with_config( + config: StorageConfig, + conversation_id: impl Into, + state: RatchetState, + ) -> Result { + let storage = RatchetStorage::with_config(config)?; + Self::create(storage, conversation_id, state) + } + /// Initializes a new session as a sender and persists the initial state. pub fn create_sender_session( - storage: &'a mut RatchetStorage, + storage: RatchetStorage, conversation_id: &str, shared_secret: SharedSecret, remote_pub: PublicKey, @@ -60,12 +83,12 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> { return Err(SessionError::ConvAlreadyExists(conversation_id.to_string())); } let state = RatchetState::::init_sender(shared_secret, remote_pub); - Ok(Self::create(storage, conversation_id, state)?) + Self::create(storage, conversation_id, state) } /// Initializes a new session as a receiver and persists the initial state. pub fn create_receiver_session( - storage: &'a mut RatchetStorage, + storage: RatchetStorage, conversation_id: &str, shared_secret: SharedSecret, dh_self: InstallationKeyPair, @@ -75,7 +98,7 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> { } let state = RatchetState::::init_receiver(shared_secret, dh_self); - Ok(Self::create(storage, conversation_id, state)?) + Self::create(storage, conversation_id, state) } /// Encrypts a message and persists the updated state. @@ -137,6 +160,12 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> { &self.conversation_id } + /// Consumes the session and returns the underlying storage. + /// Useful when you need to reuse the storage for another session. + pub fn into_storage(self) -> RatchetStorage { + self.storage + } + /// Manually saves the current state. pub fn save(&mut self) -> Result<(), SessionError> { self.storage @@ -164,30 +193,29 @@ mod tests { #[test] fn test_session_create_and_open() { - let mut storage = create_test_storage(); + let storage = create_test_storage(); let shared_secret = [0x42; 32]; let bob_keypair = InstallationKeyPair::generate(); let alice: RatchetState = RatchetState::init_sender(shared_secret, bob_keypair.public().clone()); - // Create session - { - let session = RatchetSession::create(&mut storage, "conv1", alice).unwrap(); - assert_eq!(session.conversation_id(), "conv1"); - } + // Create session - session takes ownership of storage + let session = RatchetSession::create(storage, "conv1", alice).unwrap(); + assert_eq!(session.conversation_id(), "conv1"); + + // Get storage back from session to reopen + let storage = session.into_storage(); // Open existing session - { - let session: RatchetSession = - RatchetSession::open(&mut storage, "conv1").unwrap(); - assert_eq!(session.state().msg_send, 0); - } + let session: RatchetSession = + RatchetSession::open(storage, "conv1").unwrap(); + assert_eq!(session.state().msg_send, 0); } #[test] fn test_session_encrypt_persists() { - let mut storage = create_test_storage(); + let storage = create_test_storage(); let shared_secret = [0x42; 32]; let bob_keypair = InstallationKeyPair::generate(); @@ -195,158 +223,120 @@ mod tests { RatchetState::init_sender(shared_secret, bob_keypair.public().clone()); // Create and encrypt - { - let mut session = RatchetSession::create(&mut storage, "conv1", alice).unwrap(); - session.encrypt_message(b"Hello").unwrap(); - assert_eq!(session.state().msg_send, 1); - } + let mut session = RatchetSession::create(storage, "conv1", alice).unwrap(); + session.encrypt_message(b"Hello").unwrap(); + assert_eq!(session.state().msg_send, 1); + + // Get storage back and reopen + let storage = session.into_storage(); // Reopen - state should be persisted - { - let session: RatchetSession = - RatchetSession::open(&mut storage, "conv1").unwrap(); - assert_eq!(session.state().msg_send, 1); - } + let session: RatchetSession = + RatchetSession::open(storage, "conv1").unwrap(); + assert_eq!(session.state().msg_send, 1); } #[test] fn test_session_full_conversation() { - let mut storage = create_test_storage(); + // Use separate in-memory storages for alice and bob (simulates different devices) + let alice_storage = create_test_storage(); + let bob_storage = create_test_storage(); let shared_secret = [0x42; 32]; let bob_keypair = InstallationKeyPair::generate(); - let alice: RatchetState = + let alice_state: RatchetState = RatchetState::init_sender(shared_secret, bob_keypair.public().clone()); - let bob: RatchetState = + let bob_state: RatchetState = RatchetState::init_receiver(shared_secret, bob_keypair); // Alice sends - let (ct, header) = { - let mut session = RatchetSession::create(&mut storage, "alice", alice).unwrap(); - session.encrypt_message(b"Hello Bob").unwrap() - }; + let mut alice_session = RatchetSession::create(alice_storage, "conv", alice_state).unwrap(); + let (ct, header) = alice_session.encrypt_message(b"Hello Bob").unwrap(); // Bob receives - let plaintext = { - let mut session = RatchetSession::create(&mut storage, "bob", bob).unwrap(); - session.decrypt_message(&ct, header).unwrap() - }; + let mut bob_session = RatchetSession::create(bob_storage, "conv", bob_state).unwrap(); + let plaintext = bob_session.decrypt_message(&ct, header).unwrap(); assert_eq!(plaintext, b"Hello Bob"); // Bob replies - let (ct2, header2) = { - let mut session: RatchetSession = - RatchetSession::open(&mut storage, "bob").unwrap(); - session.encrypt_message(b"Hi Alice").unwrap() - }; + let (ct2, header2) = bob_session.encrypt_message(b"Hi Alice").unwrap(); // Alice receives - let plaintext2 = { - let mut session: RatchetSession = - RatchetSession::open(&mut storage, "alice").unwrap(); - session.decrypt_message(&ct2, header2).unwrap() - }; + let plaintext2 = alice_session.decrypt_message(&ct2, header2).unwrap(); assert_eq!(plaintext2, b"Hi Alice"); } #[test] fn test_session_open_or_create() { - let mut storage = create_test_storage(); + let storage = create_test_storage(); let shared_secret = [0x42; 32]; let bob_keypair = InstallationKeyPair::generate(); let bob_pub = bob_keypair.public().clone(); // First call creates - { - let session: RatchetSession = RatchetSession::create_sender_session( - &mut storage, - "conv1", - shared_secret, - bob_pub.clone(), - ) - .unwrap(); - assert_eq!(session.state().msg_send, 0); - } + let session: RatchetSession = + RatchetSession::create_sender_session(storage, "conv1", shared_secret, bob_pub.clone()) + .unwrap(); + assert_eq!(session.state().msg_send, 0); + let storage = session.into_storage(); - // Second call opens existing - { - let mut session: RatchetSession = - RatchetSession::open(&mut storage, "conv1").unwrap(); - session.encrypt_message(b"test").unwrap(); - } + // Second call opens existing and encrypts + let mut session: RatchetSession = + RatchetSession::open(storage, "conv1").unwrap(); + session.encrypt_message(b"test").unwrap(); + let storage = session.into_storage(); // Verify persistence - { - let session: RatchetSession = - RatchetSession::open(&mut storage, "conv1").unwrap(); - assert_eq!(session.state().msg_send, 1); - } + let session: RatchetSession = + RatchetSession::open(storage, "conv1").unwrap(); + assert_eq!(session.state().msg_send, 1); } #[test] fn test_create_sender_session_fails_when_conversation_exists() { - let mut storage = create_test_storage(); + let storage = create_test_storage(); let shared_secret = [0x42; 32]; let bob_keypair = InstallationKeyPair::generate(); let bob_pub = bob_keypair.public().clone(); // First creation succeeds - { - let _session: RatchetSession = RatchetSession::create_sender_session( - &mut storage, - "conv1", - shared_secret, - bob_pub.clone(), - ) - .unwrap(); - } + let session: RatchetSession = + RatchetSession::create_sender_session(storage, "conv1", shared_secret, bob_pub.clone()) + .unwrap(); + let storage = session.into_storage(); // Second creation should fail with ConversationAlreadyExists - { - let result: Result, _> = - RatchetSession::create_sender_session( - &mut storage, - "conv1", - shared_secret, - bob_pub.clone(), - ); + let result: Result, _> = + RatchetSession::create_sender_session(storage, "conv1", shared_secret, bob_pub.clone()); - assert!(matches!(result, Err(SessionError::ConvAlreadyExists(_)))); - } + assert!(matches!(result, Err(SessionError::ConvAlreadyExists(_)))); } #[test] fn test_create_receiver_session_fails_when_conversation_exists() { - let mut storage = create_test_storage(); + let storage = create_test_storage(); let shared_secret = [0x42; 32]; let bob_keypair = InstallationKeyPair::generate(); // First creation succeeds - { - let _session: RatchetSession = RatchetSession::create_receiver_session( - &mut storage, - "conv1", - shared_secret, - bob_keypair, - ) - .unwrap(); - } + let session: RatchetSession = + RatchetSession::create_receiver_session(storage, "conv1", shared_secret, bob_keypair) + .unwrap(); + let storage = session.into_storage(); // Second creation should fail with ConversationAlreadyExists - { - let another_keypair = InstallationKeyPair::generate(); - let result: Result, _> = - RatchetSession::create_receiver_session( - &mut storage, - "conv1", - shared_secret, - another_keypair, - ); + let another_keypair = InstallationKeyPair::generate(); + let result: Result, _> = + RatchetSession::create_receiver_session( + storage, + "conv1", + shared_secret, + another_keypair, + ); - assert!(matches!(result, Err(SessionError::ConvAlreadyExists(_)))); - } + assert!(matches!(result, Err(SessionError::ConvAlreadyExists(_)))); } }