mirror of
https://github.com/logos-messaging/libchat.git
synced 2026-04-01 09:04:03 +00:00
Merge c4566bb7ce147779c372a3e3c37648993bbc2a00 into 8cddd9ddcfb446deeff96fd5a68d6e4b14927d9f
This commit is contained in:
kaichao
GitHub
commit
83803d32d5
@ -1,9 +1,11 @@
|
||||
use std::rc::Rc;
|
||||
use std::sync::Arc;
|
||||
|
||||
use double_ratchets::{RatchetState, RatchetStorage};
|
||||
use storage::StorageConfig;
|
||||
|
||||
use crate::{
|
||||
conversation::{ConversationId, ConversationStore, Convo, Id},
|
||||
conversation::{ConversationId, Convo, Id, PrivateV1Convo},
|
||||
errors::ChatError,
|
||||
identity::Identity,
|
||||
inbox::Inbox,
|
||||
@ -19,10 +21,9 @@ pub use crate::inbox::Introduction;
|
||||
// Ctx manages lifetimes of objects to process and generate payloads.
|
||||
pub struct Context {
|
||||
_identity: Rc<Identity>,
|
||||
store: ConversationStore,
|
||||
inbox: Inbox,
|
||||
#[allow(dead_code)] // Will be used for conversation persistence
|
||||
storage: ChatStorage,
|
||||
ratchet_storage: RatchetStorage,
|
||||
}
|
||||
|
||||
impl Context {
|
||||
@ -31,7 +32,8 @@ impl Context {
|
||||
/// If an identity exists in storage, it will be restored.
|
||||
/// Otherwise, a new identity will be created with the given name and saved.
|
||||
pub fn open(name: impl Into<String>, config: StorageConfig) -> Result<Self, ChatError> {
|
||||
let mut storage = ChatStorage::new(config)?;
|
||||
let mut storage = ChatStorage::new(config.clone())?;
|
||||
let ratchet_storage = RatchetStorage::from_config(config)?;
|
||||
let name = name.into();
|
||||
|
||||
// Load or create identity
|
||||
@ -48,9 +50,9 @@ impl Context {
|
||||
|
||||
Ok(Self {
|
||||
_identity: identity,
|
||||
store: ConversationStore::new(),
|
||||
inbox,
|
||||
storage,
|
||||
ratchet_storage,
|
||||
})
|
||||
}
|
||||
|
||||
@ -81,12 +83,16 @@ impl Context {
|
||||
.map(|p| p.into_envelope(remote_id.clone()))
|
||||
.collect();
|
||||
|
||||
let convo_id = self.add_convo(Box::new(convo));
|
||||
let convo_id = self.persist_convo(&convo);
|
||||
(convo_id, payload_bytes)
|
||||
}
|
||||
|
||||
pub fn list_conversations(&self) -> Result<Vec<ConversationIdOwned>, ChatError> {
|
||||
Ok(self.store.conversation_ids())
|
||||
let records = self.storage.load_conversations()?;
|
||||
Ok(records
|
||||
.into_iter()
|
||||
.map(|r| Arc::from(r.local_convo_id.as_str()))
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub fn send_content(
|
||||
@ -94,16 +100,15 @@ impl Context {
|
||||
convo_id: ConversationId,
|
||||
content: &[u8],
|
||||
) -> Result<Vec<AddressedEnvelope>, ChatError> {
|
||||
// Lookup convo by id
|
||||
let convo = self.get_convo_mut(convo_id)?;
|
||||
let mut convo = self.load_convo(convo_id)?;
|
||||
|
||||
// Generate encrypted payloads
|
||||
let payloads = convo.send_message(content)?;
|
||||
let remote_id = convo.remote_id();
|
||||
convo.save_ratchet_state(&mut self.ratchet_storage)?;
|
||||
|
||||
// Attach conversation_ids to Envelopes
|
||||
Ok(payloads
|
||||
.into_iter()
|
||||
.map(|p| p.into_envelope(convo.remote_id()))
|
||||
.map(|p| p.into_envelope(remote_id.clone()))
|
||||
.collect())
|
||||
}
|
||||
|
||||
@ -116,7 +121,7 @@ impl Context {
|
||||
let enc = EncryptedPayload::decode(env.payload)?;
|
||||
match convo_id {
|
||||
c if c == self.inbox.id() => self.dispatch_to_inbox(enc),
|
||||
c if self.store.has(&c) => self.dispatch_to_convo(&c, enc),
|
||||
c if self.storage.has_conversation(&c)? => self.dispatch_to_convo(&c, enc),
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
@ -126,8 +131,19 @@ impl Context {
|
||||
&mut self,
|
||||
enc_payload: EncryptedPayload,
|
||||
) -> Result<Option<ContentData>, ChatError> {
|
||||
let (convo, content) = self.inbox.handle_frame(enc_payload)?;
|
||||
self.add_convo(convo);
|
||||
// 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.persist_convo(convo.as_ref());
|
||||
Ok(content)
|
||||
}
|
||||
|
||||
@ -137,44 +153,57 @@ impl Context {
|
||||
convo_id: ConversationId,
|
||||
enc_payload: EncryptedPayload,
|
||||
) -> Result<Option<ContentData>, ChatError> {
|
||||
let Some(convo) = self.store.get_mut(convo_id) else {
|
||||
return Err(ChatError::Protocol("convo id not found".into()));
|
||||
};
|
||||
let mut convo = self.load_convo(convo_id)?;
|
||||
|
||||
convo.handle_frame(enc_payload)
|
||||
let result = convo.handle_frame(enc_payload)?;
|
||||
convo.save_ratchet_state(&mut self.ratchet_storage)?;
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
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 {
|
||||
self.store.insert_convo(convo)
|
||||
/// Loads a conversation from DB by constructing it from metadata + ratchet state.
|
||||
fn load_convo(&self, convo_id: ConversationId) -> Result<PrivateV1Convo, ChatError> {
|
||||
let record = self
|
||||
.storage
|
||||
.load_conversation(convo_id)?
|
||||
.ok_or_else(|| ChatError::NoConvo(convo_id.into()))?;
|
||||
|
||||
if record.convo_type != "private_v1" {
|
||||
return Err(ChatError::BadBundleValue(format!(
|
||||
"unsupported conversation type: {}",
|
||||
record.convo_type
|
||||
)));
|
||||
}
|
||||
|
||||
let dr_state: RatchetState = self.ratchet_storage.load(&record.local_convo_id)?;
|
||||
|
||||
Ok(PrivateV1Convo::from_stored(
|
||||
record.local_convo_id,
|
||||
record.remote_convo_id,
|
||||
dr_state,
|
||||
))
|
||||
}
|
||||
|
||||
// Returns a mutable reference to a Convo for a given ConvoId
|
||||
fn get_convo_mut(&mut self, convo_id: ConversationId) -> Result<&mut dyn Convo, ChatError> {
|
||||
self.store
|
||||
.get_mut(convo_id)
|
||||
.ok_or_else(|| ChatError::NoConvo(convo_id.into()))
|
||||
/// Persists a conversation's metadata and ratchet state to DB.
|
||||
fn persist_convo(&mut self, convo: &dyn Convo) -> ConversationIdOwned {
|
||||
let _ = self
|
||||
.storage
|
||||
.save_conversation(convo.id(), &convo.remote_id(), convo.convo_type());
|
||||
let _ = convo.save_ratchet_state(&mut self.ratchet_storage);
|
||||
Arc::from(convo.id())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::conversation::GroupTestConvo;
|
||||
|
||||
#[test]
|
||||
fn convo_store_get() {
|
||||
let mut store: ConversationStore = ConversationStore::new();
|
||||
|
||||
let new_convo = GroupTestConvo::new();
|
||||
let convo_id = store.insert_convo(Box::new(new_convo));
|
||||
|
||||
let convo = store.get_mut(&convo_id).ok_or(0);
|
||||
convo.unwrap();
|
||||
}
|
||||
|
||||
fn send_and_verify(
|
||||
sender: &mut Context,
|
||||
@ -228,7 +257,6 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn identity_persistence() {
|
||||
// Use file-based storage to test real persistence
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db_path = dir
|
||||
.path()
|
||||
@ -237,19 +265,136 @@ mod tests {
|
||||
.to_string();
|
||||
let config = StorageConfig::File(db_path);
|
||||
|
||||
// Create context - this should create and save a new identity
|
||||
let ctx1 = Context::open("alice", config.clone()).unwrap();
|
||||
let pubkey1 = ctx1._identity.public_key();
|
||||
let name1 = ctx1.installation_name().to_string();
|
||||
|
||||
// Drop and reopen - should load the same identity
|
||||
drop(ctx1);
|
||||
let ctx2 = Context::open("alice", config).unwrap();
|
||||
let pubkey2 = ctx2._identity.public_key();
|
||||
let name2 = ctx2.installation_name().to_string();
|
||||
|
||||
// Identity should be the same
|
||||
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);
|
||||
|
||||
let mut ctx1 = Context::open("alice", config.clone()).unwrap();
|
||||
let bundle1 = ctx1.create_intro_bundle().unwrap();
|
||||
|
||||
drop(ctx1);
|
||||
let mut ctx2 = Context::open("alice", config.clone()).unwrap();
|
||||
|
||||
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");
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn conversation_metadata_persistence() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db_path = dir
|
||||
.path()
|
||||
.join("test_convo_meta.db")
|
||||
.to_string_lossy()
|
||||
.to_string();
|
||||
let config = StorageConfig::File(db_path);
|
||||
|
||||
let mut alice = Context::open("alice", config.clone()).unwrap();
|
||||
let mut bob = Context::new_with_name("bob");
|
||||
|
||||
let bundle = alice.create_intro_bundle().unwrap();
|
||||
let intro = Introduction::try_from(bundle.as_slice()).unwrap();
|
||||
let (_, payloads) = bob.create_private_convo(&intro, b"hi");
|
||||
|
||||
let payload = payloads.first().unwrap();
|
||||
let content = alice.handle_payload(&payload.data).unwrap().unwrap();
|
||||
assert!(content.is_new_convo);
|
||||
|
||||
let convos = alice.storage.load_conversations().unwrap();
|
||||
assert_eq!(convos.len(), 1);
|
||||
assert_eq!(convos[0].convo_type, "private_v1");
|
||||
|
||||
drop(alice);
|
||||
let alice2 = Context::open("alice", config).unwrap();
|
||||
let convos = alice2.storage.load_conversations().unwrap();
|
||||
assert_eq!(convos.len(), 1, "conversation metadata should persist");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn conversation_full_persistence() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db_path = dir
|
||||
.path()
|
||||
.join("test_full_persist.db")
|
||||
.to_string_lossy()
|
||||
.to_string();
|
||||
let config = StorageConfig::File(db_path);
|
||||
|
||||
// Alice and Bob establish a conversation
|
||||
let mut alice = Context::open("alice", config.clone()).unwrap();
|
||||
let mut bob = Context::new_with_name("bob");
|
||||
|
||||
let bundle = alice.create_intro_bundle().unwrap();
|
||||
let intro = Introduction::try_from(bundle.as_slice()).unwrap();
|
||||
let (bob_convo_id, payloads) = bob.create_private_convo(&intro, b"hello");
|
||||
|
||||
let payload = payloads.first().unwrap();
|
||||
let content = alice.handle_payload(&payload.data).unwrap().unwrap();
|
||||
let alice_convo_id = content.conversation_id;
|
||||
|
||||
// Exchange a few messages to advance ratchet state
|
||||
let payloads = alice.send_content(&alice_convo_id, b"reply 1").unwrap();
|
||||
let payload = payloads.first().unwrap();
|
||||
bob.handle_payload(&payload.data).unwrap().unwrap();
|
||||
|
||||
let payloads = bob.send_content(&bob_convo_id, b"reply 2").unwrap();
|
||||
let payload = payloads.first().unwrap();
|
||||
alice.handle_payload(&payload.data).unwrap().unwrap();
|
||||
|
||||
// Drop Alice and reopen - conversation should survive
|
||||
drop(alice);
|
||||
let mut alice2 = Context::open("alice", config).unwrap();
|
||||
|
||||
// Verify conversation was restored
|
||||
let convo_ids = alice2.list_conversations().unwrap();
|
||||
assert_eq!(convo_ids.len(), 1);
|
||||
|
||||
// Bob sends a new message - Alice should be able to decrypt after restart
|
||||
let payloads = bob.send_content(&bob_convo_id, b"after restart").unwrap();
|
||||
let payload = payloads.first().unwrap();
|
||||
let content = alice2
|
||||
.handle_payload(&payload.data)
|
||||
.expect("should decrypt after restart")
|
||||
.expect("should have content");
|
||||
assert_eq!(content.data, b"after restart");
|
||||
|
||||
// Alice can also send back
|
||||
let payloads = alice2
|
||||
.send_content(&alice_convo_id, b"alice after restart")
|
||||
.unwrap();
|
||||
let payload = payloads.first().unwrap();
|
||||
let content = bob
|
||||
.handle_payload(&payload.data)
|
||||
.unwrap()
|
||||
.expect("bob should receive");
|
||||
assert_eq!(content.data, b"alice after restart");
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,9 +1,9 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fmt::Debug;
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use crate::errors::ChatError;
|
||||
use crate::types::{AddressedEncryptedPayload, ContentData};
|
||||
use double_ratchets::RatchetStorage;
|
||||
|
||||
pub type ConversationId<'a> = &'a str;
|
||||
pub type ConversationIdOwned = Arc<str>;
|
||||
@ -27,44 +27,17 @@ pub trait Convo: Id + Debug {
|
||||
) -> Result<Option<ContentData>, ChatError>;
|
||||
|
||||
fn remote_id(&self) -> String;
|
||||
}
|
||||
|
||||
pub struct ConversationStore {
|
||||
conversations: HashMap<Arc<str>, Box<dyn Convo>>,
|
||||
}
|
||||
/// Returns the conversation type identifier for storage.
|
||||
fn convo_type(&self) -> &str;
|
||||
|
||||
impl ConversationStore {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
conversations: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn insert_convo(&mut self, conversation: Box<dyn Convo>) -> ConversationIdOwned {
|
||||
let key: ConversationIdOwned = Arc::from(conversation.id());
|
||||
self.conversations.insert(key.clone(), conversation);
|
||||
key
|
||||
}
|
||||
|
||||
pub fn has(&self, id: ConversationId) -> bool {
|
||||
self.conversations.contains_key(id)
|
||||
}
|
||||
|
||||
pub fn get_mut(&mut self, id: &str) -> Option<&mut (dyn Convo + '_)> {
|
||||
Some(self.conversations.get_mut(id)?.as_mut())
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn conversation_ids(&self) -> Vec<ConversationIdOwned> {
|
||||
self.conversations.keys().cloned().collect()
|
||||
/// Persists ratchet state to storage. Default is no-op.
|
||||
fn save_ratchet_state(&self, _storage: &mut RatchetStorage) -> Result<(), ChatError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod group_test;
|
||||
mod privatev1;
|
||||
|
||||
use chat_proto::logoschat::encryption::EncryptedPayload;
|
||||
#[cfg(test)]
|
||||
pub(crate) use group_test::GroupTestConvo;
|
||||
pub use privatev1::PrivateV1Convo;
|
||||
|
||||
@ -18,6 +18,7 @@ use crate::{
|
||||
types::{AddressedEncryptedPayload, ContentData},
|
||||
utils::timestamp_millis,
|
||||
};
|
||||
use double_ratchets::RatchetStorage;
|
||||
|
||||
// Represents the potential participant roles in this Conversation
|
||||
enum Role {
|
||||
@ -77,6 +78,19 @@ impl PrivateV1Convo {
|
||||
}
|
||||
}
|
||||
|
||||
/// Reconstructs a PrivateV1Convo from persisted metadata and ratchet state.
|
||||
pub fn from_stored(
|
||||
local_convo_id: String,
|
||||
remote_convo_id: String,
|
||||
dr_state: RatchetState,
|
||||
) -> Self {
|
||||
Self {
|
||||
local_convo_id,
|
||||
remote_convo_id,
|
||||
dr_state,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_responder(seed_key: SymmetricKey32, dh_self: &PrivateKey) -> Self {
|
||||
let base_convo_id = BaseConvoId::new(&seed_key);
|
||||
let local_convo_id = base_convo_id.id_for_participant(Role::Responder);
|
||||
@ -209,6 +223,15 @@ impl Convo for PrivateV1Convo {
|
||||
fn remote_id(&self) -> String {
|
||||
self.remote_convo_id.clone()
|
||||
}
|
||||
|
||||
fn convo_type(&self) -> &str {
|
||||
"private_v1"
|
||||
}
|
||||
|
||||
fn save_ratchet_state(&self, storage: &mut RatchetStorage) -> Result<(), ChatError> {
|
||||
storage.save(&self.local_convo_id, &self.dr_state)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Debug for PrivateV1Convo {
|
||||
|
||||
@ -22,6 +22,8 @@ pub enum ChatError {
|
||||
BadParsing(&'static str),
|
||||
#[error("convo with id: {0} was not found")]
|
||||
NoConvo(String),
|
||||
#[error("unsupported conversation type: {0}")]
|
||||
UnsupportedConvoType(String),
|
||||
#[error("storage error: {0}")]
|
||||
Storage(#[from] StorageError),
|
||||
}
|
||||
|
||||
@ -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<Identity>,
|
||||
local_convo_id: String,
|
||||
ephemeral_keys: HashMap<String, PrivateKey>,
|
||||
}
|
||||
|
||||
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,18 +41,19 @@ impl Inbox {
|
||||
Self {
|
||||
ident,
|
||||
local_convo_id,
|
||||
ephemeral_keys: HashMap::<String, PrivateKey>::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn create_intro_bundle(&mut self) -> Introduction {
|
||||
/// 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();
|
||||
self.ephemeral_keys
|
||||
.insert(hex::encode(ephemeral_key.as_bytes()), ephemeral);
|
||||
let public_key_hex = hex::encode(ephemeral_key.as_bytes());
|
||||
|
||||
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,8 +109,11 @@ impl Inbox {
|
||||
Ok((convo, payloads))
|
||||
}
|
||||
|
||||
/// 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<dyn Convo>, Option<ContentData>), ChatError> {
|
||||
let handshake = Self::extract_payload(enc_payload)?;
|
||||
@ -124,10 +122,6 @@ impl Inbox {
|
||||
.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)?;
|
||||
|
||||
@ -153,6 +147,24 @@ impl Inbox {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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<String, ChatError> {
|
||||
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(),
|
||||
@ -214,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))
|
||||
@ -235,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 = 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(),
|
||||
|
||||
@ -3,10 +3,14 @@
|
||||
mod migrations;
|
||||
mod types;
|
||||
|
||||
use crypto::PrivateKey;
|
||||
use storage::{RusqliteError, SqliteDb, StorageConfig, StorageError, params};
|
||||
use zeroize::Zeroize;
|
||||
|
||||
use crate::{identity::Identity, storage::types::IdentityRecord};
|
||||
use crate::{
|
||||
identity::Identity,
|
||||
storage::types::{ConversationRecord, IdentityRecord},
|
||||
};
|
||||
|
||||
/// Chat-specific storage operations.
|
||||
///
|
||||
@ -87,6 +91,148 @@ impl ChatStorage {
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== 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 a single ephemeral key by its public key hex.
|
||||
pub fn load_ephemeral_key(
|
||||
&self,
|
||||
public_key_hex: &str,
|
||||
) -> Result<Option<PrivateKey>, StorageError> {
|
||||
let mut stmt = self
|
||||
.db
|
||||
.connection()
|
||||
.prepare("SELECT secret_key FROM ephemeral_keys WHERE public_key_hex = ?1")?;
|
||||
|
||||
let result = stmt.query_row(params![public_key_hex], |row| {
|
||||
let secret_key: Vec<u8> = row.get(0)?;
|
||||
Ok(secret_key)
|
||||
});
|
||||
|
||||
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()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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(())
|
||||
}
|
||||
|
||||
// ==================== Conversation Operations ====================
|
||||
|
||||
/// Saves conversation metadata.
|
||||
pub fn save_conversation(
|
||||
&mut self,
|
||||
local_convo_id: &str,
|
||||
remote_convo_id: &str,
|
||||
convo_type: &str,
|
||||
) -> Result<(), StorageError> {
|
||||
self.db.connection().execute(
|
||||
"INSERT OR REPLACE INTO conversations (local_convo_id, remote_convo_id, convo_type) VALUES (?1, ?2, ?3)",
|
||||
params![local_convo_id, remote_convo_id, convo_type],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Checks if a conversation exists by its local ID.
|
||||
pub fn has_conversation(&self, local_convo_id: &str) -> Result<bool, StorageError> {
|
||||
let exists: bool = self.db.connection().query_row(
|
||||
"SELECT EXISTS(SELECT 1 FROM conversations WHERE local_convo_id = ?1)",
|
||||
params![local_convo_id],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
Ok(exists)
|
||||
}
|
||||
|
||||
/// Removes a conversation by its local ID.
|
||||
#[allow(dead_code)]
|
||||
pub fn remove_conversation(&mut self, local_convo_id: &str) -> Result<(), StorageError> {
|
||||
self.db.connection().execute(
|
||||
"DELETE FROM conversations WHERE local_convo_id = ?1",
|
||||
params![local_convo_id],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Loads a single conversation record by its local ID.
|
||||
pub fn load_conversation(
|
||||
&self,
|
||||
local_convo_id: &str,
|
||||
) -> Result<Option<ConversationRecord>, StorageError> {
|
||||
let mut stmt = self.db.connection().prepare(
|
||||
"SELECT local_convo_id, remote_convo_id, convo_type FROM conversations WHERE local_convo_id = ?1",
|
||||
)?;
|
||||
|
||||
let result = stmt.query_row(params![local_convo_id], |row| {
|
||||
Ok(ConversationRecord {
|
||||
local_convo_id: row.get(0)?,
|
||||
remote_convo_id: row.get(1)?,
|
||||
convo_type: row.get(2)?,
|
||||
})
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(record) => Ok(Some(record)),
|
||||
Err(RusqliteError::QueryReturnedNoRows) => Ok(None),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads all conversation records.
|
||||
pub fn load_conversations(&self) -> Result<Vec<ConversationRecord>, StorageError> {
|
||||
let mut stmt = self
|
||||
.db
|
||||
.connection()
|
||||
.prepare("SELECT local_convo_id, remote_convo_id, convo_type FROM conversations")?;
|
||||
|
||||
let records = stmt
|
||||
.query_map([], |row| {
|
||||
Ok(ConversationRecord {
|
||||
local_convo_id: row.get(0)?,
|
||||
remote_convo_id: row.get(1)?,
|
||||
convo_type: row.get(2)?,
|
||||
})
|
||||
})?
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
|
||||
Ok(records)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@ -109,4 +255,53 @@ 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());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_conversation_roundtrip() {
|
||||
let mut storage = ChatStorage::new(StorageConfig::InMemory).unwrap();
|
||||
|
||||
// Initially empty
|
||||
let convos = storage.load_conversations().unwrap();
|
||||
assert!(convos.is_empty());
|
||||
|
||||
// Save conversations
|
||||
storage
|
||||
.save_conversation("local_1", "remote_1", "private_v1")
|
||||
.unwrap();
|
||||
storage
|
||||
.save_conversation("local_2", "remote_2", "private_v1")
|
||||
.unwrap();
|
||||
|
||||
let convos = storage.load_conversations().unwrap();
|
||||
assert_eq!(convos.len(), 2);
|
||||
|
||||
// Remove one
|
||||
storage.remove_conversation("local_1").unwrap();
|
||||
let convos = storage.load_conversations().unwrap();
|
||||
assert_eq!(convos.len(), 1);
|
||||
assert_eq!(convos[0].local_convo_id, "local_2");
|
||||
assert_eq!(convos[0].remote_convo_id, "remote_2");
|
||||
assert_eq!(convos[0].convo_type, "private_v1");
|
||||
}
|
||||
}
|
||||
|
||||
@ -7,3 +7,17 @@ CREATE TABLE IF NOT EXISTS identity (
|
||||
name TEXT NOT NULL,
|
||||
secret_key BLOB NOT NULL
|
||||
);
|
||||
|
||||
-- Ephemeral keys for inbox handshakes
|
||||
CREATE TABLE IF NOT EXISTS ephemeral_keys (
|
||||
public_key_hex TEXT PRIMARY KEY,
|
||||
secret_key BLOB NOT NULL
|
||||
);
|
||||
|
||||
-- Conversations metadata
|
||||
CREATE TABLE IF NOT EXISTS conversations (
|
||||
local_convo_id TEXT PRIMARY KEY,
|
||||
remote_convo_id TEXT NOT NULL,
|
||||
convo_type TEXT NOT NULL,
|
||||
created_at INTEGER NOT NULL DEFAULT (strftime('%s', 'now'))
|
||||
);
|
||||
|
||||
@ -22,6 +22,13 @@ impl From<IdentityRecord> for Identity {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ConversationRecord {
|
||||
pub local_convo_id: String,
|
||||
pub remote_convo_id: String,
|
||||
pub convo_type: String,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
@ -59,6 +59,12 @@ impl RatchetStorage {
|
||||
Self::run_migration(db)
|
||||
}
|
||||
|
||||
/// Creates a ratchet storage from a generic storage configuration.
|
||||
pub fn from_config(config: storage::StorageConfig) -> Result<Self, StorageError> {
|
||||
let db = SqliteDb::new(config)?;
|
||||
Self::run_migration(db)
|
||||
}
|
||||
|
||||
/// Creates a new ratchet storage with the given database.
|
||||
fn run_migration(db: SqliteDb) -> Result<Self, StorageError> {
|
||||
// Initialize schema
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user