logos-messaging/libchat
1
0
Fork 0
mirror of https://github.com/logos-messaging/libchat.git synced 2026-02-10 00:43:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: ratchet session own storage

Browse Source
This commit is contained in:
kaichaosun 2026-02-05 15:25:26 +08:00
parent 55285aa24e
commit b0c1dbca33
No known key found for this signature in database
GPG Key ID: 223E0F992F4F03BF
8 changed files with 270 additions and 258 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
conversations/src/chat.rs
View File

@ -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<Self, ChatManagerError> {
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<RatchetStorage, ChatManagerError> {
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<AddressedEnvelope>), 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<ContentData, ChatManagerError> {
// 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(())
}
}

4
conversations/src/common.rs
View File

@ -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<dyn Chat>, Vec<ContentData>), ChatError>;
}
pub trait Chat: HasChatId + Debug {
fn send_message(&mut self, content: &[u8])
-> Result<Vec<AddressedEncryptedPayload>, ChatError>;
-> Result<Vec<AddressedEncryptedPayload>, ChatError>;
fn remote_id(&self) -> String;
}

92
conversations/src/dm/privatev1.rs
View File

@ -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<Self, ChatError> {
// 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<Self, ChatError> {
// 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<Self, ChatError> {
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<EncryptedPayload, ChatError> {
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<PrivateV1Frame, EncryptionError> {
@ -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!(

2
conversations/src/errors.rs
View File

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

15
conversations/src/inbox/inbox.rs
View File

@ -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<AddressedEncryptedPayload>), 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<dyn Chat>, Vec<ContentData>), 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(),

45
double-ratchets/examples/out_of_order_demo.rs
View File

@ -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 ===");
}

99
double-ratchets/examples/storage_demo.rs
View File

@ -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={}",

220
double-ratchets/src/storage/session.rs
View File

@ -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<D: HkdfInfo + Clone = DefaultDomain> {
storage: RatchetStorage,
conversation_id: String,
state: RatchetState<D>,
}
impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> {
impl<D: HkdfInfo + Clone> RatchetSession<D> {
/// Opens an existing session from storage.
pub fn open(
storage: &'a mut RatchetStorage,
storage: RatchetStorage,
conversation_id: impl Into<String>,
) -> Result<Self, SessionError> {
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<String>,
) -> Result<Self, SessionError> {
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<String>,
state: RatchetState<D>,
) -> Result<Self, SessionError> {
@ -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<String>,
state: RatchetState<D>,
) -> Result<Self, SessionError> {
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::<D>::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::<D>::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<DefaultDomain> =
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<DefaultDomain> =
RatchetSession::open(&mut storage, "conv1").unwrap();
assert_eq!(session.state().msg_send, 0);
}
let session: RatchetSession<DefaultDomain> =
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<DefaultDomain> =
RatchetSession::open(&mut storage, "conv1").unwrap();
assert_eq!(session.state().msg_send, 1);
}
let session: RatchetSession<DefaultDomain> =
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<DefaultDomain> =
let alice_state: RatchetState<DefaultDomain> =
RatchetState::init_sender(shared_secret, bob_keypair.public().clone());
let bob: RatchetState<DefaultDomain> =
let bob_state: RatchetState<DefaultDomain> =
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<DefaultDomain> =
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<DefaultDomain> =
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<DefaultDomain> = RatchetSession::create_sender_session(
&mut storage,
"conv1",
shared_secret,
bob_pub.clone(),
)
.unwrap();
assert_eq!(session.state().msg_send, 0);
}
let session: RatchetSession<DefaultDomain> =
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<DefaultDomain> =
RatchetSession::open(&mut storage, "conv1").unwrap();
session.encrypt_message(b"test").unwrap();
}
// Second call opens existing and encrypts
let mut session: RatchetSession<DefaultDomain> =
RatchetSession::open(storage, "conv1").unwrap();
session.encrypt_message(b"test").unwrap();
let storage = session.into_storage();
// Verify persistence
{
let session: RatchetSession<DefaultDomain> =
RatchetSession::open(&mut storage, "conv1").unwrap();
assert_eq!(session.state().msg_send, 1);
}
let session: RatchetSession<DefaultDomain> =
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<DefaultDomain> = RatchetSession::create_sender_session(
&mut storage,
"conv1",
shared_secret,
bob_pub.clone(),
)
.unwrap();
}
let session: RatchetSession<DefaultDomain> =
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<DefaultDomain>, _> =
RatchetSession::create_sender_session(
&mut storage,
"conv1",
shared_secret,
bob_pub.clone(),
);
let result: Result<RatchetSession<DefaultDomain>, _> =
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<DefaultDomain> = RatchetSession::create_receiver_session(
&mut storage,
"conv1",
shared_secret,
bob_keypair,
)
.unwrap();
}
let session: RatchetSession<DefaultDomain> =
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<DefaultDomain>, _> =
RatchetSession::create_receiver_session(
&mut storage,
"conv1",
shared_secret,
another_keypair,
);
let another_keypair = InstallationKeyPair::generate();
let result: Result<RatchetSession<DefaultDomain>, _> =
RatchetSession::create_receiver_session(
storage,
"conv1",
shared_secret,
another_keypair,
);
assert!(matches!(result, Err(SessionError::ConvAlreadyExists(_))));
}
assert!(matches!(result, Err(SessionError::ConvAlreadyExists(_))));
}
}