feat: db storage for inbox ephermeral keys

This commit is contained in:
kaichaosun 2026-03-05 20:45:18 +08:00
parent 10a403e6fa
commit 5d87b1d19a
No known key found for this signature in database
GPG Key ID: 223E0F992F4F03BF
6 changed files with 133 additions and 16 deletions

View File

@ -1,9 +1,11 @@
use std::collections::HashMap;
use std::rc::Rc; use std::rc::Rc;
use storage::StorageConfig; use storage::StorageConfig;
use crate::{ use crate::{
conversation::{ConversationId, ConversationStore, Convo, Id}, conversation::{ConversationId, ConversationStore, Convo, Id},
crypto::PrivateKey,
errors::ChatError, errors::ChatError,
identity::Identity, identity::Identity,
inbox::Inbox, inbox::Inbox,
@ -44,7 +46,17 @@ impl Context {
}; };
let identity = Rc::new(identity); 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<String, PrivateKey> = stored_keys
.into_iter()
.map(|record| (record.public_key_hex.clone(), PrivateKey::from(record.secret_key)))
.collect();
inbox.restore_ephemeral_keys(keys);
}
Ok(Self { Ok(Self {
_identity: identity, _identity: identity,
@ -126,7 +138,8 @@ impl Context {
&mut self, &mut self,
enc_payload: EncryptedPayload, enc_payload: EncryptedPayload,
) -> Result<Option<ContentData>, ChatError> { ) -> Result<Option<ContentData>, 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); self.add_convo(convo);
Ok(content) Ok(content)
} }
@ -145,7 +158,10 @@ impl Context {
} }
pub fn create_intro_bundle(&mut self) -> Result<Vec<u8>, ChatError> { pub fn create_intro_bundle(&mut self) -> Result<Vec<u8>, 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<dyn Convo>) -> ConversationIdOwned { fn add_convo(&mut self, convo: Box<dyn Convo>) -> ConversationIdOwned {

View File

@ -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<String, PrivateKey>) {
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 = PrivateKey::random();
let ephemeral_key: PublicKey = (&ephemeral).into(); let ephemeral_key: PublicKey = (&ephemeral).into();
let public_key_hex = hex::encode(ephemeral_key.as_bytes());
self.ephemeral_keys 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( pub fn invite_to_private_convo(
@ -114,10 +123,12 @@ impl Inbox {
Ok((convo, payloads)) 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( pub fn handle_frame(
&mut self, &mut self,
enc_payload: EncryptedPayload, enc_payload: EncryptedPayload,
) -> Result<(Box<dyn Convo>, Option<ContentData>), ChatError> { ) -> Result<(Box<dyn Convo>, Option<ContentData>, String), ChatError> {
let handshake = Self::extract_payload(enc_payload)?; let handshake = Self::extract_payload(enc_payload)?;
let header = handshake let header = handshake
@ -148,7 +159,7 @@ impl Inbox {
None => return Err(ChatError::Protocol("expected contentData".into())), 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 raya_ident = Identity::new("raya");
let mut raya_inbox = Inbox::new(raya_ident.into()); 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 let (_, mut payloads) = saro_inbox
.invite_to_private_convo(&bundle, "hello".as_bytes()) .invite_to_private_convo(&bundle, "hello".as_bytes())
.unwrap(); .unwrap();

View File

@ -4,7 +4,8 @@ use storage::{RusqliteError, SqliteDb, StorageConfig, StorageError, params};
use zeroize::Zeroize; use zeroize::Zeroize;
use super::migrations; use super::migrations;
use super::types::IdentityRecord; use super::types::{EphemeralKeyRecord, IdentityRecord};
use crate::crypto::PrivateKey;
use crate::identity::Identity; use crate::identity::Identity;
/// Chat-specific storage operations. /// Chat-specific storage operations.
@ -47,6 +48,73 @@ impl ChatStorage {
Ok(()) 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<Vec<EphemeralKeyRecord>, 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<u8> = row.get(1)?;
Ok((public_key_hex, secret_key))
})?
.collect::<Result<Vec<_>, _>>()?;
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. /// Loads the identity if it exists.
/// ///
/// Note: Secret key bytes are zeroized after being copied into IdentityRecord, /// Note: Secret key bytes are zeroized after being copied into IdentityRecord,

View File

@ -7,10 +7,16 @@ use storage::{Connection, StorageError};
/// Embeds and returns all migration SQL files in order. /// Embeds and returns all migration SQL files in order.
pub fn get_migrations() -> Vec<(&'static str, &'static str)> { pub fn get_migrations() -> Vec<(&'static str, &'static str)> {
vec![( vec![
"001_initial_schema", (
include_str!("migrations/001_initial_schema.sql"), "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. /// Applies all migrations to the database.

View File

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

View File

@ -17,12 +17,21 @@ pub struct IdentityRecord {
impl From<IdentityRecord> for Identity { impl From<IdentityRecord> for Identity {
fn from(record: IdentityRecord) -> Self { fn from(record: IdentityRecord) -> Self {
let name = record.name.clone();
let secret = PrivateKey::from(record.secret_key); 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)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;