diff --git a/conversations/src/context.rs b/conversations/src/context.rs index 9a1df9b..af63bb2 100644 --- a/conversations/src/context.rs +++ b/conversations/src/context.rs @@ -1,11 +1,9 @@ -use std::collections::HashMap; use std::rc::Rc; use storage::StorageConfig; use crate::{ conversation::{ConversationId, ConversationStore, Convo, Id}, - crypto::PrivateKey, errors::ChatError, identity::Identity, inbox::Inbox, @@ -23,7 +21,6 @@ pub struct Context { _identity: Rc, store: ConversationStore, inbox: Inbox, - #[allow(dead_code)] // Will be used for conversation persistence storage: ChatStorage, } @@ -46,17 +43,7 @@ impl Context { }; let identity = Rc::new(identity); - let mut inbox = Inbox::new(Rc::clone(&identity)); - - // Restore ephemeral keys from storage - let stored_keys = storage.load_ephemeral_keys()?; - if !stored_keys.is_empty() { - let keys: HashMap = stored_keys - .into_iter() - .map(|record| (record.public_key_hex.clone(), PrivateKey::from(record.secret_key))) - .collect(); - inbox.restore_ephemeral_keys(keys); - } + let inbox = Inbox::new(Rc::clone(&identity)); Ok(Self { _identity: identity, @@ -138,8 +125,18 @@ impl Context { &mut self, enc_payload: EncryptedPayload, ) -> Result, ChatError> { - let (convo, content, consumed_key_hex) = self.inbox.handle_frame(enc_payload)?; - self.storage.remove_ephemeral_key(&consumed_key_hex)?; + // Look up the ephemeral key from storage + let key_hex = Inbox::extract_ephemeral_key_hex(&enc_payload)?; + let ephemeral_key = self + .storage + .load_ephemeral_key(&key_hex)? + .ok_or(ChatError::UnknownEphemeralKey())?; + + let (convo, content) = self.inbox.handle_frame(&ephemeral_key, enc_payload)?; + + // Remove consumed ephemeral key from storage + self.storage.remove_ephemeral_key(&key_hex)?; + self.add_convo(convo); Ok(content) } @@ -268,4 +265,37 @@ mod tests { assert_eq!(pubkey1, pubkey2, "public key should persist"); assert_eq!(name1, name2, "name should persist"); } + + #[test] + fn ephemeral_key_persistence() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir + .path() + .join("test_ephemeral.db") + .to_string_lossy() + .to_string(); + let config = StorageConfig::File(db_path); + + // Create context and generate an intro bundle (creates ephemeral key) + let mut ctx1 = Context::open("alice", config.clone()).unwrap(); + let bundle1 = ctx1.create_intro_bundle().unwrap(); + + // Drop and reopen - ephemeral keys should be restored from db + drop(ctx1); + let mut ctx2 = Context::open("alice", config.clone()).unwrap(); + + // Use the intro bundle from before restart to start a conversation + let intro = Introduction::try_from(bundle1.as_slice()).unwrap(); + let mut bob = Context::new_with_name("bob"); + let (_, payloads) = bob.create_private_convo(&intro, b"hello after restart"); + + // Alice (ctx2) should be able to handle the payload using the persisted ephemeral key + let payload = payloads.first().unwrap(); + let content = ctx2 + .handle_payload(&payload.data) + .expect("should handle payload with persisted ephemeral key") + .expect("should have content"); + assert_eq!(content.data, b"hello after restart"); + assert!(content.is_new_convo); + } } diff --git a/conversations/src/inbox/handler.rs b/conversations/src/inbox/handler.rs index 6535c90..b733097 100644 --- a/conversations/src/inbox/handler.rs +++ b/conversations/src/inbox/handler.rs @@ -3,7 +3,6 @@ use chat_proto::logoschat::encryption::EncryptedPayload; use prost::Message; use prost::bytes::Bytes; use rand_core::OsRng; -use std::collections::HashMap; use std::rc::Rc; use crypto::{PrekeyBundle, SymmetricKey32}; @@ -25,7 +24,6 @@ fn delivery_address_for_installation(_: PublicKey) -> String { pub struct Inbox { ident: Rc, local_convo_id: String, - ephemeral_keys: HashMap, } impl std::fmt::Debug for Inbox { @@ -33,10 +31,6 @@ impl std::fmt::Debug for Inbox { f.debug_struct("Inbox") .field("ident", &self.ident) .field("convo_id", &self.local_convo_id) - .field( - "ephemeral_keys", - &format!("<{} keys>", self.ephemeral_keys.len()), - ) .finish() } } @@ -47,24 +41,16 @@ impl Inbox { Self { ident, local_convo_id, - ephemeral_keys: HashMap::::new(), } } - /// Restores ephemeral keys from storage into the in-memory map. - pub fn restore_ephemeral_keys(&mut self, keys: HashMap) { - self.ephemeral_keys = keys; - } - - /// Creates an intro bundle and returns the (public_key_hex, private_key) pair - /// so the caller can persist it. - pub fn create_intro_bundle(&mut self) -> (Introduction, String, PrivateKey) { + /// Creates an intro bundle and returns the Introduction along with the + /// generated ephemeral key pair (public_key_hex, private_key) for the caller to persist. + pub fn create_intro_bundle(&self) -> (Introduction, String, PrivateKey) { let ephemeral = PrivateKey::random(); let ephemeral_key: PublicKey = (&ephemeral).into(); let public_key_hex = hex::encode(ephemeral_key.as_bytes()); - self.ephemeral_keys - .insert(public_key_hex.clone(), ephemeral.clone()); let intro = Introduction::new(self.ident.secret(), ephemeral_key, OsRng); (intro, public_key_hex, ephemeral) @@ -123,22 +109,19 @@ impl Inbox { Ok((convo, payloads)) } - /// Handles an incoming inbox frame. Returns the created conversation, - /// optional content data, and the consumed ephemeral key hex (for storage cleanup). + /// Handles an incoming inbox frame. The caller must provide the ephemeral private key + /// looked up from storage. Returns the created conversation and optional content data. pub fn handle_frame( - &mut self, + &self, + ephemeral_key: &PrivateKey, enc_payload: EncryptedPayload, - ) -> Result<(Box, Option, String), ChatError> { + ) -> Result<(Box, Option), ChatError> { let handshake = Self::extract_payload(enc_payload)?; let header = handshake .header .ok_or(ChatError::UnexpectedPayload("InboxV1Header".into()))?; - // Get Ephemeral key used by the initator - let key_index = hex::encode(header.responder_ephemeral.as_ref()); - let ephemeral_key = self.lookup_ephemeral_key(&key_index)?; - // Perform handshake and decrypt frame let (seed_key, frame) = self.perform_handshake(ephemeral_key, header, handshake.payload)?; @@ -159,11 +142,29 @@ impl Inbox { None => return Err(ChatError::Protocol("expected contentData".into())), }; - Ok((Box::new(convo), Some(content), key_index)) + Ok((Box::new(convo), Some(content))) } } } + /// Extracts the ephemeral key hex from an incoming encrypted payload + /// so the caller can look it up from storage before calling handle_frame. + pub fn extract_ephemeral_key_hex( + enc_payload: &EncryptedPayload, + ) -> Result { + let Some(proto::Encryption::InboxHandshake(ref handshake)) = enc_payload.encryption else { + let got = format!("{:?}", enc_payload.encryption); + return Err(ChatError::ProtocolExpectation("inboxhandshake", got)); + }; + + let header = handshake + .header + .as_ref() + .ok_or(ChatError::UnexpectedPayload("InboxV1Header".into()))?; + + Ok(hex::encode(header.responder_ephemeral.as_ref())) + } + fn wrap_in_invite(payload: proto::EncryptedPayload) -> proto::InboxV1Frame { let invite = proto::InvitePrivateV1 { discriminator: "default".into(), @@ -225,12 +226,6 @@ impl Inbox { Ok(frame) } - fn lookup_ephemeral_key(&self, key: &str) -> Result<&PrivateKey, ChatError> { - self.ephemeral_keys - .get(key) - .ok_or(ChatError::UnknownEphemeralKey()) - } - pub fn inbox_identifier_for_key(pubkey: PublicKey) -> String { // TODO: Implement ID according to spec hex::encode(Blake2b512::digest(pubkey)) @@ -246,24 +241,34 @@ impl Id for Inbox { #[cfg(test)] mod tests { use super::*; + use crate::storage::ChatStorage; + use storage::StorageConfig; #[test] fn test_invite_privatev1_roundtrip() { + let mut storage = ChatStorage::new(StorageConfig::InMemory).unwrap(); + let saro_ident = Identity::new("saro"); let saro_inbox = Inbox::new(saro_ident.into()); let raya_ident = Identity::new("raya"); - let mut raya_inbox = Inbox::new(raya_ident.into()); + let raya_inbox = Inbox::new(raya_ident.into()); + + let (bundle, key_hex, private_key) = raya_inbox.create_intro_bundle(); + storage.save_ephemeral_key(&key_hex, &private_key).unwrap(); - let (bundle, _key_hex, _private_key) = raya_inbox.create_intro_bundle(); let (_, mut payloads) = saro_inbox .invite_to_private_convo(&bundle, "hello".as_bytes()) .unwrap(); let payload = payloads.remove(0); + // Look up ephemeral key from storage + let key_hex = Inbox::extract_ephemeral_key_hex(&payload.data).unwrap(); + let ephemeral_key = storage.load_ephemeral_key(&key_hex).unwrap().unwrap(); + // Test handle_frame with valid payload - let result = raya_inbox.handle_frame(payload.data); + let result = raya_inbox.handle_frame(&ephemeral_key, payload.data); assert!( result.is_ok(), diff --git a/conversations/src/storage/db.rs b/conversations/src/storage/db.rs index f8fb133..8eaefe9 100644 --- a/conversations/src/storage/db.rs +++ b/conversations/src/storage/db.rs @@ -4,7 +4,7 @@ use storage::{RusqliteError, SqliteDb, StorageConfig, StorageError, params}; use zeroize::Zeroize; use super::migrations; -use super::types::{EphemeralKeyRecord, IdentityRecord}; +use super::types::IdentityRecord; use crate::crypto::PrivateKey; use crate::identity::Identity; @@ -66,42 +66,39 @@ impl ChatStorage { Ok(()) } - /// Loads all ephemeral keys from storage. - pub fn load_ephemeral_keys( + /// Loads a single ephemeral key by its public key hex. + pub fn load_ephemeral_key( &self, - ) -> Result, StorageError> { + public_key_hex: &str, + ) -> Result, StorageError> { let mut stmt = self .db .connection() - .prepare("SELECT public_key_hex, secret_key FROM ephemeral_keys")?; + .prepare("SELECT secret_key FROM ephemeral_keys WHERE public_key_hex = ?1")?; - let records = stmt - .query_map([], |row| { - let public_key_hex: String = row.get(0)?; - let secret_key: Vec = row.get(1)?; - Ok((public_key_hex, secret_key)) - })? - .collect::, _>>()?; + let result = stmt.query_row(params![public_key_hex], |row| { + let secret_key: Vec = row.get(0)?; + Ok(secret_key) + }); - let mut result = Vec::with_capacity(records.len()); - for (public_key_hex, mut secret_key_vec) in records { - let bytes: Result<[u8; 32], _> = secret_key_vec.as_slice().try_into(); - let bytes = match bytes { - Ok(b) => b, - Err(_) => { - secret_key_vec.zeroize(); - return Err(StorageError::InvalidData( - "Invalid ephemeral secret key length".into(), - )); - } - }; - secret_key_vec.zeroize(); - result.push(EphemeralKeyRecord { - public_key_hex, - secret_key: bytes, - }); + match result { + Ok(mut secret_key_vec) => { + let bytes: Result<[u8; 32], _> = secret_key_vec.as_slice().try_into(); + let bytes = match bytes { + Ok(b) => b, + Err(_) => { + secret_key_vec.zeroize(); + return Err(StorageError::InvalidData( + "Invalid ephemeral secret key length".into(), + )); + } + }; + secret_key_vec.zeroize(); + Ok(Some(PrivateKey::from(bytes))) + } + Err(RusqliteError::QueryReturnedNoRows) => Ok(None), + Err(e) => Err(e.into()), } - Ok(result) } /// Removes an ephemeral key from storage. @@ -176,4 +173,25 @@ mod tests { let loaded = storage.load_identity().unwrap().unwrap(); assert_eq!(loaded.public_key(), pubkey); } + + #[test] + fn test_ephemeral_key_roundtrip() { + let mut storage = ChatStorage::new(StorageConfig::InMemory).unwrap(); + + let key1 = PrivateKey::random(); + let pub1: crate::crypto::PublicKey = (&key1).into(); + let hex1 = hex::encode(pub1.as_bytes()); + + // Initially not found + assert!(storage.load_ephemeral_key(&hex1).unwrap().is_none()); + + // Save and load + storage.save_ephemeral_key(&hex1, &key1).unwrap(); + let loaded = storage.load_ephemeral_key(&hex1).unwrap().unwrap(); + assert_eq!(loaded.DANGER_to_bytes(), key1.DANGER_to_bytes()); + + // Remove and verify gone + storage.remove_ephemeral_key(&hex1).unwrap(); + assert!(storage.load_ephemeral_key(&hex1).unwrap().is_none()); + } } diff --git a/conversations/src/storage/types.rs b/conversations/src/storage/types.rs index 4a1ef72..c34f9be 100644 --- a/conversations/src/storage/types.rs +++ b/conversations/src/storage/types.rs @@ -22,16 +22,6 @@ impl From for Identity { } } -/// Record for storing an ephemeral key pair. -/// Implements ZeroizeOnDrop to securely clear secret key from memory. -#[derive(Debug, Zeroize, ZeroizeOnDrop)] -pub struct EphemeralKeyRecord { - /// Hex-encoded public key (used as lookup key). - pub public_key_hex: String, - /// The secret key bytes (32 bytes). - pub secret_key: [u8; 32], -} - #[cfg(test)] mod tests { use super::*;