From d27b439c2de6e9a1fc9c76e758c91b81793bc2d9 Mon Sep 17 00:00:00 2001 From: kaichaosun Date: Thu, 5 Feb 2026 14:44:22 +0800 Subject: [PATCH] feat: integrate double ratchet --- conversations/src/chat.rs | 140 ++++++++++++++++---- conversations/src/common.rs | 64 +-------- conversations/src/dm/privatev1.rs | 77 ++++++++++- conversations/src/inbox/inbox.rs | 12 +- conversations/src/storage/db.rs | 205 ++++++++++++++++++++++++++++- conversations/src/storage/mod.rs | 2 +- conversations/src/storage/types.rs | 21 +++ conversations/src/utils.rs | 9 ++ double-ratchets/src/state.rs | 29 ++++ 9 files changed, 455 insertions(+), 104 deletions(-) diff --git a/conversations/src/chat.rs b/conversations/src/chat.rs index 5c3425d..1e6025d 100644 --- a/conversations/src/chat.rs +++ b/conversations/src/chat.rs @@ -3,10 +3,12 @@ //! This is the main entry point for the conversations API. It handles all //! storage operations internally - users don't need to interact with storage directly. +use std::collections::HashMap; use std::rc::Rc; use crate::{ - common::{Chat, ChatStore, HasChatId, InboundMessageHandler}, + common::{Chat, HasChatId, InboundMessageHandler}, + dm::privatev1::PrivateV1Convo, errors::ChatError, identity::Identity, inbox::{Inbox, Introduction}, @@ -62,7 +64,8 @@ pub enum ChatManagerError { /// ``` pub struct ChatManager { identity: Rc, - store: ChatStore, + /// In-memory cache of active chats. Chats are loaded from storage on demand. + chats: HashMap, inbox: Inbox, storage: ChatStorage, } @@ -94,7 +97,7 @@ impl ChatManager { Ok(Self { identity, - store: ChatStore::new(), + chats: HashMap::new(), inbox, storage, }) @@ -157,8 +160,13 @@ impl ChatManager { ); self.storage.save_chat(&chat_record)?; - // Store in memory - self.store.insert_chat(convo); + // Persist ratchet state + let (state, skipped_keys) = convo.to_storage(); + self.storage + .save_ratchet_state(&chat_id, &state, &skipped_keys)?; + + // Store in memory cache + self.chats.insert(chat_id.clone(), convo); Ok((chat_id, envelopes)) } @@ -171,27 +179,50 @@ impl ChatManager { chat_id: &str, content: &[u8], ) -> Result, ChatManagerError> { - // Try to get chat from memory first - let chat = match self.store.get_mut_chat(chat_id) { - Some(chat) => chat, - None => { - // Check if chat exists in storage but not loaded - if self.storage.chat_exists(chat_id)? { - return Err(ChatManagerError::ChatNotLoaded(chat_id.to_string())); - } else { - return Err(ChatManagerError::ChatNotFound(chat_id.to_string())); - } - } - }; + // Try to load chat from storage if not in memory + self.ensure_chat_loaded(chat_id)?; + + let chat = self + .chats + .get_mut(chat_id) + .ok_or_else(|| ChatManagerError::ChatNotFound(chat_id.to_string()))?; let payloads = chat.send_message(content)?; + // Persist updated ratchet state + let (state, skipped_keys) = chat.to_storage(); + self.storage + .save_ratchet_state(chat_id, &state, &skipped_keys)?; + + let remote_id = chat.remote_id(); Ok(payloads .into_iter() - .map(|p| p.to_envelope(chat.remote_id())) + .map(|p| p.to_envelope(remote_id.clone())) .collect()) } + /// Ensure a chat is loaded into memory. Loads from storage if needed. + fn ensure_chat_loaded(&mut self, chat_id: &str) -> Result<(), ChatManagerError> { + if self.chats.contains_key(chat_id) { + return Ok(()); + } + + // Try to load from storage + if let Some((state, skipped_keys)) = self.storage.load_ratchet_state(chat_id)? { + let convo = PrivateV1Convo::from_storage(chat_id.to_string(), state, skipped_keys); + self.chats.insert(chat_id.to_string(), convo); + Ok(()) + } else if self.storage.chat_exists(chat_id)? { + // Chat metadata exists but no ratchet state - this is a data inconsistency + Err(ChatManagerError::ChatNotFound(format!( + "{} (corrupted: missing ratchet state)", + chat_id + ))) + } else { + Err(ChatManagerError::ChatNotFound(chat_id.to_string())) + } + } + /// Handle an incoming payload from the network. /// /// This processes both inbox handshakes (to establish new chats) and @@ -205,8 +236,7 @@ impl ChatManager { Ok((chat, content_data)) => { let chat_id = chat.id().to_string(); - // Persist the new chat - // Note: We don't have full remote info here, using placeholder + // Persist the new chat metadata let chat_record = ChatRecord { chat_id: chat_id.clone(), chat_type: "private_v1".to_string(), @@ -216,8 +246,10 @@ impl ChatManager { }; self.storage.save_chat(&chat_record)?; - // Store chat in memory - self.store.insert_boxed_chat(chat); + // TODO: Persist ratchet state for incoming chats + // This requires modifying InboundMessageHandler to return PrivateV1Convo + // or adding downcast support. For now, new chats from inbox won't persist + // their ratchet state until next send_message call. // Return first content if any, otherwise empty if let Some(first) = content_data.into_iter().next() { @@ -241,13 +273,15 @@ impl ChatManager { } /// Get a reference to an active chat. - pub fn get_chat(&self, chat_id: &str) -> Option<&dyn Chat> { - self.store.get_chat(chat_id) + pub fn get_chat(&mut self, chat_id: &str) -> Option<&PrivateV1Convo> { + // Try to load from storage if not in memory + let _ = self.ensure_chat_loaded(chat_id); + self.chats.get(chat_id) } /// List all active chat IDs (in memory). pub fn list_chats(&self) -> Vec { - self.store.chat_ids().map(|id| id.to_string()).collect() + self.chats.keys().cloned().collect() } /// List all chat IDs from storage. @@ -257,7 +291,7 @@ impl ChatManager { /// Check if a chat exists (in memory or storage). pub fn chat_exists(&self, chat_id: &str) -> Result { - if self.store.get_chat(chat_id).is_some() { + if self.chats.contains_key(chat_id) { return Ok(true); } Ok(self.storage.chat_exists(chat_id)?) @@ -265,7 +299,7 @@ impl ChatManager { /// Delete a chat from both memory and storage. pub fn delete_chat(&mut self, chat_id: &str) -> Result<(), ChatManagerError> { - self.store.remove_chat(chat_id); + self.chats.remove(chat_id); self.storage.delete_chat(chat_id)?; Ok(()) } @@ -369,4 +403,56 @@ mod tests { assert!(!alice.chat_exists(&chat_id).unwrap()); assert!(alice.list_chats().is_empty()); } + + #[test] + fn test_ratchet_state_persistence() { + use tempfile::tempdir; + + // Create a temporary directory for the database + let dir = tempdir().unwrap(); + let db_path = dir.path().join("test.db"); + + let mut bob = ChatManager::in_memory().unwrap(); + let bob_intro = bob.create_intro_bundle().unwrap(); + + let chat_id; + + // Scope 1: Create chat and send messages + { + let mut alice = ChatManager::open(StorageConfig::File( + db_path.to_str().unwrap().to_string(), + )) + .unwrap(); + + let result = alice.start_private_chat(&bob_intro, "Message 1").unwrap(); + chat_id = result.0; + + // Send more messages - this advances the ratchet + alice.send_message(&chat_id, b"Message 2").unwrap(); + alice.send_message(&chat_id, b"Message 3").unwrap(); + + // Chat should be in memory + assert!(alice.chats.contains_key(&chat_id)); + } + // alice is dropped here, simulating app close + + // Scope 2: Reopen and verify chat is restored + { + let mut alice2 = ChatManager::open(StorageConfig::File( + db_path.to_str().unwrap().to_string(), + )) + .unwrap(); + + // Chat is in storage but not loaded yet + assert!(alice2.list_stored_chats().unwrap().contains(&chat_id)); + assert!(!alice2.chats.contains_key(&chat_id)); + + // Send another message - this will load the chat and advance ratchet + let result = alice2.send_message(&chat_id, b"Message 4"); + assert!(result.is_ok(), "Should be able to send after restore"); + + // Chat should now be in memory + assert!(alice2.chats.contains_key(&chat_id)); + } + } } diff --git a/conversations/src/common.rs b/conversations/src/common.rs index ada278e..e8492a8 100644 --- a/conversations/src/common.rs +++ b/conversations/src/common.rs @@ -1,12 +1,9 @@ -use std::collections::HashMap; use std::fmt::Debug; -use std::sync::Arc; pub use crate::errors::ChatError; use crate::types::{AddressedEncryptedPayload, ContentData}; pub type ChatId<'a> = &'a str; -pub type ChatIdOwned = Arc; pub trait HasChatId: Debug { fn id(&self) -> ChatId<'_>; @@ -21,66 +18,7 @@ pub trait InboundMessageHandler { pub trait Chat: HasChatId + Debug { fn send_message(&mut self, content: &[u8]) - -> Result, ChatError>; + -> Result, ChatError>; fn remote_id(&self) -> String; } - -pub struct ChatStore { - chats: HashMap, Box>, - handlers: HashMap, Box>, -} - -impl ChatStore { - pub fn new() -> Self { - Self { - chats: HashMap::new(), - handlers: HashMap::new(), - } - } - - pub fn insert_chat(&mut self, conversation: impl Chat + HasChatId + 'static) -> ChatIdOwned { - let key: ChatIdOwned = Arc::from(conversation.id()); - self.chats.insert(key.clone(), Box::new(conversation)); - key - } - - pub fn insert_boxed_chat(&mut self, conversation: Box) -> ChatIdOwned { - let key: ChatIdOwned = Arc::from(conversation.id()); - self.chats.insert(key.clone(), conversation); - key - } - - pub fn remove_chat(&mut self, id: &str) -> Option> { - self.chats.remove(id) - } - - pub fn register_handler( - &mut self, - handler: impl InboundMessageHandler + HasChatId + 'static, - ) -> ChatIdOwned { - let key: ChatIdOwned = Arc::from(handler.id()); - self.handlers.insert(key.clone(), Box::new(handler)); - key - } - - pub fn get_chat(&self, id: ChatId) -> Option<&(dyn Chat + '_)> { - self.chats.get(id).map(|c| c.as_ref()) - } - - pub fn get_mut_chat(&mut self, id: &str) -> Option<&mut (dyn Chat + '_)> { - Some(self.chats.get_mut(id)?.as_mut()) - } - - pub fn get_handler(&mut self, id: ChatId) -> Option<&mut (dyn InboundMessageHandler + '_)> { - Some(self.handlers.get_mut(id)?.as_mut()) - } - - pub fn chat_ids(&self) -> impl Iterator + '_ { - self.chats.keys().cloned() - } - - pub fn handler_ids(&self) -> impl Iterator + '_ { - self.handlers.keys().cloned() - } -} diff --git a/conversations/src/dm/privatev1.rs b/conversations/src/dm/privatev1.rs index 503a747..2a6e74b 100644 --- a/conversations/src/dm/privatev1.rs +++ b/conversations/src/dm/privatev1.rs @@ -12,32 +12,94 @@ use crate::{ common::{Chat, ChatId, HasChatId}, errors::{ChatError, EncryptionError}, proto, + storage::types::{RatchetStateRecord, SkippedKeyRecord}, types::AddressedEncryptedPayload, utils::timestamp_millis, }; pub struct PrivateV1Convo { + chat_id: String, dr_state: RatchetState, } impl PrivateV1Convo { - pub fn new_initiator(seed_key: SecretKey, remote: PublicKey) -> Self { + pub fn new_initiator(chat_id: String, seed_key: SecretKey, remote: PublicKey) -> Self { // TODO: Danger - Fix double-ratchets types to Accept SecretKey // perhaps update the DH to work with cryptocrate. // init_sender doesn't take ownership of the key so a reference can be used. let shared_secret: [u8; 32] = seed_key.as_bytes().to_vec().try_into().unwrap(); Self { + chat_id, dr_state: RatchetState::init_sender(shared_secret, remote), } } - pub fn new_responder(seed_key: SecretKey, dh_self: InstallationKeyPair) -> Self { + pub fn new_responder(chat_id: String, seed_key: SecretKey, dh_self: InstallationKeyPair) -> Self { Self { + chat_id, // TODO: Danger - Fix double-ratchets types to Accept SecretKey dr_state: RatchetState::init_receiver(seed_key.as_bytes().to_owned(), dh_self), } } + /// Restore a conversation from stored ratchet state. + pub fn from_storage( + chat_id: String, + state: RatchetStateRecord, + skipped_keys: Vec, + ) -> Self { + use std::collections::HashMap; + + let dh_self = InstallationKeyPair::from_secret_bytes(state.dh_self_secret); + let dh_remote = state.dh_remote.map(PublicKey::from); + + let skipped: HashMap<(PublicKey, u32), [u8; 32]> = skipped_keys + .into_iter() + .map(|sk| ((PublicKey::from(sk.public_key), sk.msg_num), sk.message_key)) + .collect(); + + let dr_state = RatchetState::from_parts( + state.root_key, + state.sending_chain, + state.receiving_chain, + dh_self, + dh_remote, + state.msg_send, + state.msg_recv, + state.prev_chain_len, + skipped, + ); + + Self { chat_id, dr_state } + } + + /// Get the current ratchet state for storage. + pub fn to_storage(&self) -> (RatchetStateRecord, Vec) { + let state = RatchetStateRecord { + root_key: self.dr_state.root_key, + sending_chain: self.dr_state.sending_chain, + receiving_chain: self.dr_state.receiving_chain, + dh_self_secret: *self.dr_state.dh_self.secret_bytes(), + dh_remote: self.dr_state.dh_remote.map(|pk| pk.to_bytes()), + msg_send: self.dr_state.msg_send, + msg_recv: self.dr_state.msg_recv, + prev_chain_len: self.dr_state.prev_chain_len, + }; + + let skipped_keys: Vec = self + .dr_state + .skipped_keys + .iter() + .map(|((pk, msg_num), key)| SkippedKeyRecord { + public_key: pk.to_bytes(), + msg_num: *msg_num, + message_key: *key, + }) + .collect(); + + (state, skipped_keys) + } + fn encrypt(&mut self, frame: PrivateV1Frame) -> EncryptedPayload { let encoded_bytes = frame.encode_to_vec(); let (cipher_text, header) = self.dr_state.encrypt_message(&encoded_bytes); @@ -93,8 +155,7 @@ impl PrivateV1Convo { impl HasChatId for PrivateV1Convo { fn id(&self) -> ChatId<'_> { - // TODO: implementation - "private_v1_convo_id" + &self.chat_id } } @@ -147,11 +208,15 @@ mod tests { let seed_key = saro.diffie_hellman(&pub_raya); let send_content_bytes = vec![0, 2, 4, 6, 8]; - let mut sr_convo = - PrivateV1Convo::new_initiator(SecretKey::from(seed_key.to_bytes()), pub_raya); + let mut sr_convo = PrivateV1Convo::new_initiator( + "test_chat".to_string(), + SecretKey::from(seed_key.to_bytes()), + pub_raya, + ); let installation_key_pair = InstallationKeyPair::from(raya); let mut rs_convo = PrivateV1Convo::new_responder( + "test_chat".to_string(), SecretKey::from(seed_key.to_bytes()), installation_key_pair, ); diff --git a/conversations/src/inbox/inbox.rs b/conversations/src/inbox/inbox.rs index bb1cbb5..588c171 100644 --- a/conversations/src/inbox/inbox.rs +++ b/conversations/src/inbox/inbox.rs @@ -15,6 +15,7 @@ use crate::identity::{PublicKey, StaticSecret}; use crate::inbox::handshake::InboxHandshake; use crate::proto::{self, CopyBytes}; use crate::types::{AddressedEncryptedPayload, ContentData}; +use crate::utils::generate_chat_id; use super::Introduction; @@ -107,7 +108,10 @@ impl Inbox { let (seed_key, ephemeral_pub) = InboxHandshake::perform_as_initiator(&self.ident.secret(), &pkb, &mut rng); - let mut convo = PrivateV1Convo::new_initiator(seed_key, remote_bundle.ephemeral_key); + // Generate unique chat ID + let chat_id = generate_chat_id(); + let mut convo = + PrivateV1Convo::new_initiator(chat_id, seed_key, remote_bundle.ephemeral_key); let mut payloads = convo.send_message(initial_message.as_bytes())?; @@ -245,7 +249,11 @@ impl InboundMessageHandler for Inbox { match frame.frame_type.unwrap() { proto::inbox_v1_frame::FrameType::InvitePrivateV1(_invite_private_v1) => { - let convo = PrivateV1Convo::new_responder(seed_key, ephemeral_key.clone().into()); + // Generate unique chat ID for the responder + let chat_id = generate_chat_id(); + let installation_keypair = + double_ratchets::InstallationKeyPair::from(ephemeral_key.clone()); + let convo = PrivateV1Convo::new_responder(chat_id, seed_key, installation_keypair); // TODO: Update PrivateV1 Constructor with DR, initial_message Ok((Box::new(convo), vec![])) diff --git a/conversations/src/storage/db.rs b/conversations/src/storage/db.rs index e8b151f..4993e70 100644 --- a/conversations/src/storage/db.rs +++ b/conversations/src/storage/db.rs @@ -5,7 +5,7 @@ use std::collections::HashMap; use storage::{RusqliteError, SqliteDb, StorageConfig, StorageError, params}; use x25519_dalek::StaticSecret; -use super::types::{ChatRecord, IdentityRecord}; +use super::types::{ChatRecord, IdentityRecord, RatchetStateRecord, SkippedKeyRecord}; use crate::identity::Identity; /// Schema for chat storage tables. @@ -33,6 +33,33 @@ const CHAT_SCHEMA: &str = " ); CREATE INDEX IF NOT EXISTS idx_chats_type ON chats(chat_type); + + -- Ratchet state for each conversation + 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 + ); + + -- Skipped message keys (for out-of-order messages) + 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); "; /// Chat-specific storage operations. @@ -257,13 +284,181 @@ impl ChatStorage { Ok(exists) } - /// Deletes a chat record. + /// Deletes a chat record and its ratchet state. pub fn delete_chat(&mut self, chat_id: &str) -> Result<(), StorageError> { - self.db - .connection() - .execute("DELETE FROM chats WHERE chat_id = ?1", params![chat_id])?; + let tx = self.db.transaction()?; + // Delete skipped keys first (foreign key constraint) + tx.execute( + "DELETE FROM skipped_keys WHERE conversation_id = ?1", + params![chat_id], + )?; + tx.execute( + "DELETE FROM ratchet_state WHERE conversation_id = ?1", + params![chat_id], + )?; + tx.execute("DELETE FROM chats WHERE chat_id = ?1", params![chat_id])?; + tx.commit()?; Ok(()) } + + // ==================== Ratchet State Operations ==================== + + /// Saves the ratchet state for a conversation. + pub fn save_ratchet_state( + &mut self, + conversation_id: &str, + state: &RatchetStateRecord, + skipped_keys: &[SkippedKeyRecord], + ) -> Result<(), StorageError> { + let tx = self.db.transaction()?; + + // 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, + state.root_key.as_slice(), + state.sending_chain.as_ref().map(|c| c.as_slice()), + state.receiving_chain.as_ref().map(|c| c.as_slice()), + state.dh_self_secret.as_slice(), + state.dh_remote.as_ref().map(|c| c.as_slice()), + state.msg_send, + state.msg_recv, + state.prev_chain_len, + ], + )?; + + // Sync skipped keys: delete old ones and insert new + tx.execute( + "DELETE FROM skipped_keys WHERE conversation_id = ?1", + params![conversation_id], + )?; + + for sk in skipped_keys { + 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(), + ], + )?; + } + + tx.commit()?; + Ok(()) + } + + /// Loads the ratchet state for a conversation. + pub fn load_ratchet_state( + &self, + conversation_id: &str, + ) -> Result)>, StorageError> { + // Load main state + let state = self.load_ratchet_state_data(conversation_id)?; + let state = match state { + Some(s) => s, + None => return Ok(None), + }; + + // Load skipped keys + let skipped_keys = self.load_skipped_keys(conversation_id)?; + + Ok(Some((state, skipped_keys))) + } + + fn load_ratchet_state_data( + &self, + conversation_id: &str, + ) -> Result, 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 + ", + )?; + + let result = stmt.query_row(params![conversation_id], |row| { + Ok(RatchetStateRecord { + root_key: blob_to_array(row.get::<_, Vec>(0)?), + sending_chain: row.get::<_, Option>>(1)?.map(blob_to_array), + receiving_chain: row.get::<_, Option>>(2)?.map(blob_to_array), + dh_self_secret: blob_to_array(row.get::<_, Vec>(3)?), + dh_remote: row.get::<_, Option>>(4)?.map(blob_to_array), + msg_send: row.get(5)?, + msg_recv: row.get(6)?, + prev_chain_len: row.get(7)?, + }) + }); + + match result { + Ok(record) => Ok(Some(record)), + Err(RusqliteError::QueryReturnedNoRows) => Ok(None), + Err(e) => Err(StorageError::Database(e.to_string())), + } + } + + fn load_skipped_keys( + &self, + conversation_id: &str, + ) -> Result, 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(SkippedKeyRecord { + public_key: blob_to_array(row.get::<_, Vec>(0)?), + msg_num: row.get(1)?, + message_key: blob_to_array(row.get::<_, Vec>(2)?), + }) + })?; + + rows.collect::, _>>() + .map_err(|e| StorageError::Database(e.to_string())) + } + + /// Checks if a ratchet state exists for a conversation. + pub fn ratchet_state_exists(&self, conversation_id: &str) -> Result { + 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) + } +} + +/// Helper to convert a Vec to a fixed-size array. +fn blob_to_array(blob: Vec) -> [u8; 32] { + let mut arr = [0u8; 32]; + arr.copy_from_slice(&blob); + arr } #[cfg(test)] diff --git a/conversations/src/storage/mod.rs b/conversations/src/storage/mod.rs index 93bb317..5153dbf 100644 --- a/conversations/src/storage/mod.rs +++ b/conversations/src/storage/mod.rs @@ -7,7 +7,7 @@ //! handles all storage operations automatically. mod db; -mod types; +pub(crate) mod types; pub(crate) use db::ChatStorage; pub(crate) use storage::StorageError; diff --git a/conversations/src/storage/types.rs b/conversations/src/storage/types.rs index 6c39834..4c3a5bb 100644 --- a/conversations/src/storage/types.rs +++ b/conversations/src/storage/types.rs @@ -57,3 +57,24 @@ impl ChatRecord { } } } + +/// Raw ratchet state data for SQLite storage. +#[derive(Debug, Clone)] +pub struct RatchetStateRecord { + pub root_key: [u8; 32], + pub sending_chain: Option<[u8; 32]>, + pub receiving_chain: Option<[u8; 32]>, + pub dh_self_secret: [u8; 32], + pub dh_remote: Option<[u8; 32]>, + pub msg_send: u32, + pub msg_recv: u32, + pub prev_chain_len: u32, +} + +/// Skipped key record for out-of-order message handling. +#[derive(Debug, Clone)] +pub struct SkippedKeyRecord { + pub public_key: [u8; 32], + pub msg_num: u32, + pub message_key: [u8; 32], +} diff --git a/conversations/src/utils.rs b/conversations/src/utils.rs index 306e898..4d1fad0 100644 --- a/conversations/src/utils.rs +++ b/conversations/src/utils.rs @@ -1,8 +1,17 @@ use std::time::{SystemTime, UNIX_EPOCH}; +use rand_core::OsRng; +use x25519_dalek::StaticSecret; + pub fn timestamp_millis() -> i64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_millis() as i64 } + +/// Generate a unique chat ID using random bytes. +pub fn generate_chat_id() -> String { + let secret = StaticSecret::random_from_rng(OsRng); + hex::encode(&secret.as_bytes()[..16]) +} diff --git a/double-ratchets/src/state.rs b/double-ratchets/src/state.rs index 48ad359..b7fe74c 100644 --- a/double-ratchets/src/state.rs +++ b/double-ratchets/src/state.rs @@ -291,6 +291,35 @@ impl RatchetState { } } + /// Reconstructs a RatchetState from its component parts. + /// + /// This is used for restoring state from storage. + #[allow(clippy::too_many_arguments)] + pub fn from_parts( + root_key: RootKey, + sending_chain: Option, + receiving_chain: Option, + dh_self: InstallationKeyPair, + dh_remote: Option, + msg_send: u32, + msg_recv: u32, + prev_chain_len: u32, + skipped_keys: HashMap<(PublicKey, u32), MessageKey>, + ) -> Self { + Self { + root_key, + sending_chain, + receiving_chain, + dh_self, + dh_remote, + msg_send, + msg_recv, + prev_chain_len, + skipped_keys, + _domain: PhantomData, + } + } + /// Performs a receiving-side DH ratchet when a new remote DH public key is observed. /// /// # Arguments