From 7e1a4a6a792f61072dd31273f7740e9226aece1d Mon Sep 17 00:00:00 2001 From: kaichaosun Date: Wed, 28 Jan 2026 17:44:10 +0800 Subject: [PATCH] chore: clean out of order demo --- double-ratchets/Cargo.toml | 4 +- double-ratchets/examples/out_of_order_demo.rs | 251 ++++++++---------- double-ratchets/examples/storage_demo.rs | 108 +------- double-ratchets/src/lib.rs | 3 - .../src/storage/{ratchet_storage.rs => db.rs} | 24 +- double-ratchets/src/storage/mod.rs | 4 +- storage/src/sqlite.rs | 12 +- 7 files changed, 150 insertions(+), 256 deletions(-) rename double-ratchets/src/storage/{ratchet_storage.rs => db.rs} (97%) diff --git a/double-ratchets/Cargo.toml b/double-ratchets/Cargo.toml index c46fab4..02ec5ae 100644 --- a/double-ratchets/Cargo.toml +++ b/double-ratchets/Cargo.toml @@ -20,9 +20,7 @@ thiserror = "2" blake2 = "0.10.6" safer-ffi = "0.1.13" zeroize = "1.8.2" -storage = { workspace = true, optional = true } +storage = { workspace = true } [features] -default = [] -persist = ["storage"] headers = ["safer-ffi/headers"] diff --git a/double-ratchets/examples/out_of_order_demo.rs b/double-ratchets/examples/out_of_order_demo.rs index ad882d0..d40c6a5 100644 --- a/double-ratchets/examples/out_of_order_demo.rs +++ b/double-ratchets/examples/out_of_order_demo.rs @@ -1,168 +1,149 @@ //! Demonstrates out-of-order message handling with skipped keys persistence. //! -//! Run with: cargo run --example out_of_order_demo --features persist +//! Run with: cargo run --example out_of_order_demo -p double-ratchets -#[cfg(feature = "persist")] -use double_ratchets::{ - InstallationKeyPair, RatchetState, RatchetStorage, StorageConfig, hkdf::DefaultDomain, - state::Header, -}; +use double_ratchets::{InstallationKeyPair, RatchetSession, RatchetStorage, hkdf::DefaultDomain}; fn main() { println!("=== Out-of-Order Message Handling Demo ===\n"); - #[cfg(feature = "persist")] - run_demo(); - #[cfg(not(feature = "persist"))] - println!("(skipped - enable 'persist' feature)"); -} - -#[cfg(feature = "persist")] -fn run_demo() { - let mut storage = - RatchetStorage::with_config(StorageConfig::InMemory).expect("Failed to create storage"); - // Setup + ensure_tmp_directory(); + let alice_db_path = "./tmp/out_of_order_demo_alice.db"; + let bob_db_path = "./tmp/out_of_order_demo_bob.db"; + let encryption_key = "super-secret-key-123!"; + let _ = std::fs::remove_file(alice_db_path); + let _ = std::fs::remove_file(bob_db_path); + let shared_secret = [0x42u8; 32]; let bob_keypair = InstallationKeyPair::generate(); + let bob_public = bob_keypair.public().clone(); + let conv_id = "out_of_order_conv"; - let alice_state: RatchetState = - RatchetState::init_sender(shared_secret, bob_keypair.public().clone()); - let bob_state: RatchetState = - RatchetState::init_receiver(shared_secret, bob_keypair); + // Collect messages for out-of-order delivery + let mut messages: Vec<(Vec, double_ratchets::Header)> = Vec::new(); - storage.save("alice", &alice_state).unwrap(); - storage.save("bob", &bob_state).unwrap(); - - // === Alice sends 5 messages === - println!("Alice sends 5 messages..."); - let mut messages: Vec<(Vec, Header)> = Vec::new(); - - for i in 1..=5 { - let mut alice: RatchetState = storage.load("alice").unwrap(); - let msg = format!("Message #{}", i); - let (ct, header) = alice.encrypt_message(msg.as_bytes()); - storage.save("alice", &alice).unwrap(); - messages.push((ct, header)); - println!(" Sent: \"{}\"", msg); - } - - // === Bob receives messages out of order: 1, 3, 5 === - println!("\nBob receives messages 1, 3, 5 (out of order)..."); - - for &idx in &[0, 2, 4] { - let mut bob: RatchetState = storage.load("bob").unwrap(); - let (ct, header) = &messages[idx]; - let pt = bob - .decrypt_message(ct, header.clone()) - .expect("Decrypt failed"); - storage.save("bob", &bob).unwrap(); - println!(" Received: \"{}\"", String::from_utf8_lossy(&pt)); - } - - let bob: RatchetState = storage.load("bob").unwrap(); - println!("\nBob's skipped_keys count: {}", bob.skipped_keys.len()); - println!(" (Messages 2 and 4 keys are stored for later)"); - - // === Simulate Bob's app restart === - println!("\n--- Simulating Bob's app restart ---"); - drop(storage); - - // In-memory storage doesn't persist across restarts. - // Use file storage to properly demonstrate persistence: - println!(" (Using file storage to demonstrate real persistence)"); - if let Err(e) = std::fs::create_dir_all("./tmp") { - eprintln!("Failed to create tmp directory: {}", e); - return; // Or handle as needed - } - let db_path = "./tmp/out_of_order_demo.db"; - let _ = std::fs::remove_file(db_path); - - // Redo with file storage - let mut storage = RatchetStorage::with_config(StorageConfig::File(db_path.to_string())) - .expect("Failed to create storage"); - - // Re-setup - let bob_keypair = InstallationKeyPair::generate(); - let alice_state: RatchetState = - RatchetState::init_sender(shared_secret, bob_keypair.public().clone()); - let bob_state: RatchetState = - RatchetState::init_receiver(shared_secret, bob_keypair); - - storage.save("alice", &alice_state).unwrap(); - storage.save("bob", &bob_state).unwrap(); - - // Alice sends 5 messages - let mut messages: Vec<(Vec, Header)> = Vec::new(); - for i in 1..=5 { - let mut alice: RatchetState = storage.load("alice").unwrap(); - let msg = format!("Message #{}", i); - let (ct, header) = alice.encrypt_message(msg.as_bytes()); - storage.save("alice", &alice).unwrap(); - messages.push((ct, header)); - } - println!(" Alice sent 5 messages"); - - // Bob receives 1, 3, 5 (skips 2, 4) - for &idx in &[0, 2, 4] { - let mut bob: RatchetState = storage.load("bob").unwrap(); - let (ct, header) = &messages[idx]; - bob.decrypt_message(ct, header.clone()).unwrap(); - storage.save("bob", &bob).unwrap(); - } - - let bob: RatchetState = storage.load("bob").unwrap(); - println!( - " Bob received 1,3,5. Skipped keys stored: {}", - bob.skipped_keys.len() - ); - - // Close and reopen storage (simulating app restart) - drop(storage); - let mut storage = - RatchetStorage::with_config(StorageConfig::File(db_path.to_string())).expect("Failed to reopen"); - - let bob: RatchetState = storage.load("bob").unwrap(); - println!( - "\n After restart, Bob's skipped_keys: {}", - bob.skipped_keys.len() - ); - - // === Now Bob receives the delayed messages === - println!("\nBob receives delayed message 2..."); + // Phase 1: Alice sends 5 messages, Bob receives 1, 3, 5 (skipping 2, 4) { - let mut bob: RatchetState = storage.load("bob").unwrap(); + let mut 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 mut alice_session: RatchetSession = + RatchetSession::create_sender_session( + &mut alice_storage, + conv_id, + shared_secret, + bob_public, + ) + .unwrap(); + + let mut bob_session: RatchetSession = + RatchetSession::create_receiver_session( + &mut bob_storage, + conv_id, + shared_secret, + bob_keypair, + ) + .unwrap(); + + println!(" Sessions created for Alice and Bob"); + + // Alice sends 5 messages + for i in 1..=5 { + let msg = format!("Message #{}", i); + let (ct, header) = alice_session.encrypt_message(msg.as_bytes()).unwrap(); + messages.push((ct, header)); + } + println!(" Alice sent 5 messages"); + + // Bob receives 1, 3, 5 (skips 2, 4) + for &idx in &[0, 2, 4] { + let (ct, header) = &messages[idx]; + bob_session.decrypt_message(ct, header.clone()).unwrap(); + } + + println!( + " Bob received 1,3,5. Skipped keys stored: {}", + bob_session.state().skipped_keys.len() + ); + } + + // 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_session: RatchetSession = + RatchetSession::open(&mut bob_storage, conv_id).unwrap(); + println!( + " After restart, Bob's skipped_keys: {}", + bob_session.state().skipped_keys.len() + ); + } + + // Phase 3: Bob receives the delayed messages + 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 mut bob_session: RatchetSession = + RatchetSession::open(&mut bob_storage, conv_id).unwrap(); + let (ct, header) = &messages[1]; - let pt = bob.decrypt_message(ct, header.clone()).unwrap(); - storage.save("bob", &bob).unwrap(); + let pt = bob_session.decrypt_message(ct, header.clone()).unwrap(); println!(" Received: \"{}\"", String::from_utf8_lossy(&pt)); - println!(" Remaining skipped_keys: {}", bob.skipped_keys.len()); + println!( + " Remaining skipped_keys: {}", + bob_session.state().skipped_keys.len() + ); } println!("\nBob receives delayed message 4..."); - let (ct4, header4) = messages[3].clone(); { - let mut bob: RatchetState = storage.load("bob").unwrap(); - let pt = bob.decrypt_message(&ct4, header4.clone()).unwrap(); - storage.save("bob", &bob).unwrap(); + let mut 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 pt = bob_session.decrypt_message(&ct4, header4.clone()).unwrap(); println!(" Received: \"{}\"", String::from_utf8_lossy(&pt)); - println!(" Remaining skipped_keys: {}", bob.skipped_keys.len()); + println!( + " Remaining skipped_keys: {}", + bob_session.state().skipped_keys.len() + ); } - // === Demonstrate replay protection === + // Phase 4: Demonstrate replay protection println!("\n--- Replay Protection Demo ---"); println!("Trying to decrypt message 4 again (should fail)..."); - { - let mut bob: RatchetState = storage.load("bob").unwrap(); - match bob.decrypt_message(&ct4, header4) { + let mut 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(); + + match bob_session.decrypt_message(&ct4, header4) { Ok(_) => println!(" ERROR: Replay attack succeeded!"), - Err(e) => println!(" Correctly rejected: {:?}", e), + Err(e) => println!(" Correctly rejected: {}", e), } } // Cleanup - let _ = std::fs::remove_file(db_path); + let _ = std::fs::remove_file(alice_db_path); + let _ = std::fs::remove_file(bob_db_path); println!("\n=== Demo Complete ==="); } + +fn ensure_tmp_directory() { + if let Err(e) = std::fs::create_dir_all("./tmp") { + eprintln!("Failed to create tmp directory: {}", e); + } +} diff --git a/double-ratchets/examples/storage_demo.rs b/double-ratchets/examples/storage_demo.rs index cc2d692..9b08a1e 100644 --- a/double-ratchets/examples/storage_demo.rs +++ b/double-ratchets/examples/storage_demo.rs @@ -1,83 +1,12 @@ //! Demonstrates SQLite storage for Double Ratchet state persistence. //! -//! Run with: cargo run --example storage_demo --features persist +//! Run with: cargo run --example storage_demo -p double-ratchets -#[cfg(feature = "persist")] -use double_ratchets::{ - InstallationKeyPair, RatchetSession, RatchetStorage, StorageConfig, hkdf::PrivateV1Domain, -}; +use double_ratchets::{InstallationKeyPair, RatchetSession, RatchetStorage, hkdf::PrivateV1Domain}; fn main() { println!("=== Double Ratchet Storage Demo ===\n"); - // Demo 1: In-memory storage (for testing) - println!("--- Demo 1: In-Memory Storage ---"); - #[cfg(feature = "persist")] - demo_in_memory(); - #[cfg(not(feature = "persist"))] - println!(" (skipped - enable 'persist' feature)"); - - // Demo 2: File-based storage (for local development) - println!("\n--- Demo 2: File-Based Storage ---"); - #[cfg(feature = "persist")] - demo_file_storage(); - #[cfg(not(feature = "persist"))] - println!(" (skipped - enable 'persist' feature)"); - - // Demo 3: SQLCipher encrypted storage (for production) - println!("\n--- Demo 3: SQLCipher Encrypted Storage ---"); - #[cfg(feature = "persist")] - demo_sqlcipher(); - #[cfg(not(feature = "persist"))] - println!(" (skipped - enable 'persist' feature)"); -} - -#[cfg(feature = "persist")] -fn demo_in_memory() { - let mut alice_storage = - RatchetStorage::with_config(StorageConfig::InMemory).expect("Failed to create storage"); - let mut bob_storage = - RatchetStorage::with_config(StorageConfig::InMemory).expect("Failed to create storage"); - run_conversation(&mut alice_storage, &mut bob_storage); -} - -#[cfg(feature = "persist")] -fn demo_file_storage() { - ensure_tmp_directory(); - - let db_path_alice = "./tmp/double_ratchet_demo_alice.db"; - let db_path_bob = "./tmp/double_ratchet_demo_bob.db"; - let _ = std::fs::remove_file(db_path_alice); - let _ = std::fs::remove_file(db_path_bob); - - // Initial conversation - { - let mut alice_storage = RatchetStorage::with_config(StorageConfig::File(db_path_alice.to_string())) - .expect("Failed to create storage"); - - let mut bob_storage = RatchetStorage::with_config(StorageConfig::File(db_path_bob.to_string())) - .expect("Failed to create storage"); - - println!(" Database created at: {}, {}", db_path_alice, db_path_bob); - run_conversation(&mut alice_storage, &mut bob_storage); - } - - // Simulate restart - reopen and continue - println!("\n Simulating application restart..."); - { - let mut alice_storage = RatchetStorage::with_config(StorageConfig::File(db_path_alice.to_string())) - .expect("Failed to reopen storage"); - let mut bob_storage = RatchetStorage::with_config(StorageConfig::File(db_path_bob.to_string())) - .expect("Failed to reopen storage"); - continue_after_restart(&mut alice_storage, &mut bob_storage); - } - - let _ = std::fs::remove_file(db_path_alice); - let _ = std::fs::remove_file(db_path_bob); -} - -#[cfg(feature = "persist")] -fn demo_sqlcipher() { ensure_tmp_directory(); let alice_db_path = "./tmp/double_ratchet_encrypted_alice.db"; let bob_db_path = "./tmp/double_ratchet_encrypted_bob.db"; @@ -87,16 +16,10 @@ fn demo_sqlcipher() { // Initial conversation with encryption { - let mut alice_storage = RatchetStorage::with_config(StorageConfig::Encrypted { - path: alice_db_path.to_string(), - key: encryption_key.to_string(), - }) - .expect("Failed to create encrypted storage"); - let mut bob_storage = RatchetStorage::with_config(StorageConfig::Encrypted { - path: bob_db_path.to_string(), - key: encryption_key.to_string(), - }) - .expect("Failed to create encrypted storage"); + let mut 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) + .expect("Failed to create bob encrypted storage"); println!( " Encrypted database created at: {}, {}", alice_db_path, bob_db_path @@ -107,16 +30,10 @@ fn demo_sqlcipher() { // Restart with correct key println!("\n Simulating restart with encryption key..."); { - let mut alice_storage = RatchetStorage::with_config(StorageConfig::Encrypted { - path: alice_db_path.to_string(), - key: encryption_key.to_string(), - }) - .expect("Failed to create encrypted storage"); - let mut bob_storage = RatchetStorage::with_config(StorageConfig::Encrypted { - path: bob_db_path.to_string(), - key: encryption_key.to_string(), - }) - .expect("Failed to create encrypted storage"); + let mut 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) + .expect("Failed to create bob encrypted storage"); continue_after_restart(&mut alice_storage, &mut bob_storage); } @@ -124,17 +41,15 @@ fn demo_sqlcipher() { let _ = std::fs::remove_file(bob_db_path); } -#[allow(dead_code)] fn ensure_tmp_directory() { if let Err(e) = std::fs::create_dir_all("./tmp") { eprintln!("Failed to create tmp directory: {}", e); - return; // Or handle as needed + return; } } /// Simulates a conversation between Alice and Bob. /// Each party saves/loads state from storage for each operation. -#[cfg(feature = "persist")] fn run_conversation(alice_storage: &mut RatchetStorage, bob_storage: &mut RatchetStorage) { // === Setup: Simulate X3DH key exchange === let shared_secret = [0x42u8; 32]; // In reality, this comes from X3DH @@ -206,7 +121,6 @@ fn run_conversation(alice_storage: &mut RatchetStorage, bob_storage: &mut Ratche ); } -#[cfg(feature = "persist")] fn continue_after_restart(alice_storage: &mut RatchetStorage, bob_storage: &mut RatchetStorage) { // Load persisted states let conv_id = "conv1"; diff --git a/double-ratchets/src/lib.rs b/double-ratchets/src/lib.rs index c055348..21bb81a 100644 --- a/double-ratchets/src/lib.rs +++ b/double-ratchets/src/lib.rs @@ -4,13 +4,10 @@ pub mod ffi; pub mod hkdf; pub mod keypair; pub mod state; -#[cfg(feature = "persist")] pub mod storage; pub mod types; pub use keypair::InstallationKeyPair; pub use state::{Header, RatchetState, SkippedKey}; -#[cfg(feature = "persist")] pub use storage::StorageConfig; -#[cfg(feature = "persist")] pub use storage::{RatchetSession, RatchetStorage, SessionError}; diff --git a/double-ratchets/src/storage/ratchet_storage.rs b/double-ratchets/src/storage/db.rs similarity index 97% rename from double-ratchets/src/storage/ratchet_storage.rs rename to double-ratchets/src/storage/db.rs index ad3b549..d69c651 100644 --- a/double-ratchets/src/storage/ratchet_storage.rs +++ b/double-ratchets/src/storage/db.rs @@ -47,23 +47,23 @@ pub struct RatchetStorage { } impl RatchetStorage { - /// Creates a new ratchet storage with the given database. - pub fn new(db: SqliteDb) -> Result { - // Initialize schema - db.execute_batch(RATCHET_SCHEMA)?; - Ok(Self { db }) - } - - /// Creates a new ratchet storage with the given configuration. - pub fn with_config(config: storage::StorageConfig) -> Result { - let db = SqliteDb::new(config)?; - Self::new(db) + /// Opens an existing encrypted database file. + pub fn new(path: &str, key: &str) -> Result { + let db = SqliteDb::sqlcipher(path.to_string(), key.to_string())?; + Self::new_internal(db) } /// Creates an in-memory storage (useful for testing). pub fn in_memory() -> Result { let db = SqliteDb::in_memory()?; - Self::new(db) + Self::new_internal(db) + } + + /// Creates a new ratchet storage with the given database. + fn new_internal(db: SqliteDb) -> Result { + // Initialize schema + db.execute_batch(RATCHET_SCHEMA)?; + Ok(Self { db }) } /// Saves the ratchet state for a conversation. diff --git a/double-ratchets/src/storage/mod.rs b/double-ratchets/src/storage/mod.rs index 1625f95..490799b 100644 --- a/double-ratchets/src/storage/mod.rs +++ b/double-ratchets/src/storage/mod.rs @@ -3,11 +3,11 @@ //! This module provides storage implementations for the double ratchet state, //! built on top of the shared `storage` crate. -mod ratchet_storage; +mod db; mod session; mod types; -pub use ratchet_storage::RatchetStorage; +pub use db::RatchetStorage; pub use session::{RatchetSession, SessionError}; pub use storage::{SqliteDb, StorageConfig, StorageError}; pub use types::RatchetStateRecord; diff --git a/storage/src/sqlite.rs b/storage/src/sqlite.rs index 4b0c225..829fe55 100644 --- a/storage/src/sqlite.rs +++ b/storage/src/sqlite.rs @@ -13,10 +13,7 @@ pub enum StorageConfig { /// File-based SQLite database. File(String), /// SQLCipher encrypted database. - Encrypted { - path: String, - key: String, - }, + Encrypted { path: String, key: String }, } /// SQLite database wrapper. @@ -58,6 +55,13 @@ impl SqliteDb { Self::new(StorageConfig::InMemory) } + pub fn sqlcipher(path: String, key: String) -> Result { + Self::new(StorageConfig::Encrypted { + path: path, + key: key, + }) + } + /// Returns a reference to the underlying connection. /// /// Use this for domain-specific storage operations.