mirror of
https://github.com/logos-messaging/libchat.git
synced 2026-02-09 16:33:10 +00:00
Split storage crate to abstract database layer (#30)
* feat: shared storage crate * chore: remove backup codes * chore: remove feature gates * chore: clean out of order demo * chore: refactor create session * chore: shorten error name * chore: clean errors * chore: remove table exist check * chore: remove unused traits * chore: remove unused functions. * chore: use tempfile for examples
This commit is contained in:
kaichao
GitHub
parent
7c580b5896
commit
74695877fa
127
Cargo.lock
generated
127
Cargo.lock
generated
@ -221,9 +221,10 @@ dependencies = [
|
||||
"hkdf",
|
||||
"rand",
|
||||
"rand_core",
|
||||
"rusqlite",
|
||||
"safer-ffi",
|
||||
"serde",
|
||||
"storage",
|
||||
"tempfile",
|
||||
"thiserror",
|
||||
"x25519-dalek",
|
||||
"zeroize",
|
||||
@ -265,6 +266,16 @@ version = "1.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
|
||||
|
||||
[[package]]
|
||||
name = "errno"
|
||||
version = "0.3.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ext-trait"
|
||||
version = "1.0.1"
|
||||
@ -312,6 +323,12 @@ version = "0.1.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a"
|
||||
|
||||
[[package]]
|
||||
name = "fastrand"
|
||||
version = "2.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
|
||||
|
||||
[[package]]
|
||||
name = "fiat-crypto"
|
||||
version = "0.2.9"
|
||||
@ -361,6 +378,18 @@ dependencies = [
|
||||
"wasi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "getrandom"
|
||||
version = "0.3.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"libc",
|
||||
"r-efi",
|
||||
"wasip2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hashbrown"
|
||||
version = "0.15.5"
|
||||
@ -464,6 +493,12 @@ dependencies = [
|
||||
"vcpkg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "linux-raw-sys"
|
||||
version = "0.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "df1d3c3b53da64cf5760482273a98e575c651a67eec7f77df96b5b642de8f039"
|
||||
|
||||
[[package]]
|
||||
name = "logos-chat"
|
||||
version = "0.1.0"
|
||||
@ -501,6 +536,12 @@ version = "2.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273"
|
||||
|
||||
[[package]]
|
||||
name = "once_cell"
|
||||
version = "1.21.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
|
||||
|
||||
[[package]]
|
||||
name = "opaque-debug"
|
||||
version = "0.3.1"
|
||||
@ -509,9 +550,9 @@ checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381"
|
||||
|
||||
[[package]]
|
||||
name = "openssl-src"
|
||||
version = "300.5.4+3.5.4"
|
||||
version = "300.5.5+3.5.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a507b3792995dae9b0df8a1c1e3771e8418b7c2d9f0baeba32e6fe8b06c7cb72"
|
||||
checksum = "3f1787d533e03597a7934fd0a765f0d28e94ecc5fb7789f8053b1e699a56f709"
|
||||
dependencies = [
|
||||
"cc",
|
||||
]
|
||||
@ -631,6 +672,12 @@ dependencies = [
|
||||
"proc-macro2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "r-efi"
|
||||
version = "5.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.8.5"
|
||||
@ -658,7 +705,7 @@ version = "0.6.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
|
||||
dependencies = [
|
||||
"getrandom",
|
||||
"getrandom 0.2.17",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@ -684,6 +731,19 @@ dependencies = [
|
||||
"semver",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustix"
|
||||
version = "1.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "146c9e247ccc180c1f61615433868c99f3de3ae256a30a43b49f67c2d9171f34"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustversion"
|
||||
version = "1.0.22"
|
||||
@ -847,6 +907,14 @@ dependencies = [
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "storage"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"rusqlite",
|
||||
"thiserror",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "subtle"
|
||||
version = "2.6.1"
|
||||
@ -875,6 +943,19 @@ dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tempfile"
|
||||
version = "3.24.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "655da9c7eb6305c55742045d5a8d2037996d61d8de95806335c7c86ce0f82e9c"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
"getrandom 0.3.4",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "2.0.18"
|
||||
@ -980,6 +1061,30 @@ version = "0.11.1+wasi-snapshot-preview1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
|
||||
|
||||
[[package]]
|
||||
name = "wasip2"
|
||||
version = "1.0.2+wasi-0.2.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9517f9239f02c069db75e65f174b3da828fe5f5b945c4dd26bd25d89c03ebcf5"
|
||||
dependencies = [
|
||||
"wit-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-link"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.61.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "winnow"
|
||||
version = "0.7.14"
|
||||
@ -989,6 +1094,12 @@ dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wit-bindgen"
|
||||
version = "0.51.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5"
|
||||
|
||||
[[package]]
|
||||
name = "with_builtin_macros"
|
||||
version = "0.0.3"
|
||||
@ -1039,18 +1150,18 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "zerocopy"
|
||||
version = "0.8.34"
|
||||
version = "0.8.35"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "71ddd76bcebeed25db614f82bf31a9f4222d3fbba300e6fb6c00afa26cbd4d9d"
|
||||
checksum = "fdea86ddd5568519879b8187e1cf04e24fce28f7fe046ceecbce472ff19a2572"
|
||||
dependencies = [
|
||||
"zerocopy-derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zerocopy-derive"
|
||||
version = "0.8.34"
|
||||
version = "0.8.35"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d8187381b52e32220d50b255276aa16a084ec0a9017a0ca2152a1f55c539758d"
|
||||
checksum = "0c15e1b46eff7c6c91195752e0eeed8ef040e391cdece7c25376957d5f15df22"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
|
||||
@ -6,7 +6,9 @@ members = [
|
||||
"conversations",
|
||||
"crypto",
|
||||
"double-ratchets",
|
||||
"storage",
|
||||
]
|
||||
|
||||
[workspace.dependencies]
|
||||
blake2 = "0.10"
|
||||
storage = { path = "storage" }
|
||||
|
||||
@ -20,11 +20,11 @@ thiserror = "2"
|
||||
blake2 = "0.10.6"
|
||||
safer-ffi = "0.1.13"
|
||||
zeroize = "1.8.2"
|
||||
storage = { workspace = true }
|
||||
serde = "1.0"
|
||||
rusqlite = { version = "0.35", optional = true, features = ["bundled"] }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
storage = ["rusqlite"]
|
||||
sqlcipher = ["storage", "rusqlite/bundled-sqlcipher-vendored-openssl"]
|
||||
headers = ["safer-ffi/headers"]
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
@ -1,4 +1,4 @@
|
||||
use double_ratchets::{InstallationKeyPair, RatchetState, hkdf::PrivateV1Domain};
|
||||
use double_ratchets::{InstallationKeyPair, RatchetState};
|
||||
|
||||
fn main() {
|
||||
// === Initial shared secret (X3DH / prekey result in real systems) ===
|
||||
@ -6,9 +6,8 @@ fn main() {
|
||||
|
||||
let bob_dh = InstallationKeyPair::generate();
|
||||
|
||||
let mut alice: RatchetState<PrivateV1Domain> =
|
||||
RatchetState::init_sender(shared_secret, bob_dh.public().clone());
|
||||
let mut bob: RatchetState<PrivateV1Domain> = RatchetState::init_receiver(shared_secret, bob_dh);
|
||||
let mut alice: RatchetState = RatchetState::init_sender(shared_secret, bob_dh.public().clone());
|
||||
let mut bob: RatchetState = RatchetState::init_receiver(shared_secret, bob_dh);
|
||||
|
||||
let (ciphertext, header) = alice.encrypt_message(b"Hello Bob!");
|
||||
|
||||
|
||||
@ -1,166 +1,140 @@
|
||||
//! Demonstrates out-of-order message handling with skipped keys persistence.
|
||||
//!
|
||||
//! Run with: cargo run --example out_of_order_demo --features storage
|
||||
//! Run with: cargo run --example out_of_order_demo -p double-ratchets
|
||||
|
||||
#[cfg(feature = "storage")]
|
||||
use double_ratchets::{
|
||||
InstallationKeyPair, RatchetState, SqliteStorage, StorageConfig, hkdf::DefaultDomain,
|
||||
state::Header,
|
||||
};
|
||||
use double_ratchets::{InstallationKeyPair, RatchetSession, RatchetStorage};
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
fn main() {
|
||||
println!("=== Out-of-Order Message Handling Demo (skipped - enable 'storage' feature) ===\n");
|
||||
println!("=== Out-of-Order Message Handling Demo ===\n");
|
||||
|
||||
#[cfg(feature = "storage")]
|
||||
run_demo();
|
||||
}
|
||||
let alice_db_file = NamedTempFile::new().unwrap();
|
||||
let alice_db_path = alice_db_file.path().to_str().unwrap();
|
||||
let bob_db_file = NamedTempFile::new().unwrap();
|
||||
let bob_db_path = bob_db_file.path().to_str().unwrap();
|
||||
|
||||
#[cfg(feature = "storage")]
|
||||
fn run_demo() {
|
||||
let mut storage =
|
||||
SqliteStorage::new(StorageConfig::InMemory).expect("Failed to create storage");
|
||||
|
||||
// Setup
|
||||
let shared_secret = [0x42u8; 32];
|
||||
let bob_keypair = InstallationKeyPair::generate();
|
||||
let bob_public = bob_keypair.public().clone();
|
||||
|
||||
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);
|
||||
let conv_id = "out_of_order_conv";
|
||||
let encryption_key = "super-secret-key-123!";
|
||||
|
||||
storage.save("alice", &alice_state).unwrap();
|
||||
storage.save("bob", &bob_state).unwrap();
|
||||
// Collect messages for out-of-order delivery
|
||||
let mut messages: Vec<(Vec<u8>, double_ratchets::Header)> = Vec::new();
|
||||
|
||||
// === 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 = SqliteStorage::new(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 =
|
||||
SqliteStorage::new(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 = 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<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 =
|
||||
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 =
|
||||
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 ===");
|
||||
}
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
use double_ratchets::{InstallationKeyPair, RatchetState, hkdf::PrivateV1Domain};
|
||||
use double_ratchets::{InstallationKeyPair, RatchetState};
|
||||
|
||||
fn main() {
|
||||
// === Initial shared secret (X3DH / prekey result in real systems) ===
|
||||
@ -6,9 +6,8 @@ fn main() {
|
||||
|
||||
let bob_dh = InstallationKeyPair::generate();
|
||||
|
||||
let mut alice: RatchetState<PrivateV1Domain> =
|
||||
RatchetState::init_sender(shared_secret, bob_dh.public().clone());
|
||||
let mut bob: RatchetState<PrivateV1Domain> = RatchetState::init_receiver(shared_secret, bob_dh);
|
||||
let mut alice: RatchetState = RatchetState::init_sender(shared_secret, bob_dh.public().clone());
|
||||
let mut bob: RatchetState = RatchetState::init_receiver(shared_secret, bob_dh);
|
||||
|
||||
let (ciphertext, header) = alice.encrypt_message(b"Hello Bob!");
|
||||
|
||||
@ -35,9 +34,8 @@ fn main() {
|
||||
|
||||
// === Deserialize alice and bob state from bytes ===
|
||||
println!("Restart alice and bob");
|
||||
let mut alice_new: RatchetState<PrivateV1Domain> =
|
||||
RatchetState::from_bytes(&alice_state).unwrap();
|
||||
let mut bob_new: RatchetState<PrivateV1Domain> = RatchetState::from_bytes(&bob_state).unwrap();
|
||||
let mut alice_new: RatchetState = RatchetState::from_bytes(&alice_state).unwrap();
|
||||
let mut bob_new: RatchetState = RatchetState::from_bytes(&bob_state).unwrap();
|
||||
|
||||
// === Alice sends a message ===
|
||||
let (ciphertext, header) = alice_new.encrypt_message(b"Hello Bob!");
|
||||
|
||||
@ -1,104 +1,26 @@
|
||||
//! Demonstrates SQLite storage for Double Ratchet state persistence.
|
||||
//!
|
||||
//! Run with: cargo run --example storage_demo --features storage
|
||||
//! For SQLCipher: cargo run --example storage_demo --features sqlcipher
|
||||
//! Run with: cargo run --example storage_demo -p double-ratchets
|
||||
|
||||
#[cfg(feature = "storage")]
|
||||
use double_ratchets::{
|
||||
InstallationKeyPair, RatchetSession, SqliteStorage, StorageConfig, hkdf::PrivateV1Domain,
|
||||
};
|
||||
use double_ratchets::{InstallationKeyPair, RatchetSession, RatchetStorage};
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
fn main() {
|
||||
println!("=== Double Ratchet Storage Demo ===\n");
|
||||
|
||||
// Demo 1: In-memory storage (for testing)
|
||||
println!("--- Demo 1: In-Memory Storage (skipped - enable 'storage' feature) ---");
|
||||
#[cfg(feature = "storage")]
|
||||
demo_in_memory();
|
||||
let alice_db_file = NamedTempFile::new().unwrap();
|
||||
let alice_db_path = alice_db_file.path().to_str().unwrap();
|
||||
let bob_db_file = NamedTempFile::new().unwrap();
|
||||
let bob_db_path = bob_db_file.path().to_str().unwrap();
|
||||
|
||||
// Demo 2: File-based storage (for local development)
|
||||
println!("\n--- Demo 2: File-Based Storage (skipped - enable 'storage' feature) ---");
|
||||
#[cfg(feature = "storage")]
|
||||
demo_file_storage();
|
||||
|
||||
// Demo 3: SQLCipher encrypted storage (for production)
|
||||
#[cfg(feature = "sqlcipher")]
|
||||
{
|
||||
println!("\n--- Demo 3: SQLCipher Encrypted Storage ---");
|
||||
demo_sqlcipher();
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "sqlcipher"))]
|
||||
{
|
||||
println!("\n--- Demo 3: SQLCipher (skipped - enable 'sqlcipher' feature) ---");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "storage")]
|
||||
fn demo_in_memory() {
|
||||
let mut alice_storage =
|
||||
SqliteStorage::new(StorageConfig::InMemory).expect("Failed to create storage");
|
||||
let mut bob_storage =
|
||||
SqliteStorage::new(StorageConfig::InMemory).expect("Failed to create storage");
|
||||
run_conversation(&mut alice_storage, &mut bob_storage);
|
||||
}
|
||||
|
||||
#[cfg(feature = "storage")]
|
||||
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 = SqliteStorage::new(StorageConfig::File(db_path_alice.to_string()))
|
||||
.expect("Failed to create storage");
|
||||
|
||||
let mut bob_storage = SqliteStorage::new(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 = SqliteStorage::new(StorageConfig::File(db_path_alice.to_string()))
|
||||
.expect("Failed to reopen storage");
|
||||
let mut bob_storage = SqliteStorage::new(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 = "sqlcipher")]
|
||||
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";
|
||||
let encryption_key = "super-secret-key-123!";
|
||||
let _ = std::fs::remove_file(alice_db_path);
|
||||
let _ = std::fs::remove_file(bob_db_path);
|
||||
|
||||
// Initial conversation with encryption
|
||||
{
|
||||
let mut alice_storage = SqliteStorage::new(StorageConfig::Encrypted {
|
||||
path: alice_db_path.to_string(),
|
||||
key: encryption_key.to_string(),
|
||||
})
|
||||
.expect("Failed to create encrypted storage");
|
||||
let mut bob_storage = SqliteStorage::new(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
|
||||
@ -109,16 +31,10 @@ fn demo_sqlcipher() {
|
||||
// Restart with correct key
|
||||
println!("\n Simulating restart with encryption key...");
|
||||
{
|
||||
let mut alice_storage = SqliteStorage::new(StorageConfig::Encrypted {
|
||||
path: alice_db_path.to_string(),
|
||||
key: encryption_key.to_string(),
|
||||
})
|
||||
.expect("Failed to create encrypted storage");
|
||||
let mut bob_storage = SqliteStorage::new(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);
|
||||
}
|
||||
|
||||
@ -126,25 +42,16 @@ 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
|
||||
}
|
||||
}
|
||||
|
||||
/// Simulates a conversation between Alice and Bob.
|
||||
/// Each party saves/loads state from storage for each operation.
|
||||
#[cfg(feature = "storage")]
|
||||
fn run_conversation(alice_storage: &mut SqliteStorage, bob_storage: &mut SqliteStorage) {
|
||||
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
|
||||
let bob_keypair = InstallationKeyPair::generate();
|
||||
|
||||
let conv_id = "conv1";
|
||||
|
||||
let mut alice_session: RatchetSession<PrivateV1Domain> = RatchetSession::create_sender_session(
|
||||
let mut alice_session: RatchetSession = RatchetSession::create_sender_session(
|
||||
alice_storage,
|
||||
conv_id,
|
||||
shared_secret,
|
||||
@ -152,7 +59,7 @@ fn run_conversation(alice_storage: &mut SqliteStorage, bob_storage: &mut SqliteS
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let mut bob_session: RatchetSession<PrivateV1Domain> =
|
||||
let mut bob_session: RatchetSession =
|
||||
RatchetSession::create_receiver_session(bob_storage, conv_id, shared_secret, bob_keypair)
|
||||
.unwrap();
|
||||
|
||||
@ -208,15 +115,12 @@ fn run_conversation(alice_storage: &mut SqliteStorage, bob_storage: &mut SqliteS
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "storage")]
|
||||
fn continue_after_restart(alice_storage: &mut SqliteStorage, bob_storage: &mut SqliteStorage) {
|
||||
fn continue_after_restart(alice_storage: &mut RatchetStorage, bob_storage: &mut RatchetStorage) {
|
||||
// Load persisted states
|
||||
let conv_id = "conv1";
|
||||
|
||||
let mut alice_session: RatchetSession<PrivateV1Domain> =
|
||||
RatchetSession::open(alice_storage, conv_id).unwrap();
|
||||
let mut bob_session: RatchetSession<PrivateV1Domain> =
|
||||
RatchetSession::open(bob_storage, conv_id).unwrap();
|
||||
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",);
|
||||
|
||||
// Continue conversation
|
||||
|
||||
@ -5,11 +5,10 @@ pub mod hkdf;
|
||||
pub mod keypair;
|
||||
pub mod reader;
|
||||
pub mod state;
|
||||
#[cfg(feature = "storage")]
|
||||
pub mod storage;
|
||||
pub mod types;
|
||||
|
||||
pub use keypair::InstallationKeyPair;
|
||||
pub use state::{Header, RatchetState};
|
||||
#[cfg(feature = "storage")]
|
||||
pub use storage::{RatchetSession, SessionError, SqliteStorage, StorageConfig, StorageError};
|
||||
pub use state::{Header, RatchetState, SkippedKey};
|
||||
pub use storage::StorageConfig;
|
||||
pub use storage::{RatchetSession, RatchetStorage, SessionError};
|
||||
|
||||
320
double-ratchets/src/storage/db.rs
Normal file
320
double-ratchets/src/storage/db.rs
Normal file
@ -0,0 +1,320 @@
|
||||
//! Ratchet-specific storage implementation.
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use storage::{SqliteDb, StorageError, params};
|
||||
|
||||
use super::types::RatchetStateRecord;
|
||||
use crate::{
|
||||
hkdf::HkdfInfo,
|
||||
state::{RatchetState, SkippedKey},
|
||||
};
|
||||
|
||||
/// Schema for ratchet state tables.
|
||||
const RATCHET_SCHEMA: &str = "
|
||||
CREATE TABLE IF NOT EXISTS ratchet_state (
|
||||
conversation_id TEXT PRIMARY KEY,
|
||||
root_key BLOB NOT NULL,
|
||||
sending_chain BLOB,
|
||||
receiving_chain BLOB,
|
||||
dh_self_secret BLOB NOT NULL,
|
||||
dh_remote BLOB,
|
||||
msg_send INTEGER NOT NULL,
|
||||
msg_recv INTEGER NOT NULL,
|
||||
prev_chain_len INTEGER NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS skipped_keys (
|
||||
conversation_id TEXT NOT NULL,
|
||||
public_key BLOB NOT NULL,
|
||||
msg_num INTEGER NOT NULL,
|
||||
message_key BLOB NOT NULL,
|
||||
created_at INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
|
||||
PRIMARY KEY (conversation_id, public_key, msg_num),
|
||||
FOREIGN KEY (conversation_id) REFERENCES ratchet_state(conversation_id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_skipped_keys_conversation
|
||||
ON skipped_keys(conversation_id);
|
||||
";
|
||||
|
||||
/// Ratchet-specific storage operations.
|
||||
///
|
||||
/// This struct wraps a `SqliteDb` and provides domain-specific
|
||||
/// storage operations for ratchet state.
|
||||
pub struct RatchetStorage {
|
||||
db: SqliteDb,
|
||||
}
|
||||
|
||||
impl RatchetStorage {
|
||||
/// 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::run_migration(db)
|
||||
}
|
||||
|
||||
/// Creates an in-memory storage (useful for testing).
|
||||
pub fn in_memory() -> Result<Self, StorageError> {
|
||||
let db = SqliteDb::in_memory()?;
|
||||
Self::run_migration(db)
|
||||
}
|
||||
|
||||
/// Creates a new ratchet storage with the given database.
|
||||
fn run_migration(db: SqliteDb) -> Result<Self, StorageError> {
|
||||
// Initialize schema
|
||||
db.connection().execute_batch(RATCHET_SCHEMA)?;
|
||||
Ok(Self { db })
|
||||
}
|
||||
|
||||
/// Saves the ratchet state for a conversation.
|
||||
pub fn save<D: HkdfInfo>(
|
||||
&mut self,
|
||||
conversation_id: &str,
|
||||
state: &RatchetState<D>,
|
||||
) -> Result<(), StorageError> {
|
||||
let tx = self.db.transaction()?;
|
||||
|
||||
let data = RatchetStateRecord::from(state);
|
||||
let skipped_keys: Vec<SkippedKey> = state.skipped_keys();
|
||||
|
||||
// Upsert main state
|
||||
tx.execute(
|
||||
"
|
||||
INSERT INTO ratchet_state (
|
||||
conversation_id, root_key, sending_chain, receiving_chain,
|
||||
dh_self_secret, dh_remote, msg_send, msg_recv, prev_chain_len
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
|
||||
ON CONFLICT(conversation_id) DO UPDATE SET
|
||||
root_key = excluded.root_key,
|
||||
sending_chain = excluded.sending_chain,
|
||||
receiving_chain = excluded.receiving_chain,
|
||||
dh_self_secret = excluded.dh_self_secret,
|
||||
dh_remote = excluded.dh_remote,
|
||||
msg_send = excluded.msg_send,
|
||||
msg_recv = excluded.msg_recv,
|
||||
prev_chain_len = excluded.prev_chain_len
|
||||
",
|
||||
params![
|
||||
conversation_id,
|
||||
data.root_key.as_slice(),
|
||||
data.sending_chain.as_ref().map(|c| c.as_slice()),
|
||||
data.receiving_chain.as_ref().map(|c| c.as_slice()),
|
||||
data.dh_self_secret.as_slice(),
|
||||
data.dh_remote.as_ref().map(|c| c.as_slice()),
|
||||
data.msg_send,
|
||||
data.msg_recv,
|
||||
data.prev_chain_len,
|
||||
],
|
||||
)?;
|
||||
|
||||
// Sync skipped keys
|
||||
sync_skipped_keys(&tx, conversation_id, skipped_keys)?;
|
||||
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Loads the ratchet state for a conversation.
|
||||
pub fn load<D: HkdfInfo>(
|
||||
&self,
|
||||
conversation_id: &str,
|
||||
) -> Result<RatchetState<D>, StorageError> {
|
||||
let data = self.load_state_data(conversation_id)?;
|
||||
let skipped_keys = self.load_skipped_keys(conversation_id)?;
|
||||
Ok(data.into_ratchet_state(skipped_keys))
|
||||
}
|
||||
|
||||
fn load_state_data(&self, conversation_id: &str) -> Result<RatchetStateRecord, StorageError> {
|
||||
let conn = self.db.connection();
|
||||
let mut stmt = conn.prepare(
|
||||
"
|
||||
SELECT root_key, sending_chain, receiving_chain, dh_self_secret,
|
||||
dh_remote, msg_send, msg_recv, prev_chain_len
|
||||
FROM ratchet_state
|
||||
WHERE conversation_id = ?1
|
||||
",
|
||||
)?;
|
||||
|
||||
stmt.query_row(params![conversation_id], |row| {
|
||||
Ok(RatchetStateRecord {
|
||||
root_key: blob_to_array(row.get::<_, Vec<u8>>(0)?),
|
||||
sending_chain: row.get::<_, Option<Vec<u8>>>(1)?.map(blob_to_array),
|
||||
receiving_chain: row.get::<_, Option<Vec<u8>>>(2)?.map(blob_to_array),
|
||||
dh_self_secret: blob_to_array(row.get::<_, Vec<u8>>(3)?),
|
||||
dh_remote: row.get::<_, Option<Vec<u8>>>(4)?.map(blob_to_array),
|
||||
msg_send: row.get(5)?,
|
||||
msg_recv: row.get(6)?,
|
||||
prev_chain_len: row.get(7)?,
|
||||
})
|
||||
})
|
||||
.map_err(|e| match e {
|
||||
storage::RusqliteError::QueryReturnedNoRows => {
|
||||
StorageError::NotFound(conversation_id.to_string())
|
||||
}
|
||||
e => StorageError::Database(e.to_string()),
|
||||
})
|
||||
}
|
||||
|
||||
fn load_skipped_keys(&self, conversation_id: &str) -> Result<Vec<SkippedKey>, StorageError> {
|
||||
let conn = self.db.connection();
|
||||
let mut stmt = conn.prepare(
|
||||
"
|
||||
SELECT public_key, msg_num, message_key
|
||||
FROM skipped_keys
|
||||
WHERE conversation_id = ?1
|
||||
",
|
||||
)?;
|
||||
|
||||
let rows = stmt.query_map(params![conversation_id], |row| {
|
||||
Ok(SkippedKey {
|
||||
public_key: blob_to_array(row.get::<_, Vec<u8>>(0)?),
|
||||
msg_num: row.get(1)?,
|
||||
message_key: blob_to_array(row.get::<_, Vec<u8>>(2)?),
|
||||
})
|
||||
})?;
|
||||
|
||||
rows.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(|e| StorageError::Database(e.to_string()))
|
||||
}
|
||||
|
||||
/// Checks if a conversation exists.
|
||||
pub fn exists(&self, conversation_id: &str) -> Result<bool, StorageError> {
|
||||
let conn = self.db.connection();
|
||||
let count: i64 = conn.query_row(
|
||||
"SELECT COUNT(*) FROM ratchet_state WHERE conversation_id = ?1",
|
||||
params![conversation_id],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
Ok(count > 0)
|
||||
}
|
||||
|
||||
/// Deletes a conversation and its skipped keys.
|
||||
pub fn delete(&mut self, conversation_id: &str) -> Result<(), StorageError> {
|
||||
let tx = self.db.transaction()?;
|
||||
tx.execute(
|
||||
"DELETE FROM skipped_keys WHERE conversation_id = ?1",
|
||||
params![conversation_id],
|
||||
)?;
|
||||
tx.execute(
|
||||
"DELETE FROM ratchet_state WHERE conversation_id = ?1",
|
||||
params![conversation_id],
|
||||
)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Cleans up old skipped keys older than the given age in seconds.
|
||||
pub fn cleanup_old_skipped_keys(&mut self, max_age_secs: i64) -> Result<usize, StorageError> {
|
||||
let conn = self.db.connection();
|
||||
let deleted = conn.execute(
|
||||
"DELETE FROM skipped_keys WHERE created_at < strftime('%s', 'now') - ?1",
|
||||
params![max_age_secs],
|
||||
)?;
|
||||
Ok(deleted)
|
||||
}
|
||||
}
|
||||
|
||||
/// Syncs skipped keys efficiently by computing diff and only inserting/deleting changes.
|
||||
fn sync_skipped_keys(
|
||||
tx: &storage::Transaction,
|
||||
conversation_id: &str,
|
||||
current_keys: Vec<SkippedKey>,
|
||||
) -> Result<(), StorageError> {
|
||||
// Get existing keys from DB (just the identifiers)
|
||||
let mut stmt =
|
||||
tx.prepare("SELECT public_key, msg_num FROM skipped_keys WHERE conversation_id = ?1")?;
|
||||
let existing: HashSet<([u8; 32], u32)> = stmt
|
||||
.query_map(params![conversation_id], |row| {
|
||||
Ok((
|
||||
blob_to_array(row.get::<_, Vec<u8>>(0)?),
|
||||
row.get::<_, u32>(1)?,
|
||||
))
|
||||
})?
|
||||
.filter_map(|r| r.ok())
|
||||
.collect();
|
||||
|
||||
// Build set of current keys
|
||||
let current_set: HashSet<([u8; 32], u32)> = current_keys
|
||||
.iter()
|
||||
.map(|sk| (sk.public_key, sk.msg_num))
|
||||
.collect();
|
||||
|
||||
// Delete keys that were removed (used for decryption)
|
||||
for (pk, msg_num) in existing.difference(¤t_set) {
|
||||
tx.execute(
|
||||
"DELETE FROM skipped_keys WHERE conversation_id = ?1 AND public_key = ?2 AND msg_num = ?3",
|
||||
params![conversation_id, pk.as_slice(), msg_num],
|
||||
)?;
|
||||
}
|
||||
|
||||
// Insert new keys
|
||||
for sk in ¤t_keys {
|
||||
let key = (sk.public_key, sk.msg_num);
|
||||
if !existing.contains(&key) {
|
||||
tx.execute(
|
||||
"INSERT INTO skipped_keys (conversation_id, public_key, msg_num, message_key)
|
||||
VALUES (?1, ?2, ?3, ?4)",
|
||||
params![
|
||||
conversation_id,
|
||||
sk.public_key.as_slice(),
|
||||
sk.msg_num,
|
||||
sk.message_key.as_slice(),
|
||||
],
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn blob_to_array<const N: usize>(blob: Vec<u8>) -> [u8; N] {
|
||||
blob.try_into()
|
||||
.unwrap_or_else(|v: Vec<u8>| panic!("Expected {} bytes, got {}", N, v.len()))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::{keypair::InstallationKeyPair, state::RatchetState, types::SharedSecret};
|
||||
|
||||
fn create_test_state() -> (RatchetState, SharedSecret) {
|
||||
let shared_secret = [0x42u8; 32];
|
||||
let bob_keypair = InstallationKeyPair::generate();
|
||||
let state = RatchetState::init_sender(shared_secret, bob_keypair.public().clone());
|
||||
(state, shared_secret)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_save_and_load() {
|
||||
let mut storage = RatchetStorage::in_memory().unwrap();
|
||||
let (state, _) = create_test_state();
|
||||
|
||||
storage.save("conv1", &state).unwrap();
|
||||
let loaded: RatchetState = storage.load("conv1").unwrap();
|
||||
|
||||
assert_eq!(state.root_key, loaded.root_key);
|
||||
assert_eq!(state.msg_send, loaded.msg_send);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exists() {
|
||||
let mut storage = RatchetStorage::in_memory().unwrap();
|
||||
let (state, _) = create_test_state();
|
||||
|
||||
assert!(!storage.exists("conv1").unwrap());
|
||||
storage.save("conv1", &state).unwrap();
|
||||
assert!(storage.exists("conv1").unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delete() {
|
||||
let mut storage = RatchetStorage::in_memory().unwrap();
|
||||
let (state, _) = create_test_state();
|
||||
|
||||
storage.save("conv1", &state).unwrap();
|
||||
assert!(storage.exists("conv1").unwrap());
|
||||
|
||||
storage.delete("conv1").unwrap();
|
||||
assert!(!storage.exists("conv1").unwrap());
|
||||
}
|
||||
}
|
||||
16
double-ratchets/src/storage/errors.rs
Normal file
16
double-ratchets/src/storage/errors.rs
Normal file
@ -0,0 +1,16 @@
|
||||
use storage::StorageError;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::errors::RatchetError;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum SessionError {
|
||||
#[error("storage error: {0}")]
|
||||
Storage(#[from] StorageError),
|
||||
|
||||
#[error("ratchet error: {0}")]
|
||||
Ratchet(#[from] RatchetError),
|
||||
|
||||
#[error("conversation already exists: {0}")]
|
||||
ConvAlreadyExists(String),
|
||||
}
|
||||
@ -1,5 +1,15 @@
|
||||
mod session;
|
||||
mod sqlite;
|
||||
//! Storage module for persisting ratchet state.
|
||||
//!
|
||||
//! This module provides storage implementations for the double ratchet state,
|
||||
//! built on top of the shared `storage` crate.
|
||||
|
||||
pub use session::{RatchetSession, SessionError};
|
||||
pub use sqlite::{SqliteStorage, StorageConfig};
|
||||
mod db;
|
||||
mod errors;
|
||||
mod session;
|
||||
mod types;
|
||||
|
||||
pub use db::RatchetStorage;
|
||||
pub use errors::SessionError;
|
||||
pub use session::RatchetSession;
|
||||
pub use storage::{SqliteDb, StorageConfig, StorageError};
|
||||
pub use types::RatchetStateRecord;
|
||||
|
||||
@ -1,58 +1,30 @@
|
||||
//! Session wrapper for automatic state persistence.
|
||||
|
||||
use x25519_dalek::PublicKey;
|
||||
|
||||
use crate::{
|
||||
InstallationKeyPair,
|
||||
errors::RatchetError,
|
||||
hkdf::HkdfInfo,
|
||||
InstallationKeyPair, SessionError,
|
||||
hkdf::{DefaultDomain, HkdfInfo},
|
||||
state::{Header, RatchetState},
|
||||
types::SharedSecret,
|
||||
};
|
||||
|
||||
use super::{SqliteStorage, StorageError};
|
||||
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> {
|
||||
storage: &'a mut SqliteStorage,
|
||||
pub struct RatchetSession<'a, D: HkdfInfo + Clone = DefaultDomain> {
|
||||
storage: &'a mut RatchetStorage,
|
||||
conversation_id: String,
|
||||
state: RatchetState<D>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum SessionError {
|
||||
Storage(StorageError),
|
||||
Ratchet(RatchetError),
|
||||
}
|
||||
|
||||
impl From<StorageError> for SessionError {
|
||||
fn from(e: StorageError) -> Self {
|
||||
SessionError::Storage(e)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<RatchetError> for SessionError {
|
||||
fn from(e: RatchetError) -> Self {
|
||||
SessionError::Ratchet(e)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for SessionError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
SessionError::Storage(e) => write!(f, "storage error: {}", e),
|
||||
SessionError::Ratchet(e) => write!(f, "ratchet error: {}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for SessionError {}
|
||||
|
||||
impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> {
|
||||
/// Opens an existing session from storage.
|
||||
pub fn open(
|
||||
storage: &'a mut SqliteStorage,
|
||||
storage: &'a mut RatchetStorage,
|
||||
conversation_id: impl Into<String>,
|
||||
) -> Result<Self, StorageError> {
|
||||
) -> Result<Self, SessionError> {
|
||||
let conversation_id = conversation_id.into();
|
||||
let state = storage.load(&conversation_id)?;
|
||||
Ok(Self {
|
||||
@ -64,10 +36,10 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> {
|
||||
|
||||
/// Creates a new session and persists the initial state.
|
||||
pub fn create(
|
||||
storage: &'a mut SqliteStorage,
|
||||
storage: &'a mut RatchetStorage,
|
||||
conversation_id: impl Into<String>,
|
||||
state: RatchetState<D>,
|
||||
) -> Result<Self, StorageError> {
|
||||
) -> Result<Self, SessionError> {
|
||||
let conversation_id = conversation_id.into();
|
||||
storage.save(&conversation_id, &state)?;
|
||||
Ok(Self {
|
||||
@ -79,29 +51,31 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> {
|
||||
|
||||
/// Initializes a new session as a sender and persists the initial state.
|
||||
pub fn create_sender_session(
|
||||
storage: &'a mut SqliteStorage,
|
||||
conversation_id: impl Into<String>,
|
||||
storage: &'a mut RatchetStorage,
|
||||
conversation_id: &str,
|
||||
shared_secret: SharedSecret,
|
||||
remote_pub: PublicKey,
|
||||
) -> Result<Self, StorageError> {
|
||||
) -> Result<Self, SessionError> {
|
||||
if storage.exists(conversation_id)? {
|
||||
return Err(SessionError::ConvAlreadyExists(conversation_id.to_string()));
|
||||
}
|
||||
let state = RatchetState::<D>::init_sender(shared_secret, remote_pub);
|
||||
Self::create(storage, conversation_id, state)
|
||||
Ok(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 SqliteStorage,
|
||||
conversation_id: impl Into<String>,
|
||||
storage: &'a mut RatchetStorage,
|
||||
conversation_id: &str,
|
||||
shared_secret: SharedSecret,
|
||||
dh_self: InstallationKeyPair,
|
||||
) -> Result<Self, StorageError> {
|
||||
let conversation_id = conversation_id.into();
|
||||
if storage.exists(&conversation_id)? {
|
||||
return Self::open(storage, conversation_id);
|
||||
) -> Result<Self, SessionError> {
|
||||
if storage.exists(conversation_id)? {
|
||||
return Err(SessionError::ConvAlreadyExists(conversation_id.to_string()));
|
||||
}
|
||||
|
||||
let state = RatchetState::<D>::init_receiver(shared_secret, dh_self);
|
||||
Self::create(storage, conversation_id, state)
|
||||
Ok(Self::create(storage, conversation_id, state)?)
|
||||
}
|
||||
|
||||
/// Encrypts a message and persists the updated state.
|
||||
@ -117,7 +91,7 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> {
|
||||
if let Err(e) = self.storage.save(&self.conversation_id, &self.state) {
|
||||
// Rollback
|
||||
self.state = state_backup;
|
||||
return Err(SessionError::Storage(e));
|
||||
return Err(e.into());
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
@ -139,7 +113,7 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> {
|
||||
Err(e) => {
|
||||
// Rollback on decrypt failure
|
||||
self.state = state_backup;
|
||||
return Err(SessionError::Ratchet(e));
|
||||
return Err(e.into());
|
||||
}
|
||||
};
|
||||
|
||||
@ -147,7 +121,7 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> {
|
||||
if let Err(e) = self.storage.save(&self.conversation_id, &self.state) {
|
||||
// Rollback
|
||||
self.state = state_backup;
|
||||
return Err(SessionError::Storage(e));
|
||||
return Err(e.into());
|
||||
}
|
||||
|
||||
Ok(plaintext)
|
||||
@ -164,8 +138,10 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> {
|
||||
}
|
||||
|
||||
/// Manually saves the current state.
|
||||
pub fn save(&mut self) -> Result<(), StorageError> {
|
||||
self.storage.save(&self.conversation_id, &self.state)
|
||||
pub fn save(&mut self) -> Result<(), SessionError> {
|
||||
self.storage
|
||||
.save(&self.conversation_id, &self.state)
|
||||
.map_err(|error| error.into())
|
||||
}
|
||||
|
||||
pub fn msg_send(&self) -> u32 {
|
||||
@ -180,10 +156,10 @@ impl<'a, D: HkdfInfo + Clone> RatchetSession<'a, D> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::{hkdf::DefaultDomain, keypair::InstallationKeyPair, storage::StorageConfig};
|
||||
use crate::hkdf::DefaultDomain;
|
||||
|
||||
fn create_test_storage() -> SqliteStorage {
|
||||
SqliteStorage::new(StorageConfig::InMemory).unwrap()
|
||||
fn create_test_storage() -> RatchetStorage {
|
||||
RatchetStorage::in_memory().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -307,4 +283,70 @@ mod tests {
|
||||
assert_eq!(session.state().msg_send, 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_sender_session_fails_when_conversation_exists() {
|
||||
let mut 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();
|
||||
}
|
||||
|
||||
// Second creation should fail with ConversationAlreadyExists
|
||||
{
|
||||
let result: Result<RatchetSession<DefaultDomain>, _> =
|
||||
RatchetSession::create_sender_session(
|
||||
&mut storage,
|
||||
"conv1",
|
||||
shared_secret,
|
||||
bob_pub.clone(),
|
||||
);
|
||||
|
||||
assert!(matches!(result, Err(SessionError::ConvAlreadyExists(_))));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_receiver_session_fails_when_conversation_exists() {
|
||||
let mut 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();
|
||||
}
|
||||
|
||||
// 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,
|
||||
);
|
||||
|
||||
assert!(matches!(result, Err(SessionError::ConvAlreadyExists(_))));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,437 +0,0 @@
|
||||
use rusqlite::{Connection, params};
|
||||
|
||||
use super::{RatchetStateRecord, SkippedKey, StorageError};
|
||||
use crate::{hkdf::HkdfInfo, state::RatchetState};
|
||||
|
||||
/// Configuration for SQLite storage.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum StorageConfig {
|
||||
/// In-memory database (for testing).
|
||||
InMemory,
|
||||
/// File-based SQLite database (unencrypted, for local dev).
|
||||
File(String),
|
||||
/// SQLCipher encrypted database (for production).
|
||||
/// Requires the `sqlcipher` feature.
|
||||
#[cfg(feature = "sqlcipher")]
|
||||
Encrypted { path: String, key: String },
|
||||
}
|
||||
|
||||
/// SQLite-based storage for ratchet state.
|
||||
pub struct SqliteStorage {
|
||||
conn: Connection,
|
||||
}
|
||||
|
||||
impl SqliteStorage {
|
||||
/// Creates a new SQLite storage with the given configuration.
|
||||
pub fn new(config: StorageConfig) -> Result<Self, StorageError> {
|
||||
let conn = match config {
|
||||
StorageConfig::InMemory => Connection::open_in_memory()?,
|
||||
StorageConfig::File(path) => Connection::open(path)?,
|
||||
#[cfg(feature = "sqlcipher")]
|
||||
StorageConfig::Encrypted { path, key } => {
|
||||
let conn = Connection::open(path)?;
|
||||
conn.pragma_update(None, "key", &key)?;
|
||||
conn
|
||||
}
|
||||
};
|
||||
|
||||
let storage = Self { conn };
|
||||
storage.init_schema()?;
|
||||
Ok(storage)
|
||||
}
|
||||
|
||||
fn init_schema(&self) -> Result<(), StorageError> {
|
||||
self.conn.execute_batch(
|
||||
"
|
||||
CREATE TABLE IF NOT EXISTS ratchet_state (
|
||||
conversation_id TEXT PRIMARY KEY,
|
||||
root_key BLOB NOT NULL,
|
||||
sending_chain BLOB,
|
||||
receiving_chain BLOB,
|
||||
dh_self_secret BLOB NOT NULL,
|
||||
dh_remote BLOB,
|
||||
msg_send INTEGER NOT NULL,
|
||||
msg_recv INTEGER NOT NULL,
|
||||
prev_chain_len INTEGER NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS skipped_keys (
|
||||
conversation_id TEXT NOT NULL,
|
||||
public_key BLOB NOT NULL,
|
||||
msg_num INTEGER NOT NULL,
|
||||
message_key BLOB NOT NULL,
|
||||
created_at INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
|
||||
PRIMARY KEY (conversation_id, public_key, msg_num),
|
||||
FOREIGN KEY (conversation_id) REFERENCES ratchet_state(conversation_id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_skipped_keys_conversation
|
||||
ON skipped_keys(conversation_id);
|
||||
",
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Saves the ratchet state for a conversation within a transaction.
|
||||
/// Rolls back automatically if any error occurs.
|
||||
pub fn save<D: HkdfInfo>(
|
||||
&mut self,
|
||||
conversation_id: &str,
|
||||
state: &RatchetState<D>,
|
||||
) -> Result<(), StorageError> {
|
||||
let tx = self.conn.transaction()?;
|
||||
|
||||
let data = RatchetStateRecord::from(state);
|
||||
let skipped_keys: Vec<SkippedKey> = state.skipped_keys();
|
||||
|
||||
// Upsert main state
|
||||
tx.execute(
|
||||
"
|
||||
INSERT INTO ratchet_state (
|
||||
conversation_id, root_key, sending_chain, receiving_chain,
|
||||
dh_self_secret, dh_remote, msg_send, msg_recv, prev_chain_len
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
|
||||
ON CONFLICT(conversation_id) DO UPDATE SET
|
||||
root_key = excluded.root_key,
|
||||
sending_chain = excluded.sending_chain,
|
||||
receiving_chain = excluded.receiving_chain,
|
||||
dh_self_secret = excluded.dh_self_secret,
|
||||
dh_remote = excluded.dh_remote,
|
||||
msg_send = excluded.msg_send,
|
||||
msg_recv = excluded.msg_recv,
|
||||
prev_chain_len = excluded.prev_chain_len
|
||||
",
|
||||
params![
|
||||
conversation_id,
|
||||
data.root_key.as_slice(),
|
||||
data.sending_chain.as_ref().map(|c| c.as_slice()),
|
||||
data.receiving_chain.as_ref().map(|c| c.as_slice()),
|
||||
data.dh_self_secret.as_slice(),
|
||||
data.dh_remote.as_ref().map(|c| c.as_slice()),
|
||||
data.msg_send,
|
||||
data.msg_recv,
|
||||
data.prev_chain_len,
|
||||
],
|
||||
)?;
|
||||
|
||||
// Sync skipped keys efficiently - only insert new, delete removed
|
||||
sync_skipped_keys(&tx, conversation_id, skipped_keys)?;
|
||||
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Loads the ratchet state for a conversation.
|
||||
pub fn load<D: HkdfInfo>(
|
||||
&self,
|
||||
conversation_id: &str,
|
||||
) -> Result<RatchetState<D>, StorageError> {
|
||||
let data = self.load_state_data(conversation_id)?;
|
||||
let skipped_keys = self.load_skipped_keys(conversation_id)?;
|
||||
Ok(data.into_ratchet_state(skipped_keys))
|
||||
}
|
||||
|
||||
fn load_state_data(&self, conversation_id: &str) -> Result<RatchetStateRecord, StorageError> {
|
||||
let mut stmt = self.conn.prepare(
|
||||
"
|
||||
SELECT root_key, sending_chain, receiving_chain, dh_self_secret,
|
||||
dh_remote, msg_send, msg_recv, prev_chain_len
|
||||
FROM ratchet_state
|
||||
WHERE conversation_id = ?1
|
||||
",
|
||||
)?;
|
||||
|
||||
stmt.query_row(params![conversation_id], |row| {
|
||||
Ok(RatchetStateRecord {
|
||||
root_key: blob_to_array(row.get::<_, Vec<u8>>(0)?),
|
||||
sending_chain: row.get::<_, Option<Vec<u8>>>(1)?.map(blob_to_array),
|
||||
receiving_chain: row.get::<_, Option<Vec<u8>>>(2)?.map(blob_to_array),
|
||||
dh_self_secret: blob_to_array(row.get::<_, Vec<u8>>(3)?),
|
||||
dh_remote: row.get::<_, Option<Vec<u8>>>(4)?.map(blob_to_array),
|
||||
msg_send: row.get(5)?,
|
||||
msg_recv: row.get(6)?,
|
||||
prev_chain_len: row.get(7)?,
|
||||
})
|
||||
})
|
||||
.map_err(|e| match e {
|
||||
rusqlite::Error::QueryReturnedNoRows => {
|
||||
StorageError::ConversationNotFound(conversation_id.to_string())
|
||||
}
|
||||
e => StorageError::Database(e),
|
||||
})
|
||||
}
|
||||
|
||||
fn load_skipped_keys(&self, conversation_id: &str) -> Result<Vec<SkippedKey>, StorageError> {
|
||||
let mut stmt = self.conn.prepare(
|
||||
"
|
||||
SELECT public_key, msg_num, message_key
|
||||
FROM skipped_keys
|
||||
WHERE conversation_id = ?1
|
||||
",
|
||||
)?;
|
||||
|
||||
let rows = stmt.query_map(params![conversation_id], |row| {
|
||||
Ok(SkippedKey {
|
||||
public_key: blob_to_array(row.get::<_, Vec<u8>>(0)?),
|
||||
msg_num: row.get(1)?,
|
||||
message_key: blob_to_array(row.get::<_, Vec<u8>>(2)?),
|
||||
})
|
||||
})?;
|
||||
|
||||
rows.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(StorageError::Database)
|
||||
}
|
||||
|
||||
/// Checks if a conversation exists.
|
||||
pub fn exists(&self, conversation_id: &str) -> Result<bool, StorageError> {
|
||||
let count: i64 = self.conn.query_row(
|
||||
"SELECT COUNT(*) FROM ratchet_state WHERE conversation_id = ?1",
|
||||
params![conversation_id],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
Ok(count > 0)
|
||||
}
|
||||
|
||||
/// Deletes a conversation and its skipped keys.
|
||||
pub fn delete(&mut self, conversation_id: &str) -> Result<(), StorageError> {
|
||||
let tx = self.conn.transaction()?;
|
||||
tx.execute(
|
||||
"DELETE FROM skipped_keys WHERE conversation_id = ?1",
|
||||
params![conversation_id],
|
||||
)?;
|
||||
tx.execute(
|
||||
"DELETE FROM ratchet_state WHERE conversation_id = ?1",
|
||||
params![conversation_id],
|
||||
)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Cleans up old skipped keys older than the given age in seconds.
|
||||
pub fn cleanup_old_skipped_keys(&mut self, max_age_secs: i64) -> Result<usize, StorageError> {
|
||||
let deleted = self.conn.execute(
|
||||
"DELETE FROM skipped_keys WHERE created_at < strftime('%s', 'now') - ?1",
|
||||
params![max_age_secs],
|
||||
)?;
|
||||
Ok(deleted)
|
||||
}
|
||||
}
|
||||
|
||||
/// Syncs skipped keys efficiently by computing diff and only inserting/deleting changes.
|
||||
fn sync_skipped_keys(
|
||||
tx: &rusqlite::Transaction,
|
||||
conversation_id: &str,
|
||||
current_keys: Vec<SkippedKey>,
|
||||
) -> Result<(), StorageError> {
|
||||
use std::collections::HashSet;
|
||||
|
||||
// Get existing keys from DB (just the identifiers)
|
||||
let mut stmt =
|
||||
tx.prepare("SELECT public_key, msg_num FROM skipped_keys WHERE conversation_id = ?1")?;
|
||||
let existing: HashSet<([u8; 32], u32)> = stmt
|
||||
.query_map(params![conversation_id], |row| {
|
||||
Ok((
|
||||
blob_to_array(row.get::<_, Vec<u8>>(0)?),
|
||||
row.get::<_, u32>(1)?,
|
||||
))
|
||||
})?
|
||||
.filter_map(|r| r.ok())
|
||||
.collect();
|
||||
|
||||
// Build set of current keys
|
||||
let current_set: HashSet<([u8; 32], u32)> = current_keys
|
||||
.iter()
|
||||
.map(|sk| (sk.public_key, sk.msg_num))
|
||||
.collect();
|
||||
|
||||
// Delete keys that were removed (used for decryption)
|
||||
for (pk, msg_num) in existing.difference(¤t_set) {
|
||||
tx.execute(
|
||||
"DELETE FROM skipped_keys WHERE conversation_id = ?1 AND public_key = ?2 AND msg_num = ?3",
|
||||
params![conversation_id, pk.as_slice(), msg_num],
|
||||
)?;
|
||||
}
|
||||
|
||||
// Insert new keys
|
||||
for sk in ¤t_keys {
|
||||
let key = (sk.public_key, sk.msg_num);
|
||||
if !existing.contains(&key) {
|
||||
tx.execute(
|
||||
"INSERT INTO skipped_keys (conversation_id, public_key, msg_num, message_key)
|
||||
VALUES (?1, ?2, ?3, ?4)",
|
||||
params![
|
||||
conversation_id,
|
||||
sk.public_key.as_slice(),
|
||||
sk.msg_num,
|
||||
sk.message_key.as_slice(),
|
||||
],
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn blob_to_array<const N: usize>(blob: Vec<u8>) -> [u8; N] {
|
||||
blob.try_into()
|
||||
.unwrap_or_else(|v: Vec<u8>| panic!("Expected {} bytes, got {}", N, v.len()))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::{hkdf::DefaultDomain, keypair::InstallationKeyPair};
|
||||
|
||||
fn create_test_storage() -> SqliteStorage {
|
||||
SqliteStorage::new(StorageConfig::InMemory).unwrap()
|
||||
}
|
||||
|
||||
fn create_test_state() -> (RatchetState<DefaultDomain>, RatchetState<DefaultDomain>) {
|
||||
let shared_secret = [0x42; 32];
|
||||
let bob_keypair = InstallationKeyPair::generate();
|
||||
let alice = RatchetState::init_sender(shared_secret, bob_keypair.public().clone());
|
||||
let bob = RatchetState::init_receiver(shared_secret, bob_keypair);
|
||||
(alice, bob)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_save_and_load_sender() {
|
||||
let mut storage = create_test_storage();
|
||||
let (alice, _) = create_test_state();
|
||||
|
||||
storage.save("conv1", &alice).unwrap();
|
||||
let loaded: RatchetState<DefaultDomain> = storage.load("conv1").unwrap();
|
||||
|
||||
assert_eq!(alice.root_key, loaded.root_key);
|
||||
assert_eq!(alice.sending_chain, loaded.sending_chain);
|
||||
assert_eq!(alice.receiving_chain, loaded.receiving_chain);
|
||||
assert_eq!(alice.msg_send, loaded.msg_send);
|
||||
assert_eq!(alice.msg_recv, loaded.msg_recv);
|
||||
assert_eq!(alice.prev_chain_len, loaded.prev_chain_len);
|
||||
assert_eq!(
|
||||
alice.dh_self.public().to_bytes(),
|
||||
loaded.dh_self.public().to_bytes()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_save_and_load_receiver() {
|
||||
let mut storage = create_test_storage();
|
||||
let (_, bob) = create_test_state();
|
||||
|
||||
storage.save("conv1", &bob).unwrap();
|
||||
let loaded: RatchetState<DefaultDomain> = storage.load("conv1").unwrap();
|
||||
|
||||
assert_eq!(bob.root_key, loaded.root_key);
|
||||
assert!(loaded.dh_remote.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_load_not_found() {
|
||||
let storage = create_test_storage();
|
||||
let result: Result<RatchetState<DefaultDomain>, _> = storage.load("nonexistent");
|
||||
assert!(matches!(result, Err(StorageError::ConversationNotFound(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_save_with_skipped_keys() {
|
||||
let mut storage = create_test_storage();
|
||||
let (mut alice, mut bob) = create_test_state();
|
||||
|
||||
// Alice sends 3 messages
|
||||
let mut sent = vec![];
|
||||
for i in 0..3 {
|
||||
let plaintext = format!("Message {}", i + 1).into_bytes();
|
||||
let (ct, header) = alice.encrypt_message(&plaintext);
|
||||
sent.push((ct, header, plaintext));
|
||||
}
|
||||
|
||||
// Bob receives 0 and 2, skipping 1
|
||||
bob.decrypt_message(&sent[0].0, sent[0].1.clone()).unwrap();
|
||||
bob.decrypt_message(&sent[2].0, sent[2].1.clone()).unwrap();
|
||||
|
||||
assert_eq!(bob.skipped_keys.len(), 1);
|
||||
|
||||
// Save and reload
|
||||
storage.save("conv1", &bob).unwrap();
|
||||
let mut loaded: RatchetState<DefaultDomain> = storage.load("conv1").unwrap();
|
||||
|
||||
assert_eq!(loaded.skipped_keys.len(), 1);
|
||||
|
||||
// Should be able to decrypt skipped message
|
||||
let pt = loaded
|
||||
.decrypt_message(&sent[1].0, sent[1].1.clone())
|
||||
.unwrap();
|
||||
assert_eq!(pt, sent[1].2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_update_existing() {
|
||||
let mut storage = create_test_storage();
|
||||
let (mut alice, mut bob) = create_test_state();
|
||||
|
||||
storage.save("conv1", &alice).unwrap();
|
||||
|
||||
// Exchange a message
|
||||
let (ct, header) = alice.encrypt_message(b"Hello");
|
||||
bob.decrypt_message(&ct, header).unwrap();
|
||||
|
||||
// Update Alice's state
|
||||
storage.save("conv1", &alice).unwrap();
|
||||
|
||||
let loaded: RatchetState<DefaultDomain> = storage.load("conv1").unwrap();
|
||||
assert_eq!(loaded.msg_send, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exists() {
|
||||
let mut storage = create_test_storage();
|
||||
let (alice, _) = create_test_state();
|
||||
|
||||
assert!(!storage.exists("conv1").unwrap());
|
||||
storage.save("conv1", &alice).unwrap();
|
||||
assert!(storage.exists("conv1").unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delete() {
|
||||
let mut storage = create_test_storage();
|
||||
let (alice, _) = create_test_state();
|
||||
|
||||
storage.save("conv1", &alice).unwrap();
|
||||
assert!(storage.exists("conv1").unwrap());
|
||||
|
||||
storage.delete("conv1").unwrap();
|
||||
assert!(!storage.exists("conv1").unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_continue_conversation_after_reload() {
|
||||
let mut storage = create_test_storage();
|
||||
let (mut alice, mut bob) = create_test_state();
|
||||
|
||||
// Exchange messages
|
||||
let (ct1, h1) = alice.encrypt_message(b"Hello Bob");
|
||||
bob.decrypt_message(&ct1, h1).unwrap();
|
||||
|
||||
let (ct2, h2) = bob.encrypt_message(b"Hello Alice");
|
||||
alice.decrypt_message(&ct2, h2).unwrap();
|
||||
|
||||
// Save both
|
||||
storage.save("alice", &alice).unwrap();
|
||||
storage.save("bob", &bob).unwrap();
|
||||
|
||||
// Reload
|
||||
let mut alice_new: RatchetState<DefaultDomain> = storage.load("alice").unwrap();
|
||||
let mut bob_new: RatchetState<DefaultDomain> = storage.load("bob").unwrap();
|
||||
|
||||
// Continue conversation
|
||||
let (ct3, h3) = alice_new.encrypt_message(b"After reload");
|
||||
let pt3 = bob_new.decrypt_message(&ct3, h3).unwrap();
|
||||
assert_eq!(pt3, b"After reload");
|
||||
|
||||
let (ct4, h4) = bob_new.encrypt_message(b"Reply after reload");
|
||||
let pt4 = alice_new.decrypt_message(&ct4, h4).unwrap();
|
||||
assert_eq!(pt4, b"Reply after reload");
|
||||
}
|
||||
}
|
||||
@ -1,28 +1,12 @@
|
||||
//! Storage types for ratchet state.
|
||||
|
||||
use crate::{
|
||||
hkdf::HkdfInfo,
|
||||
state::{RatchetState, SkippedKey},
|
||||
types::MessageKey,
|
||||
};
|
||||
use thiserror::Error;
|
||||
use x25519_dalek::PublicKey;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum StorageError {
|
||||
#[error("database error: {0}")]
|
||||
Database(#[from] rusqlite::Error),
|
||||
|
||||
#[error("conversation not found: {0}")]
|
||||
ConversationNotFound(String),
|
||||
|
||||
#[error("serialization error")]
|
||||
Serialization,
|
||||
|
||||
#[error("deserialization error")]
|
||||
Deserialization,
|
||||
}
|
||||
|
||||
/// Stored representation of a skipped message key.
|
||||
|
||||
/// Raw state data for storage (without generic parameter).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RatchetStateRecord {
|
||||
@ -42,7 +26,7 @@ impl<D: HkdfInfo> From<&RatchetState<D>> for RatchetStateRecord {
|
||||
root_key: state.root_key,
|
||||
sending_chain: state.sending_chain,
|
||||
receiving_chain: state.receiving_chain,
|
||||
dh_self_secret: state.dh_self.secret_bytes(),
|
||||
dh_self_secret: *state.dh_self.secret_bytes(),
|
||||
dh_remote: state.dh_remote.map(|pk| pk.to_bytes()),
|
||||
msg_send: state.msg_send,
|
||||
msg_recv: state.msg_recv,
|
||||
|
||||
9
storage/Cargo.toml
Normal file
9
storage/Cargo.toml
Normal file
@ -0,0 +1,9 @@
|
||||
[package]
|
||||
name = "storage"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
description = "Shared storage layer for libchat"
|
||||
|
||||
[dependencies]
|
||||
thiserror = "2"
|
||||
rusqlite = { version = "0.35", features = ["bundled-sqlcipher-vendored-openssl"] }
|
||||
35
storage/src/errors.rs
Normal file
35
storage/src/errors.rs
Normal file
@ -0,0 +1,35 @@
|
||||
use thiserror::Error;
|
||||
|
||||
/// Common storage errors.
|
||||
#[derive(Debug, Error)]
|
||||
pub enum StorageError {
|
||||
/// Database error (wraps rusqlite::Error when sqlite feature is enabled).
|
||||
#[error("database error: {0}")]
|
||||
Database(String),
|
||||
|
||||
/// Record not found.
|
||||
#[error("not found: {0}")]
|
||||
NotFound(String),
|
||||
|
||||
/// Serialization error.
|
||||
#[error("serialization error: {0}")]
|
||||
Serialization(String),
|
||||
|
||||
/// Deserialization error.
|
||||
#[error("deserialization error: {0}")]
|
||||
Deserialization(String),
|
||||
|
||||
/// Schema migration error.
|
||||
#[error("migration error: {0}")]
|
||||
Migration(String),
|
||||
|
||||
/// Transaction error.
|
||||
#[error("transaction error: {0}")]
|
||||
Transaction(String),
|
||||
}
|
||||
|
||||
impl From<rusqlite::Error> for StorageError {
|
||||
fn from(e: rusqlite::Error) -> Self {
|
||||
StorageError::Database(e.to_string())
|
||||
}
|
||||
}
|
||||
15
storage/src/lib.rs
Normal file
15
storage/src/lib.rs
Normal file
@ -0,0 +1,15 @@
|
||||
//! Shared storage layer for libchat.
|
||||
//!
|
||||
//! This crate provides a common storage abstraction that can be used by
|
||||
//! multiple crates in the libchat workspace (double-ratchets, conversations, etc.).
|
||||
//!
|
||||
//! Uses SQLCipher for encrypted SQLite storage.
|
||||
|
||||
mod errors;
|
||||
mod sqlite;
|
||||
|
||||
pub use errors::StorageError;
|
||||
pub use sqlite::{SqliteDb, StorageConfig};
|
||||
|
||||
// Re-export rusqlite types that domain crates will need
|
||||
pub use rusqlite::{Error as RusqliteError, Transaction, params};
|
||||
76
storage/src/sqlite.rs
Normal file
76
storage/src/sqlite.rs
Normal file
@ -0,0 +1,76 @@
|
||||
//! SQLite storage backend.
|
||||
|
||||
use rusqlite::Connection;
|
||||
use std::path::Path;
|
||||
|
||||
use crate::StorageError;
|
||||
|
||||
/// Configuration for SQLite storage.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum StorageConfig {
|
||||
/// In-memory database (for testing).
|
||||
InMemory,
|
||||
/// File-based SQLite database.
|
||||
File(String),
|
||||
/// SQLCipher encrypted database.
|
||||
Encrypted { path: String, key: String },
|
||||
}
|
||||
|
||||
/// SQLite database wrapper.
|
||||
///
|
||||
/// This provides the core database connection and can be shared
|
||||
/// across different domain-specific storage implementations.
|
||||
pub struct SqliteDb {
|
||||
conn: Connection,
|
||||
}
|
||||
|
||||
impl SqliteDb {
|
||||
/// Creates a new SQLite database with the given configuration.
|
||||
pub fn new(config: StorageConfig) -> Result<Self, StorageError> {
|
||||
let conn = match config {
|
||||
StorageConfig::InMemory => Connection::open_in_memory()?,
|
||||
StorageConfig::File(ref path) => Connection::open(path)?,
|
||||
StorageConfig::Encrypted { ref path, ref key } => {
|
||||
let conn = Connection::open(path)?;
|
||||
conn.pragma_update(None, "key", key)?;
|
||||
conn
|
||||
}
|
||||
};
|
||||
|
||||
// Enable foreign keys
|
||||
conn.execute_batch("PRAGMA foreign_keys = ON;")?;
|
||||
|
||||
Ok(Self { conn })
|
||||
}
|
||||
|
||||
/// Opens an existing database file.
|
||||
pub fn open<P: AsRef<Path>>(path: P) -> Result<Self, StorageError> {
|
||||
let conn = Connection::open(path)?;
|
||||
conn.execute_batch("PRAGMA foreign_keys = ON;")?;
|
||||
Ok(Self { conn })
|
||||
}
|
||||
|
||||
/// Creates an in-memory database (useful for testing).
|
||||
pub fn in_memory() -> Result<Self, StorageError> {
|
||||
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.
|
||||
pub fn connection(&self) -> &Connection {
|
||||
&self.conn
|
||||
}
|
||||
|
||||
/// Begins a transaction.
|
||||
pub fn transaction(&mut self) -> Result<rusqlite::Transaction<'_>, StorageError> {
|
||||
Ok(self.conn.transaction()?)
|
||||
}
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user