chore: clean out of order demo

This commit is contained in:
kaichaosun 2026-01-28 17:44:10 +08:00
parent 810dcc28e9
commit 7e1a4a6a79
No known key found for this signature in database
GPG Key ID: 223E0F992F4F03BF
7 changed files with 150 additions and 256 deletions

View File

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

View File

@ -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<DefaultDomain> =
RatchetState::init_sender(shared_secret, bob_keypair.public().clone());
let bob_state: RatchetState<DefaultDomain> =
RatchetState::init_receiver(shared_secret, bob_keypair);
// Collect messages for out-of-order delivery
let mut messages: Vec<(Vec<u8>, 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<u8>, Header)> = Vec::new();
for i in 1..=5 {
let mut alice: RatchetState<DefaultDomain> = 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<DefaultDomain> = 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<DefaultDomain> = 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<DefaultDomain> =
RatchetState::init_sender(shared_secret, bob_keypair.public().clone());
let bob_state: RatchetState<DefaultDomain> =
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<u8>, Header)> = Vec::new();
for i in 1..=5 {
let mut alice: RatchetState<DefaultDomain> = 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<DefaultDomain> = storage.load("bob").unwrap();
let (ct, header) = &messages[idx];
bob.decrypt_message(ct, header.clone()).unwrap();
storage.save("bob", &bob).unwrap();
}
let bob: RatchetState<DefaultDomain> = 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<DefaultDomain> = 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<DefaultDomain> = 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<DefaultDomain> =
RatchetSession::create_sender_session(
&mut alice_storage,
conv_id,
shared_secret,
bob_public,
)
.unwrap();
let mut bob_session: RatchetSession<DefaultDomain> =
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<DefaultDomain> =
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<DefaultDomain> =
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<DefaultDomain> = 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<DefaultDomain> =
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<DefaultDomain> = 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<DefaultDomain> =
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);
}
}

View File

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

View File

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

View File

@ -47,23 +47,23 @@ pub struct RatchetStorage {
}
impl RatchetStorage {
/// Creates a new ratchet storage with the given database.
pub fn new(db: SqliteDb) -> Result<Self, StorageError> {
// 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<Self, StorageError> {
let db = SqliteDb::new(config)?;
Self::new(db)
/// Opens an existing encrypted database file.
pub fn new(path: &str, key: &str) -> Result<Self, StorageError> {
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<Self, StorageError> {
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<Self, StorageError> {
// Initialize schema
db.execute_batch(RATCHET_SCHEMA)?;
Ok(Self { db })
}
/// Saves the ratchet state for a conversation.

View File

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

View File

@ -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, StorageError> {
Self::new(StorageConfig::Encrypted {
path: path,
key: key,
})
}
/// Returns a reference to the underlying connection.
///
/// Use this for domain-specific storage operations.