From 5d87b1d19ac42bb98471d71cb7cb120b0ce2b591 Mon Sep 17 00:00:00 2001 From: kaichaosun Date: Thu, 5 Mar 2026 20:45:18 +0800 Subject: [PATCH] feat: db storage for inbox ephermeral keys --- conversations/src/context.rs | 22 +++++- conversations/src/inbox/handler.rs | 23 ++++-- conversations/src/storage/db.rs | 70 ++++++++++++++++++- conversations/src/storage/migrations.rs | 14 ++-- .../storage/migrations/002_ephemeral_keys.sql | 7 ++ conversations/src/storage/types.rs | 13 +++- 6 files changed, 133 insertions(+), 16 deletions(-) create mode 100644 conversations/src/storage/migrations/002_ephemeral_keys.sql diff --git a/conversations/src/context.rs b/conversations/src/context.rs index bec37a5..9a1df9b 100644 --- a/conversations/src/context.rs +++ b/conversations/src/context.rs @@ -1,9 +1,11 @@ +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, @@ -44,7 +46,17 @@ impl Context { }; let identity = Rc::new(identity); - let inbox = Inbox::new(Rc::clone(&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); + } Ok(Self { _identity: identity, @@ -126,7 +138,8 @@ impl Context { &mut self, enc_payload: EncryptedPayload, ) -> Result, ChatError> { - let (convo, content) = self.inbox.handle_frame(enc_payload)?; + let (convo, content, consumed_key_hex) = self.inbox.handle_frame(enc_payload)?; + self.storage.remove_ephemeral_key(&consumed_key_hex)?; self.add_convo(convo); Ok(content) } @@ -145,7 +158,10 @@ impl Context { } pub fn create_intro_bundle(&mut self) -> Result, ChatError> { - Ok(self.inbox.create_intro_bundle().into()) + let (intro, public_key_hex, private_key) = self.inbox.create_intro_bundle(); + self.storage + .save_ephemeral_key(&public_key_hex, &private_key)?; + Ok(intro.into()) } fn add_convo(&mut self, convo: Box) -> ConversationIdOwned { diff --git a/conversations/src/inbox/handler.rs b/conversations/src/inbox/handler.rs index 278ae16..6535c90 100644 --- a/conversations/src/inbox/handler.rs +++ b/conversations/src/inbox/handler.rs @@ -51,14 +51,23 @@ impl Inbox { } } - pub fn create_intro_bundle(&mut self) -> Introduction { + /// 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) { let ephemeral = PrivateKey::random(); let ephemeral_key: PublicKey = (&ephemeral).into(); + let public_key_hex = hex::encode(ephemeral_key.as_bytes()); self.ephemeral_keys - .insert(hex::encode(ephemeral_key.as_bytes()), ephemeral); + .insert(public_key_hex.clone(), ephemeral.clone()); - Introduction::new(self.ident.secret(), ephemeral_key, OsRng) + let intro = Introduction::new(self.ident.secret(), ephemeral_key, OsRng); + (intro, public_key_hex, ephemeral) } pub fn invite_to_private_convo( @@ -114,10 +123,12 @@ 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). pub fn handle_frame( &mut self, enc_payload: EncryptedPayload, - ) -> Result<(Box, Option), ChatError> { + ) -> Result<(Box, Option, String), ChatError> { let handshake = Self::extract_payload(enc_payload)?; let header = handshake @@ -148,7 +159,7 @@ impl Inbox { None => return Err(ChatError::Protocol("expected contentData".into())), }; - Ok((Box::new(convo), Some(content))) + Ok((Box::new(convo), Some(content), key_index)) } } } @@ -244,7 +255,7 @@ mod tests { let raya_ident = Identity::new("raya"); let mut raya_inbox = Inbox::new(raya_ident.into()); - let bundle = raya_inbox.create_intro_bundle(); + 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(); diff --git a/conversations/src/storage/db.rs b/conversations/src/storage/db.rs index c855416..f8fb133 100644 --- a/conversations/src/storage/db.rs +++ b/conversations/src/storage/db.rs @@ -4,7 +4,8 @@ use storage::{RusqliteError, SqliteDb, StorageConfig, StorageError, params}; use zeroize::Zeroize; use super::migrations; -use super::types::IdentityRecord; +use super::types::{EphemeralKeyRecord, IdentityRecord}; +use crate::crypto::PrivateKey; use crate::identity::Identity; /// Chat-specific storage operations. @@ -47,6 +48,73 @@ impl ChatStorage { Ok(()) } + // ==================== Ephemeral Key Operations ==================== + + /// Saves an ephemeral key pair to storage. + pub fn save_ephemeral_key( + &mut self, + public_key_hex: &str, + private_key: &PrivateKey, + ) -> Result<(), StorageError> { + let mut secret_bytes = private_key.DANGER_to_bytes(); + let result = self.db.connection().execute( + "INSERT OR REPLACE INTO ephemeral_keys (public_key_hex, secret_key) VALUES (?1, ?2)", + params![public_key_hex, secret_bytes.as_slice()], + ); + secret_bytes.zeroize(); + result?; + Ok(()) + } + + /// Loads all ephemeral keys from storage. + pub fn load_ephemeral_keys( + &self, + ) -> Result, StorageError> { + let mut stmt = self + .db + .connection() + .prepare("SELECT public_key_hex, secret_key FROM ephemeral_keys")?; + + 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 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, + }); + } + Ok(result) + } + + /// Removes an ephemeral key from storage. + pub fn remove_ephemeral_key(&mut self, public_key_hex: &str) -> Result<(), StorageError> { + self.db.connection().execute( + "DELETE FROM ephemeral_keys WHERE public_key_hex = ?1", + params![public_key_hex], + )?; + Ok(()) + } + + // ==================== Identity Operations (continued) ==================== + /// Loads the identity if it exists. /// /// Note: Secret key bytes are zeroized after being copied into IdentityRecord, diff --git a/conversations/src/storage/migrations.rs b/conversations/src/storage/migrations.rs index 014bb96..5de4737 100644 --- a/conversations/src/storage/migrations.rs +++ b/conversations/src/storage/migrations.rs @@ -7,10 +7,16 @@ use storage::{Connection, StorageError}; /// Embeds and returns all migration SQL files in order. pub fn get_migrations() -> Vec<(&'static str, &'static str)> { - vec![( - "001_initial_schema", - include_str!("migrations/001_initial_schema.sql"), - )] + vec![ + ( + "001_initial_schema", + include_str!("migrations/001_initial_schema.sql"), + ), + ( + "002_ephemeral_keys", + include_str!("migrations/002_ephemeral_keys.sql"), + ), + ] } /// Applies all migrations to the database. diff --git a/conversations/src/storage/migrations/002_ephemeral_keys.sql b/conversations/src/storage/migrations/002_ephemeral_keys.sql new file mode 100644 index 0000000..14ecbb0 --- /dev/null +++ b/conversations/src/storage/migrations/002_ephemeral_keys.sql @@ -0,0 +1,7 @@ +-- Ephemeral keys for inbox handshakes +-- Migration: 002_ephemeral_keys + +CREATE TABLE IF NOT EXISTS ephemeral_keys ( + public_key_hex TEXT PRIMARY KEY, + secret_key BLOB NOT NULL +); diff --git a/conversations/src/storage/types.rs b/conversations/src/storage/types.rs index 8767324..4a1ef72 100644 --- a/conversations/src/storage/types.rs +++ b/conversations/src/storage/types.rs @@ -17,12 +17,21 @@ pub struct IdentityRecord { impl From for Identity { fn from(record: IdentityRecord) -> Self { - let name = record.name.clone(); let secret = PrivateKey::from(record.secret_key); - Identity::from_secret(name, secret) + Identity::from_secret(record.name.clone(), secret) } } +/// 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::*;