From db813baf25dac03a9a684b52642335cd918e9d27 Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Tue, 24 Oct 2023 15:16:52 -0700 Subject: [PATCH 01/24] Add key_store table and generate --- xmtp_mls/{src => }/diesel.toml | 0 .../down.sql | 1 + .../2023-10-24-213844_create_key_store/up.sql | 4 +++ xmtp_mls/migrations/README.md | 25 +++++++++++++++++++ xmtp_mls/migrations/ignore.md | 1 - .../src/storage/encrypted_store/schema.rs | 7 ++++++ 6 files changed, 37 insertions(+), 1 deletion(-) rename xmtp_mls/{src => }/diesel.toml (100%) create mode 100644 xmtp_mls/migrations/2023-10-24-213844_create_key_store/down.sql create mode 100644 xmtp_mls/migrations/2023-10-24-213844_create_key_store/up.sql create mode 100644 xmtp_mls/migrations/README.md delete mode 100644 xmtp_mls/migrations/ignore.md diff --git a/xmtp_mls/src/diesel.toml b/xmtp_mls/diesel.toml similarity index 100% rename from xmtp_mls/src/diesel.toml rename to xmtp_mls/diesel.toml diff --git a/xmtp_mls/migrations/2023-10-24-213844_create_key_store/down.sql b/xmtp_mls/migrations/2023-10-24-213844_create_key_store/down.sql new file mode 100644 index 000000000..d7c804b80 --- /dev/null +++ b/xmtp_mls/migrations/2023-10-24-213844_create_key_store/down.sql @@ -0,0 +1 @@ +DROP TABLE openmls_key_store; diff --git a/xmtp_mls/migrations/2023-10-24-213844_create_key_store/up.sql b/xmtp_mls/migrations/2023-10-24-213844_create_key_store/up.sql new file mode 100644 index 000000000..8f4b81a93 --- /dev/null +++ b/xmtp_mls/migrations/2023-10-24-213844_create_key_store/up.sql @@ -0,0 +1,4 @@ +CREATE TABLE IF NOT EXISTS openmls_key_store ( + key_bytes BLOB PRIMARY KEY NOT NULL, + value_bytes BLOB NOT NULL +); diff --git a/xmtp_mls/migrations/README.md b/xmtp_mls/migrations/README.md new file mode 100644 index 000000000..dfe5006d9 --- /dev/null +++ b/xmtp_mls/migrations/README.md @@ -0,0 +1,25 @@ +# Steps for setting up a diesel migration + +### Install the CLI onto your local system (one-time) + +``` +cargo install diesel_cli --no-default-features --features sqlite +``` + +### Create your migration SQL + +In this example the migration is called `create_key_store`: + +``` +diesel migration generate create_key_store +``` + +Edit the `up.sql` and `down.sql` files created + +### Generate application code + +``` +cargo run --bin update-schema +``` + +This updates the generated `schema.rs` file. You can now update `models.rs` to reference it and consume your new model in the rest of the crate. diff --git a/xmtp_mls/migrations/ignore.md b/xmtp_mls/migrations/ignore.md deleted file mode 100644 index 67e54d9c8..000000000 --- a/xmtp_mls/migrations/ignore.md +++ /dev/null @@ -1 +0,0 @@ -Placeholder file to create migrations directory - remove after first migration. diff --git a/xmtp_mls/src/storage/encrypted_store/schema.rs b/xmtp_mls/src/storage/encrypted_store/schema.rs index d9a52af7e..9de21ee27 100644 --- a/xmtp_mls/src/storage/encrypted_store/schema.rs +++ b/xmtp_mls/src/storage/encrypted_store/schema.rs @@ -1 +1,8 @@ // @generated automatically by Diesel CLI. + +diesel::table! { + openmls_key_store (key_bytes) { + key_bytes -> Binary, + value_bytes -> Binary, + } +} From f9c6ec87d5189ad50e83945ba9ae24e7948b8f8f Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Tue, 24 Oct 2023 15:49:03 -0700 Subject: [PATCH 02/24] Implement store/fetch/delete traits --- xmtp_mls/migrations/README.md | 2 +- xmtp_mls/src/lib.rs | 10 +++++-- xmtp_mls/src/storage/encrypted_store/mod.rs | 30 +++++++++++++++++++ .../src/storage/encrypted_store/models.rs | 10 +++++++ 4 files changed, 48 insertions(+), 4 deletions(-) diff --git a/xmtp_mls/migrations/README.md b/xmtp_mls/migrations/README.md index dfe5006d9..6c8ed1457 100644 --- a/xmtp_mls/migrations/README.md +++ b/xmtp_mls/migrations/README.md @@ -22,4 +22,4 @@ Edit the `up.sql` and `down.sql` files created cargo run --bin update-schema ``` -This updates the generated `schema.rs` file. You can now update `models.rs` to reference it and consume your new model in the rest of the crate. +This updates the generated `schema.rs` file. You can now update `models.rs` to reference it and `encrypted_store/mod.rs` to define queries against the model. diff --git a/xmtp_mls/src/lib.rs b/xmtp_mls/src/lib.rs index 64b5dbae4..4dd0a6b01 100644 --- a/xmtp_mls/src/lib.rs +++ b/xmtp_mls/src/lib.rs @@ -26,10 +26,14 @@ pub trait Store { pub trait Fetch { type Key<'a>; - // Fetches all instances of a model from the underlying data store - fn fetch_all(&mut self) -> Result, StorageError>; + // Fetches a single instance by key of a model from the underlying data store + #[allow(clippy::needless_lifetimes)] + fn fetch<'a>(&mut self, key: Self::Key<'a>) -> Result, StorageError>; +} +pub trait Delete { + type Key<'a>; // Fetches a single instance by key of a model from the underlying data store #[allow(clippy::needless_lifetimes)] - fn fetch_one<'a>(&mut self, key: Self::Key<'a>) -> Result, StorageError>; + fn delete<'a>(&mut self, key: Self::Key<'a>) -> Result; } diff --git a/xmtp_mls/src/storage/encrypted_store/mod.rs b/xmtp_mls/src/storage/encrypted_store/mod.rs index 45bec5433..9e06e9bdb 100644 --- a/xmtp_mls/src/storage/encrypted_store/mod.rs +++ b/xmtp_mls/src/storage/encrypted_store/mod.rs @@ -14,6 +14,10 @@ pub mod models; pub mod schema; +use crate::{Delete, Fetch, Store}; + +use self::{models::StoredKeyStoreEntry, schema::openmls_key_store}; + use super::StorageError; use diesel::{ connection::SimpleConnection, @@ -163,6 +167,32 @@ fn warn_length(list: &Vec, str_id: &str, max_length: usize) { } } +impl Store for StoredKeyStoreEntry { + fn store(&self, into: &mut DbConnection) -> Result<(), StorageError> { + diesel::insert_into(openmls_key_store::table) + .values(self) + .execute(into)?; + + Ok(()) + } +} + +impl Fetch for DbConnection { + type Key<'a> = Vec; + fn fetch(&mut self, key: Vec) -> Result, StorageError> where { + use self::schema::openmls_key_store::dsl::*; + Ok(openmls_key_store.find(key).first(self).optional()?) + } +} + +impl Delete for DbConnection { + type Key<'a> = Vec; + fn delete(&mut self, key: Vec) -> Result where { + use self::schema::openmls_key_store::dsl::*; + Ok(diesel::delete(openmls_key_store.filter(key_bytes.eq(key))).execute(self)?) + } +} + #[cfg(test)] mod tests { #[test] diff --git a/xmtp_mls/src/storage/encrypted_store/models.rs b/xmtp_mls/src/storage/encrypted_store/models.rs index 8b1378917..49b06e5dc 100644 --- a/xmtp_mls/src/storage/encrypted_store/models.rs +++ b/xmtp_mls/src/storage/encrypted_store/models.rs @@ -1 +1,11 @@ +use diesel::prelude::*; +use super::schema::*; + +#[derive(Insertable, Queryable, Debug, Clone)] +#[diesel(table_name = openmls_key_store)] +#[diesel(primary_key(key_bytes))] +pub struct StoredKeyStoreEntry { + pub key_bytes: Vec, + pub value_bytes: Vec, +} From 257d5b6b9e3ed4c7eb373ddc08098055aae13560 Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Tue, 24 Oct 2023 16:29:25 -0700 Subject: [PATCH 03/24] Add SQL key store --- xmtp_mls/Cargo.toml | 3 +- xmtp_mls/src/storage/encrypted_store/mod.rs | 2 +- xmtp_mls/src/storage/mod.rs | 1 + xmtp_mls/src/storage/sql_key_store.rs | 102 ++++++++++++++++++++ 4 files changed, 106 insertions(+), 2 deletions(-) create mode 100644 xmtp_mls/src/storage/sql_key_store.rs diff --git a/xmtp_mls/Cargo.toml b/xmtp_mls/Cargo.toml index 2ca0f8bf2..84d629f11 100644 --- a/xmtp_mls/Cargo.toml +++ b/xmtp_mls/Cargo.toml @@ -11,6 +11,7 @@ path = "src/bin/update-schema.rs" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +anyhow = "1.0.71" async-trait = "0.1.68" diesel = { version = "2.0.4", features = ["sqlite", "r2d2", "returning_clauses_for_sqlite_3_35"] } diesel_migrations = { version = "2.0.0", features = ["sqlite"] } @@ -33,7 +34,7 @@ tokio = { version = "1.28.1", features = ["macros"] } toml = "0.7.4" xmtp_cryptography = { path = "../xmtp_cryptography"} xmtp_proto = { path = "../xmtp_proto", features = ["proto_full"] } -anyhow = "1.0.71" [dev-dependencies] +rand = "0.8.5" tempfile = "3.5.0" diff --git a/xmtp_mls/src/storage/encrypted_store/mod.rs b/xmtp_mls/src/storage/encrypted_store/mod.rs index 9e06e9bdb..b6cb8c43f 100644 --- a/xmtp_mls/src/storage/encrypted_store/mod.rs +++ b/xmtp_mls/src/storage/encrypted_store/mod.rs @@ -56,7 +56,7 @@ pub fn ignore_unique_violation( } #[allow(dead_code)] -#[derive(Clone)] +#[derive(Clone, Debug)] /// Manages a Sqlite db for persisting messages and other objects. pub struct EncryptedMessageStore { connect_opt: StorageOption, diff --git a/xmtp_mls/src/storage/mod.rs b/xmtp_mls/src/storage/mod.rs index 4f60edd20..c70f22a51 100644 --- a/xmtp_mls/src/storage/mod.rs +++ b/xmtp_mls/src/storage/mod.rs @@ -1,6 +1,7 @@ mod encrypted_store; mod errors; pub mod in_memory_key_store; +pub mod sql_key_store; pub use encrypted_store::{DbConnection, EncryptedMessageStore, EncryptionKey, StorageOption}; pub use errors::StorageError; diff --git a/xmtp_mls/src/storage/sql_key_store.rs b/xmtp_mls/src/storage/sql_key_store.rs new file mode 100644 index 000000000..242263aab --- /dev/null +++ b/xmtp_mls/src/storage/sql_key_store.rs @@ -0,0 +1,102 @@ +use log::{debug, error}; +use openmls_traits::key_store::{MlsEntity, OpenMlsKeyStore}; + +use crate::{Delete, Fetch, Store}; + +use super::{encrypted_store::models::StoredKeyStoreEntry, EncryptedMessageStore, StorageError}; + +#[derive(Debug)] +pub struct SqlKeyStore { + store: EncryptedMessageStore, +} + +impl OpenMlsKeyStore for SqlKeyStore { + /// The error type returned by the [`OpenMlsKeyStore`]. + type Error = StorageError; + + /// Store a value `v` that implements the [`ToKeyStoreValue`] trait for + /// serialization for ID `k`. + /// + /// Returns an error if storing fails. + fn store(&self, k: &[u8], v: &V) -> Result<(), Self::Error> { + let entry = StoredKeyStoreEntry { + key_bytes: k.to_vec(), + value_bytes: serde_json::to_vec(v).map_err(|_| StorageError::SerializationError)?, + }; + entry.store(&mut self.store.conn()?)?; + Ok(()) + } + + /// Read and return a value stored for ID `k` that implements the + /// [`FromKeyStoreValue`] trait for deserialization. + /// + /// Returns [`None`] if no value is stored for `k` or reading fails. + fn read(&self, k: &[u8]) -> Option { + let conn_result = self.store.conn(); + if let Err(e) = conn_result { + error!("Failed to get connection: {:?}", e); + return None; + } + let mut conn = conn_result.unwrap(); + let fetch_result = conn.fetch(k.to_vec()); + if let Err(e) = fetch_result { + error!("Failed to fetch key: {:?}", e); + return None; + } + let entry_option = fetch_result.unwrap(); + if entry_option.is_none() { + debug!("No entry to read for key {:?}", k); + return None; + } + serde_json::from_slice(&entry_option.unwrap().value_bytes).ok() + } + + /// Delete a value stored for ID `k`. + /// + /// Interface is unclear on expected behavior when item is already deleted - + /// we choose to not surface an error if this is the case. + fn delete(&self, k: &[u8]) -> Result<(), Self::Error> { + let num_deleted = self.store.conn()?.delete(k.to_vec())?; + if num_deleted == 0 { + debug!("No entry to delete for key {:?}", k); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use openmls_basic_credential::SignatureKeyPair; + use openmls_traits::key_store::OpenMlsKeyStore; + use rand::distributions::{Alphanumeric, DistString}; + + use crate::{ + configuration::CIPHERSUITE, + storage::{EncryptedMessageStore, StorageOption}, + }; + + use super::SqlKeyStore; + + fn rand_string() -> String { + Alphanumeric.sample_string(&mut rand::thread_rng(), 16) + } + + #[test] + fn store_read_delete() { + let db_path = format!("{}.db3", rand_string()); + let key_store = SqlKeyStore { + store: EncryptedMessageStore::new( + StorageOption::Persistent(db_path), + EncryptedMessageStore::generate_enc_key(), + ) + .unwrap(), + }; + let signature_keys = SignatureKeyPair::new(CIPHERSUITE.signature_algorithm()).unwrap(); + let index = "index".as_bytes(); + assert!(key_store.read::(index).is_none()); + key_store.store(index, &signature_keys).unwrap(); + assert!(key_store.read::(index).is_some()); + key_store.delete::(index).unwrap(); + assert!(key_store.read::(index).is_none()); + } +} From 6af1384af851f8b1af1ea3d070791ea2ae21070a Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Wed, 25 Oct 2023 15:33:34 -0700 Subject: [PATCH 04/24] Use SQL key store for provider, remove in-memory key store --- xmtp_mls/src/storage/in_memory_key_store.rs | 76 --------------------- xmtp_mls/src/storage/mod.rs | 1 - xmtp_mls/src/storage/sql_key_store.rs | 2 +- xmtp_mls/src/xmtp_openmls_provider.rs | 6 +- 4 files changed, 4 insertions(+), 81 deletions(-) delete mode 100644 xmtp_mls/src/storage/in_memory_key_store.rs diff --git a/xmtp_mls/src/storage/in_memory_key_store.rs b/xmtp_mls/src/storage/in_memory_key_store.rs deleted file mode 100644 index 7f6d0f4a7..000000000 --- a/xmtp_mls/src/storage/in_memory_key_store.rs +++ /dev/null @@ -1,76 +0,0 @@ -use openmls_traits::key_store::{MlsEntity, OpenMlsKeyStore}; -use std::{collections::HashMap, sync::RwLock}; - -use super::StorageError; - -#[derive(Debug, Default)] -pub struct InMemoryKeyStore { - values: RwLock, Vec>>, -} - -impl OpenMlsKeyStore for InMemoryKeyStore { - /// The error type returned by the [`OpenMlsKeyStore`]. - type Error = StorageError; - - /// Store a value `v` that implements the [`ToKeyStoreValue`] trait for - /// serialization for ID `k`. - /// - /// Returns an error if storing fails. - fn store(&self, k: &[u8], v: &V) -> Result<(), Self::Error> { - let value = serde_json::to_vec(v).map_err(|_| StorageError::SerializationError)?; - // We unwrap here, because this is the only function claiming a write - // lock on `credential_bundles`. It only holds the lock very briefly and - // should not panic during that period. - let mut values = self.values.write().unwrap(); - values.insert(k.to_vec(), value); - Ok(()) - } - - /// Read and return a value stored for ID `k` that implements the - /// [`FromKeyStoreValue`] trait for deserialization. - /// - /// Returns [`None`] if no value is stored for `k` or reading fails. - fn read(&self, k: &[u8]) -> Option { - // We unwrap here, because the two functions claiming a write lock on - // `init_key_package_bundles` (this one and `generate_key_package_bundle`) only - // hold the lock very briefly and should not panic during that period. - let values = self.values.read().unwrap(); - if let Some(value) = values.get(k) { - serde_json::from_slice(value).ok() - } else { - None - } - } - - /// Delete a value stored for ID `k`. - /// - /// Returns an error if storing fails. - fn delete(&self, k: &[u8]) -> Result<(), Self::Error> { - // We just delete both ... - let mut values = self.values.write().unwrap(); - values.remove(k); - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use openmls_basic_credential::SignatureKeyPair; - use openmls_traits::key_store::OpenMlsKeyStore; - - use crate::configuration::CIPHERSUITE; - - use super::InMemoryKeyStore; - - #[test] - fn store_read_delete() { - let key_store = InMemoryKeyStore::default(); - let signature_keys = SignatureKeyPair::new(CIPHERSUITE.signature_algorithm()).unwrap(); - let index = "index".as_bytes(); - assert!(key_store.read::(index).is_none()); - key_store.store(index, &signature_keys).unwrap(); - assert!(key_store.read::(index).is_some()); - key_store.delete::(index).unwrap(); - assert!(key_store.read::(index).is_none()); - } -} diff --git a/xmtp_mls/src/storage/mod.rs b/xmtp_mls/src/storage/mod.rs index c70f22a51..6710adb90 100644 --- a/xmtp_mls/src/storage/mod.rs +++ b/xmtp_mls/src/storage/mod.rs @@ -1,6 +1,5 @@ mod encrypted_store; mod errors; -pub mod in_memory_key_store; pub mod sql_key_store; pub use encrypted_store::{DbConnection, EncryptedMessageStore, EncryptionKey, StorageOption}; diff --git a/xmtp_mls/src/storage/sql_key_store.rs b/xmtp_mls/src/storage/sql_key_store.rs index 242263aab..d0545b57d 100644 --- a/xmtp_mls/src/storage/sql_key_store.rs +++ b/xmtp_mls/src/storage/sql_key_store.rs @@ -5,7 +5,7 @@ use crate::{Delete, Fetch, Store}; use super::{encrypted_store::models::StoredKeyStoreEntry, EncryptedMessageStore, StorageError}; -#[derive(Debug)] +#[derive(Debug, Default)] pub struct SqlKeyStore { store: EncryptedMessageStore, } diff --git a/xmtp_mls/src/xmtp_openmls_provider.rs b/xmtp_mls/src/xmtp_openmls_provider.rs index f050eaa19..32df7706e 100644 --- a/xmtp_mls/src/xmtp_openmls_provider.rs +++ b/xmtp_mls/src/xmtp_openmls_provider.rs @@ -1,18 +1,18 @@ use openmls_rust_crypto::RustCrypto; use openmls_traits::OpenMlsProvider; -use crate::storage::in_memory_key_store::InMemoryKeyStore; +use crate::storage::sql_key_store::SqlKeyStore; #[derive(Default, Debug)] pub struct XmtpOpenMlsProvider { crypto: RustCrypto, - key_store: InMemoryKeyStore, + key_store: SqlKeyStore, } impl OpenMlsProvider for XmtpOpenMlsProvider { type CryptoProvider = RustCrypto; type RandProvider = RustCrypto; - type KeyStoreProvider = InMemoryKeyStore; + type KeyStoreProvider = SqlKeyStore; fn crypto(&self) -> &Self::CryptoProvider { &self.crypto From a404c205d190c9650f0fa5da93ae127962b33f78 Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Wed, 25 Oct 2023 16:16:28 -0700 Subject: [PATCH 05/24] Use reference to encrypted store to allow multiple consumers --- xmtp_mls/src/builder.rs | 3 ++- xmtp_mls/src/identity.rs | 7 +++++-- xmtp_mls/src/storage/sql_key_store.rs | 22 +++++++++++++++++----- xmtp_mls/src/xmtp_openmls_provider.rs | 21 +++++++++++++++------ 4 files changed, 39 insertions(+), 14 deletions(-) diff --git a/xmtp_mls/src/builder.rs b/xmtp_mls/src/builder.rs index c9a476ad4..c6425ad96 100644 --- a/xmtp_mls/src/builder.rs +++ b/xmtp_mls/src/builder.rs @@ -108,6 +108,7 @@ where parameter: "api_client", })?; let store = self.store.take().unwrap_or_default(); + let provider = XmtpOpenMlsProvider::new(&store); // Fetch the Identity based upon the identity strategy. let identity = match self.identity_strategy { IdentityStrategy::CachedOnly(_) => { @@ -116,7 +117,7 @@ where } IdentityStrategy::CreateIfNotFound(owner) => { // TODO: persistence/retrieval - Identity::new(CIPHERSUITE, &XmtpOpenMlsProvider::default(), &owner)? + Identity::new(CIPHERSUITE, &provider, &owner)? } #[cfg(test)] IdentityStrategy::ExternalIdentity(a) => a, diff --git a/xmtp_mls/src/identity.rs b/xmtp_mls/src/identity.rs index f0025db74..6fc0ea9ea 100644 --- a/xmtp_mls/src/identity.rs +++ b/xmtp_mls/src/identity.rs @@ -101,7 +101,10 @@ impl Identity { mod tests { use xmtp_cryptography::utils::generate_local_wallet; - use crate::{configuration::CIPHERSUITE, xmtp_openmls_provider::XmtpOpenMlsProvider}; + use crate::{ + configuration::CIPHERSUITE, storage::EncryptedMessageStore, + xmtp_openmls_provider::XmtpOpenMlsProvider, + }; use super::Identity; @@ -109,7 +112,7 @@ mod tests { fn does_not_error() { Identity::new( CIPHERSUITE, - &XmtpOpenMlsProvider::default(), + &XmtpOpenMlsProvider::new(&EncryptedMessageStore::default()), &generate_local_wallet(), ) .unwrap(); diff --git a/xmtp_mls/src/storage/sql_key_store.rs b/xmtp_mls/src/storage/sql_key_store.rs index d0545b57d..8adec0340 100644 --- a/xmtp_mls/src/storage/sql_key_store.rs +++ b/xmtp_mls/src/storage/sql_key_store.rs @@ -5,12 +5,24 @@ use crate::{Delete, Fetch, Store}; use super::{encrypted_store::models::StoredKeyStoreEntry, EncryptedMessageStore, StorageError}; -#[derive(Debug, Default)] -pub struct SqlKeyStore { - store: EncryptedMessageStore, +#[derive(Debug)] +pub struct SqlKeyStore<'a> { + store: &'a EncryptedMessageStore, } -impl OpenMlsKeyStore for SqlKeyStore { +impl Default for SqlKeyStore<'_> { + fn default() -> Self { + unimplemented!() + } +} + +impl<'a> SqlKeyStore<'a> { + pub fn new(store: &'a EncryptedMessageStore) -> Self { + SqlKeyStore { store } + } +} + +impl OpenMlsKeyStore for SqlKeyStore<'_> { /// The error type returned by the [`OpenMlsKeyStore`]. type Error = StorageError; @@ -85,7 +97,7 @@ mod tests { fn store_read_delete() { let db_path = format!("{}.db3", rand_string()); let key_store = SqlKeyStore { - store: EncryptedMessageStore::new( + store: &EncryptedMessageStore::new( StorageOption::Persistent(db_path), EncryptedMessageStore::generate_enc_key(), ) diff --git a/xmtp_mls/src/xmtp_openmls_provider.rs b/xmtp_mls/src/xmtp_openmls_provider.rs index 32df7706e..eab5b3582 100644 --- a/xmtp_mls/src/xmtp_openmls_provider.rs +++ b/xmtp_mls/src/xmtp_openmls_provider.rs @@ -1,18 +1,27 @@ use openmls_rust_crypto::RustCrypto; use openmls_traits::OpenMlsProvider; -use crate::storage::sql_key_store::SqlKeyStore; +use crate::storage::{sql_key_store::SqlKeyStore, EncryptedMessageStore}; -#[derive(Default, Debug)] -pub struct XmtpOpenMlsProvider { +#[derive(Debug)] +pub struct XmtpOpenMlsProvider<'a> { crypto: RustCrypto, - key_store: SqlKeyStore, + key_store: SqlKeyStore<'a>, } -impl OpenMlsProvider for XmtpOpenMlsProvider { +impl<'a> XmtpOpenMlsProvider<'a> { + pub fn new(store: &'a EncryptedMessageStore) -> Self { + Self { + crypto: RustCrypto::default(), + key_store: SqlKeyStore::new(store), + } + } +} + +impl<'a> OpenMlsProvider for XmtpOpenMlsProvider<'a> { type CryptoProvider = RustCrypto; type RandProvider = RustCrypto; - type KeyStoreProvider = SqlKeyStore; + type KeyStoreProvider = SqlKeyStore<'a>; fn crypto(&self) -> &Self::CryptoProvider { &self.crypto From acd42b834144b8990358d69f253207cf48683cae Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Wed, 25 Oct 2023 16:42:16 -0700 Subject: [PATCH 06/24] Tidy identity --- xmtp_mls/src/identity.rs | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/xmtp_mls/src/identity.rs b/xmtp_mls/src/identity.rs index 6fc0ea9ea..a9fae1c70 100644 --- a/xmtp_mls/src/identity.rs +++ b/xmtp_mls/src/identity.rs @@ -37,8 +37,9 @@ pub enum IdentityError { #[derive(serde::Serialize, serde::Deserialize)] pub struct Identity { - pub(crate) credential_with_key: CredentialWithKey, - pub(crate) signer: SignatureKeyPair, + account_address: String, + signer: SignatureKeyPair, + credential: Credential, } impl Identity { @@ -50,7 +51,7 @@ impl Identity { let signature_keys = SignatureKeyPair::new(ciphersuite.signature_algorithm())?; signature_keys.store(provider.key_store())?; - let credential_with_key = Identity::create_credential(&signature_keys, owner)?; + let credential = Identity::create_credential(&signature_keys, owner)?; // The builder automatically stores it in the key store // TODO: Make OpenMLS not delete this once used @@ -61,22 +62,26 @@ impl Identity { }, provider, &signature_keys, - credential_with_key.clone(), + CredentialWithKey { + credential: credential.clone(), + signature_key: signature_keys.to_public_vec().into(), + }, )?; // TODO: persist identity // TODO: upload credential_with_key and last_resort_key_package Ok(Self { - credential_with_key, + account_address: owner.get_address(), signer: signature_keys, + credential, }) } fn create_credential( signature_keys: &SignatureKeyPair, owner: &impl InboxOwner, - ) -> Result { + ) -> Result { // Generate association let assoc_text = AssociationText::Static { blockchain_address: owner.get_address(), @@ -88,12 +93,7 @@ impl Identity { let association_proto: Eip191AssociationProto = association.into(); // Serialize into credential - let credential = - Credential::new(association_proto.encode_to_vec(), CredentialType::Basic).unwrap(); - Ok(CredentialWithKey { - credential, - signature_key: signature_keys.to_public_vec().into(), - }) + Ok(Credential::new(association_proto.encode_to_vec(), CredentialType::Basic).unwrap()) } } From 0a53a11376fc8ef95c9eac5445976415655771f5 Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Wed, 25 Oct 2023 20:06:32 -0700 Subject: [PATCH 07/24] Add identity table, queries, and unit tests --- xmtp/Cargo.toml | 4 +- xmtp_mls/Cargo.toml | 6 + .../down.sql | 1 + .../2023-10-25-234319_create_identity/up.sql | 6 + xmtp_mls/src/lib.rs | 23 ++-- xmtp_mls/src/storage/encrypted_store/mod.rs | 127 +++++++++++++++++- .../src/storage/encrypted_store/models.rs | 24 ++++ .../src/storage/encrypted_store/schema.rs | 14 ++ xmtp_mls/src/storage/sql_key_store.rs | 2 +- 9 files changed, 189 insertions(+), 18 deletions(-) create mode 100644 xmtp_mls/migrations/2023-10-25-234319_create_identity/down.sql create mode 100644 xmtp_mls/migrations/2023-10-25-234319_create_identity/up.sql diff --git a/xmtp/Cargo.toml b/xmtp/Cargo.toml index c449e2103..47f2aea77 100644 --- a/xmtp/Cargo.toml +++ b/xmtp/Cargo.toml @@ -11,7 +11,7 @@ path = "src/bin/update-schema.rs" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [features] -default = ["types"] +default = ["types", "native"] types = [] grpc = ["xmtp_proto/grpc"] native = ["libsqlite3-sys/bundled-sqlcipher-vendored-openssl"] @@ -36,7 +36,7 @@ ethers-core = "2.0.4" prost = { version = "0.11", features = ["prost-derive"] } futures = "0.3.28" base64 = "0.21.1" -tokio = "1.28.1" +tokio = { version = "1.28.1", features = ["rt", "macros"] } anyhow = "1.0.71" [dev-dependencies] diff --git a/xmtp_mls/Cargo.toml b/xmtp_mls/Cargo.toml index 84d629f11..7f850db84 100644 --- a/xmtp_mls/Cargo.toml +++ b/xmtp_mls/Cargo.toml @@ -10,6 +10,12 @@ path = "src/bin/update-schema.rs" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[features] +default = ["types", "native"] +types = [] +grpc = ["xmtp_proto/grpc"] +native = ["libsqlite3-sys/bundled-sqlcipher-vendored-openssl"] + [dependencies] anyhow = "1.0.71" async-trait = "0.1.68" diff --git a/xmtp_mls/migrations/2023-10-25-234319_create_identity/down.sql b/xmtp_mls/migrations/2023-10-25-234319_create_identity/down.sql new file mode 100644 index 000000000..ca9f66bf9 --- /dev/null +++ b/xmtp_mls/migrations/2023-10-25-234319_create_identity/down.sql @@ -0,0 +1 @@ +DROP TABLE identity; diff --git a/xmtp_mls/migrations/2023-10-25-234319_create_identity/up.sql b/xmtp_mls/migrations/2023-10-25-234319_create_identity/up.sql new file mode 100644 index 000000000..5de241cb7 --- /dev/null +++ b/xmtp_mls/migrations/2023-10-25-234319_create_identity/up.sql @@ -0,0 +1,6 @@ +CREATE TABLE identity ( + account_address TEXT NOT NULL, + signature_keypair BLOB NOT NULL, + credential_bytes BLOB NOT NULL, + rowid INTEGER PRIMARY KEY CHECK (rowid = 1) -- There can only be one identity +); diff --git a/xmtp_mls/src/lib.rs b/xmtp_mls/src/lib.rs index 4dd0a6b01..558ec0dc2 100644 --- a/xmtp_mls/src/lib.rs +++ b/xmtp_mls/src/lib.rs @@ -20,20 +20,19 @@ pub trait InboxOwner { } // Inserts a model to the underlying data store -pub trait Store { - fn store(&self, into: &mut I) -> Result<(), StorageError>; +pub trait Store { + fn store(&self, into: &mut StorageConnection) -> Result<(), StorageError>; } -pub trait Fetch { - type Key<'a>; - // Fetches a single instance by key of a model from the underlying data store - #[allow(clippy::needless_lifetimes)] - fn fetch<'a>(&mut self, key: Self::Key<'a>) -> Result, StorageError>; +pub trait Fetch { + type Key; + fn fetch(&mut self, key: Self::Key) -> Result, StorageError>; } -pub trait Delete { - type Key<'a>; - // Fetches a single instance by key of a model from the underlying data store - #[allow(clippy::needless_lifetimes)] - fn delete<'a>(&mut self, key: Self::Key<'a>) -> Result; +pub trait Delete { + type Key; + fn delete(&mut self, key: Self::Key) -> Result; } + +#[cfg(test)] +mod tests {} diff --git a/xmtp_mls/src/storage/encrypted_store/mod.rs b/xmtp_mls/src/storage/encrypted_store/mod.rs index b6cb8c43f..c9367019f 100644 --- a/xmtp_mls/src/storage/encrypted_store/mod.rs +++ b/xmtp_mls/src/storage/encrypted_store/mod.rs @@ -16,7 +16,7 @@ pub mod schema; use crate::{Delete, Fetch, Store}; -use self::{models::StoredKeyStoreEntry, schema::openmls_key_store}; +use self::{models::*, schema::*}; use super::StorageError; use diesel::{ @@ -178,7 +178,7 @@ impl Store for StoredKeyStoreEntry { } impl Fetch for DbConnection { - type Key<'a> = Vec; + type Key = Vec; fn fetch(&mut self, key: Vec) -> Result, StorageError> where { use self::schema::openmls_key_store::dsl::*; Ok(openmls_key_store.find(key).first(self).optional()?) @@ -186,15 +186,136 @@ impl Fetch for DbConnection { } impl Delete for DbConnection { - type Key<'a> = Vec; + type Key = Vec; fn delete(&mut self, key: Vec) -> Result where { use self::schema::openmls_key_store::dsl::*; Ok(diesel::delete(openmls_key_store.filter(key_bytes.eq(key))).execute(self)?) } } +impl Store for StoredIdentity { + fn store(&self, into: &mut DbConnection) -> Result<(), StorageError> { + diesel::insert_into(identity::table) + .values(self) + .execute(into)?; + Ok(()) + } +} + +impl Fetch for DbConnection { + type Key = (); + fn fetch(&mut self, _key: ()) -> Result, StorageError> where { + use self::schema::identity::dsl::*; + Ok(identity.first(self).optional()?) + } +} + #[cfg(test)] mod tests { + use super::{models::*, EncryptedMessageStore, StorageError, StorageOption}; + use crate::{Fetch, Store}; + use rand::{ + distributions::{Alphanumeric, DistString}, + Rng, + }; + use std::boxed::Box; + use std::fs; + + fn rand_string() -> String { + Alphanumeric.sample_string(&mut rand::thread_rng(), 16) + } + + fn rand_vec() -> Vec { + rand::thread_rng().gen::<[u8; 16]>().to_vec() + } + + #[test] + fn ephemeral_store() { + let store = EncryptedMessageStore::new( + StorageOption::Ephemeral, + EncryptedMessageStore::generate_enc_key(), + ) + .unwrap(); + let conn = &mut store.conn().unwrap(); + + let account_address = "address"; + StoredIdentity::new(account_address.to_string(), rand_vec(), rand_vec()) + .store(conn) + .unwrap(); + + let fetched_identity: StoredIdentity = conn.fetch(()).unwrap().unwrap(); + assert_eq!(fetched_identity.account_address, account_address); + } + + #[test] + fn persistent_store() { + let db_path = format!("{}.db3", rand_string()); + { + let store = EncryptedMessageStore::new( + StorageOption::Persistent(db_path.clone()), + EncryptedMessageStore::generate_enc_key(), + ) + .unwrap(); + let conn = &mut store.conn().unwrap(); + + let account_address = "address"; + StoredIdentity::new(account_address.to_string(), rand_vec(), rand_vec()) + .store(conn) + .unwrap(); + + let fetched_identity: StoredIdentity = conn.fetch(()).unwrap().unwrap(); + assert_eq!(fetched_identity.account_address, account_address); + } + + fs::remove_file(db_path).unwrap(); + } + + #[test] + fn mismatched_encryption_key() { + let mut enc_key = [1u8; 32]; + + let db_path = format!("{}.db3", rand_string()); + { + // Setup a persistent store + let store = EncryptedMessageStore::new( + StorageOption::Persistent(db_path.clone()), + enc_key.clone(), + ) + .unwrap(); + + StoredIdentity::new("dummy_address".to_string(), rand_vec(), rand_vec()) + .store(&mut store.conn().unwrap()) + .unwrap(); + } // Drop it + + enc_key[3] = 145; // Alter the enc_key + let res = EncryptedMessageStore::new(StorageOption::Persistent(db_path.clone()), enc_key); + // Ensure it fails + match res.err() { + Some(StorageError::DbInitError(_)) => (), + _ => panic!("Expected a DbInitError"), + } + fs::remove_file(db_path).unwrap(); + } + + #[test] + fn can_only_store_one_identity() { + let store = EncryptedMessageStore::new( + StorageOption::Ephemeral, + EncryptedMessageStore::generate_enc_key(), + ) + .unwrap(); + let conn = &mut store.conn().unwrap(); + + StoredIdentity::new("".to_string(), rand_vec(), rand_vec()) + .store(conn) + .unwrap(); + + let duplicate_insertion = + StoredIdentity::new("".to_string(), rand_vec(), rand_vec()).store(conn); + assert!(duplicate_insertion.is_err()); + } + #[test] fn it_returns_ok_when_given_ok_result() { let result: Result<(), diesel::result::Error> = Ok(()); diff --git a/xmtp_mls/src/storage/encrypted_store/models.rs b/xmtp_mls/src/storage/encrypted_store/models.rs index 49b06e5dc..0489a270b 100644 --- a/xmtp_mls/src/storage/encrypted_store/models.rs +++ b/xmtp_mls/src/storage/encrypted_store/models.rs @@ -9,3 +9,27 @@ pub struct StoredKeyStoreEntry { pub key_bytes: Vec, pub value_bytes: Vec, } + +#[derive(Insertable, Queryable, Debug, Clone)] +#[diesel(table_name = identity)] +pub struct StoredIdentity { + pub account_address: String, + pub signature_keypair: Vec, + pub credential_bytes: Vec, + rowid: Option, +} + +impl StoredIdentity { + pub fn new( + account_address: String, + signature_keypair: Vec, + credential_bytes: Vec, + ) -> Self { + Self { + account_address, + signature_keypair, + credential_bytes, + rowid: None, + } + } +} diff --git a/xmtp_mls/src/storage/encrypted_store/schema.rs b/xmtp_mls/src/storage/encrypted_store/schema.rs index 9de21ee27..bb26d1082 100644 --- a/xmtp_mls/src/storage/encrypted_store/schema.rs +++ b/xmtp_mls/src/storage/encrypted_store/schema.rs @@ -1,8 +1,22 @@ // @generated automatically by Diesel CLI. +diesel::table! { + identity (rowid) { + account_address -> Text, + signature_keypair -> Binary, + credential_bytes -> Binary, + rowid -> Nullable, + } +} + diesel::table! { openmls_key_store (key_bytes) { key_bytes -> Binary, value_bytes -> Binary, } } + +diesel::allow_tables_to_appear_in_same_query!( + identity, + openmls_key_store, +); diff --git a/xmtp_mls/src/storage/sql_key_store.rs b/xmtp_mls/src/storage/sql_key_store.rs index 8adec0340..a30382e10 100644 --- a/xmtp_mls/src/storage/sql_key_store.rs +++ b/xmtp_mls/src/storage/sql_key_store.rs @@ -55,7 +55,7 @@ impl OpenMlsKeyStore for SqlKeyStore<'_> { error!("Failed to fetch key: {:?}", e); return None; } - let entry_option = fetch_result.unwrap(); + let entry_option: Option = fetch_result.unwrap(); if entry_option.is_none() { debug!("No entry to read for key {:?}", k); return None; From 4fb105969782ed0638af871cb6466851f3cc0333 Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Wed, 25 Oct 2023 21:34:38 -0700 Subject: [PATCH 08/24] Persist and retrieve identity from storage inside builder --- xmtp_mls/src/builder.rs | 50 +++++++++++-------- xmtp_mls/src/identity.rs | 7 ++- .../src/storage/encrypted_store/models.rs | 26 ++++++++++ xmtp_mls/src/storage/errors.rs | 2 + xmtp_mls/src/storage/mod.rs | 5 +- xmtp_mls/src/storage/serialization.rs | 17 +++++++ xmtp_mls/src/storage/sql_key_store.rs | 10 ++-- 7 files changed, 88 insertions(+), 29 deletions(-) create mode 100644 xmtp_mls/src/storage/serialization.rs diff --git a/xmtp_mls/src/builder.rs b/xmtp_mls/src/builder.rs index c6425ad96..e1181f53a 100644 --- a/xmtp_mls/src/builder.rs +++ b/xmtp_mls/src/builder.rs @@ -1,13 +1,13 @@ use crate::configuration::CIPHERSUITE; +use crate::storage::StoredIdentity; use crate::xmtp_openmls_provider::XmtpOpenMlsProvider; -use crate::StorageError; use crate::{ client::{Client, Network}, identity::{Identity, IdentityError}, storage::EncryptedMessageStore, - types::Address, InboxOwner, }; +use crate::{Fetch, StorageError}; use thiserror::Error; use xmtp_proto::api_client::XmtpApiClient; @@ -37,17 +37,11 @@ pub enum ClientBuilderError { pub enum IdentityStrategy { CreateIfNotFound(Owner), - CachedOnly(Address), + CachedOnly, #[cfg(test)] ExternalIdentity(Identity), } -impl From for IdentityStrategy { - fn from(value: String) -> Self { - IdentityStrategy::CachedOnly(value) - } -} - impl From for IdentityStrategy where Owner: InboxOwner, @@ -107,22 +101,36 @@ where .ok_or(ClientBuilderError::MissingParameter { parameter: "api_client", })?; + let network = self.network; let store = self.store.take().unwrap_or_default(); let provider = XmtpOpenMlsProvider::new(&store); - // Fetch the Identity based upon the identity strategy. - let identity = match self.identity_strategy { - IdentityStrategy::CachedOnly(_) => { - // TODO: persistence/retrieval - unimplemented!() - } - IdentityStrategy::CreateIfNotFound(owner) => { - // TODO: persistence/retrieval - Identity::new(CIPHERSUITE, &provider, &owner)? + let identity = self.initialize_identity(&store, &provider)?; + Ok(Client::new(api_client, network, identity, store)) + } + + fn initialize_identity( + self, + store: &EncryptedMessageStore, + provider: &XmtpOpenMlsProvider, + ) -> Result { + let conn = &mut store.conn()?; + let identity_option: Option = conn.fetch(())?.map(|i: StoredIdentity| i.into()); + match self.identity_strategy { + IdentityStrategy::CachedOnly => { + identity_option.ok_or(ClientBuilderError::RequiredIdentityNotFound) } + IdentityStrategy::CreateIfNotFound(owner) => match identity_option { + Some(identity) => { + if identity.account_address != owner.get_address() { + return Err(ClientBuilderError::StoredIdentityMismatch); + } + Ok(identity) + } + None => Ok(Identity::new(CIPHERSUITE, &provider, &owner)?), + }, #[cfg(test)] - IdentityStrategy::ExternalIdentity(a) => a, - }; - Ok(Client::new(api_client, self.network, identity, store)) + IdentityStrategy::ExternalIdentity(identity) => Ok(identity), + } } } diff --git a/xmtp_mls/src/identity.rs b/xmtp_mls/src/identity.rs index a9fae1c70..098e4e9c6 100644 --- a/xmtp_mls/src/identity.rs +++ b/xmtp_mls/src/identity.rs @@ -35,11 +35,10 @@ pub enum IdentityError { KeyPackageGenerationError(#[from] KeyPackageNewError), } -#[derive(serde::Serialize, serde::Deserialize)] pub struct Identity { - account_address: String, - signer: SignatureKeyPair, - credential: Credential, + pub(crate) account_address: String, + pub(crate) signer: SignatureKeyPair, + pub(crate) credential: Credential, } impl Identity { diff --git a/xmtp_mls/src/storage/encrypted_store/models.rs b/xmtp_mls/src/storage/encrypted_store/models.rs index 0489a270b..b5a877d01 100644 --- a/xmtp_mls/src/storage/encrypted_store/models.rs +++ b/xmtp_mls/src/storage/encrypted_store/models.rs @@ -1,5 +1,10 @@ use diesel::prelude::*; +use crate::{ + identity::Identity, + storage::serialization::{db_deserialize, db_serialize}, +}; + use super::schema::*; #[derive(Insertable, Queryable, Debug, Clone)] @@ -33,3 +38,24 @@ impl StoredIdentity { } } } + +impl From for StoredIdentity { + fn from(identity: Identity) -> Self { + StoredIdentity { + account_address: identity.account_address, + signature_keypair: db_serialize(&identity.signer).unwrap(), + credential_bytes: db_serialize(&identity.credential).unwrap(), + rowid: None, + } + } +} + +impl Into for StoredIdentity { + fn into(self) -> Identity { + Identity { + account_address: self.account_address, + signer: db_deserialize(&self.signature_keypair).unwrap(), + credential: db_deserialize(&self.credential_bytes).unwrap(), + } + } +} diff --git a/xmtp_mls/src/storage/errors.rs b/xmtp_mls/src/storage/errors.rs index be4292d90..c44ab4e24 100644 --- a/xmtp_mls/src/storage/errors.rs +++ b/xmtp_mls/src/storage/errors.rs @@ -14,6 +14,8 @@ pub enum StorageError { Store(String), #[error("serialization error")] SerializationError, + #[error("deserialization error")] + DeserializationError, #[error("unknown storage error: {0}")] Unknown(String), } diff --git a/xmtp_mls/src/storage/mod.rs b/xmtp_mls/src/storage/mod.rs index 6710adb90..676545425 100644 --- a/xmtp_mls/src/storage/mod.rs +++ b/xmtp_mls/src/storage/mod.rs @@ -1,6 +1,9 @@ mod encrypted_store; mod errors; +mod serialization; pub mod sql_key_store; -pub use encrypted_store::{DbConnection, EncryptedMessageStore, EncryptionKey, StorageOption}; +pub use encrypted_store::{ + models::*, DbConnection, EncryptedMessageStore, EncryptionKey, StorageOption, +}; pub use errors::StorageError; diff --git a/xmtp_mls/src/storage/serialization.rs b/xmtp_mls/src/storage/serialization.rs new file mode 100644 index 000000000..911a024be --- /dev/null +++ b/xmtp_mls/src/storage/serialization.rs @@ -0,0 +1,17 @@ +use serde::Serialize; + +use super::StorageError; + +pub fn db_serialize(value: &T) -> Result, StorageError> +where + T: ?Sized + Serialize, +{ + serde_json::to_vec(value).map_err(|_| StorageError::SerializationError) +} + +pub fn db_deserialize(bytes: &[u8]) -> Result +where + T: serde::de::DeserializeOwned, +{ + serde_json::from_slice(&bytes).map_err(|_| StorageError::DeserializationError) +} diff --git a/xmtp_mls/src/storage/sql_key_store.rs b/xmtp_mls/src/storage/sql_key_store.rs index a30382e10..1539d9b08 100644 --- a/xmtp_mls/src/storage/sql_key_store.rs +++ b/xmtp_mls/src/storage/sql_key_store.rs @@ -3,7 +3,11 @@ use openmls_traits::key_store::{MlsEntity, OpenMlsKeyStore}; use crate::{Delete, Fetch, Store}; -use super::{encrypted_store::models::StoredKeyStoreEntry, EncryptedMessageStore, StorageError}; +use super::{ + encrypted_store::models::StoredKeyStoreEntry, + serialization::{db_deserialize, db_serialize}, + EncryptedMessageStore, StorageError, +}; #[derive(Debug)] pub struct SqlKeyStore<'a> { @@ -33,7 +37,7 @@ impl OpenMlsKeyStore for SqlKeyStore<'_> { fn store(&self, k: &[u8], v: &V) -> Result<(), Self::Error> { let entry = StoredKeyStoreEntry { key_bytes: k.to_vec(), - value_bytes: serde_json::to_vec(v).map_err(|_| StorageError::SerializationError)?, + value_bytes: db_serialize(v)?, }; entry.store(&mut self.store.conn()?)?; Ok(()) @@ -60,7 +64,7 @@ impl OpenMlsKeyStore for SqlKeyStore<'_> { debug!("No entry to read for key {:?}", k); return None; } - serde_json::from_slice(&entry_option.unwrap().value_bytes).ok() + db_deserialize(&entry_option.unwrap().value_bytes).ok() } /// Delete a value stored for ID `k`. From 598b5ea86b4ba359285476815031a36a6c560878 Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Wed, 25 Oct 2023 21:42:17 -0700 Subject: [PATCH 09/24] Tidy up names and types --- .../2023-10-25-234319_create_identity/up.sql | 2 +- xmtp_mls/src/identity.rs | 14 ++++++++------ xmtp_mls/src/storage/encrypted_store/models.rs | 10 +++++----- xmtp_mls/src/storage/encrypted_store/schema.rs | 2 +- 4 files changed, 15 insertions(+), 13 deletions(-) diff --git a/xmtp_mls/migrations/2023-10-25-234319_create_identity/up.sql b/xmtp_mls/migrations/2023-10-25-234319_create_identity/up.sql index 5de241cb7..9934ee26b 100644 --- a/xmtp_mls/migrations/2023-10-25-234319_create_identity/up.sql +++ b/xmtp_mls/migrations/2023-10-25-234319_create_identity/up.sql @@ -1,6 +1,6 @@ CREATE TABLE identity ( account_address TEXT NOT NULL, - signature_keypair BLOB NOT NULL, + installation_keys BLOB NOT NULL, credential_bytes BLOB NOT NULL, rowid INTEGER PRIMARY KEY CHECK (rowid = 1) -- There can only be one identity ); diff --git a/xmtp_mls/src/identity.rs b/xmtp_mls/src/identity.rs index 098e4e9c6..59199b859 100644 --- a/xmtp_mls/src/identity.rs +++ b/xmtp_mls/src/identity.rs @@ -16,6 +16,7 @@ use xmtp_cryptography::signature::SignatureError; use crate::{ association::{AssociationError, AssociationText, Eip191Association}, storage::StorageError, + types::Address, xmtp_openmls_provider::XmtpOpenMlsProvider, InboxOwner, }; @@ -36,8 +37,8 @@ pub enum IdentityError { } pub struct Identity { - pub(crate) account_address: String, - pub(crate) signer: SignatureKeyPair, + pub(crate) account_address: Address, + pub(crate) installation_keys: SignatureKeyPair, pub(crate) credential: Credential, } @@ -72,22 +73,23 @@ impl Identity { Ok(Self { account_address: owner.get_address(), - signer: signature_keys, + installation_keys: signature_keys, credential, }) } fn create_credential( - signature_keys: &SignatureKeyPair, + installation_keys: &SignatureKeyPair, owner: &impl InboxOwner, ) -> Result { // Generate association let assoc_text = AssociationText::Static { blockchain_address: owner.get_address(), - installation_public_key: signature_keys.to_public_vec(), + installation_public_key: installation_keys.to_public_vec(), }; let signature = owner.sign(&assoc_text.text())?; - let association = Eip191Association::new(signature_keys.public(), assoc_text, signature)?; + let association = + Eip191Association::new(installation_keys.public(), assoc_text, signature)?; // TODO wrap in a Credential proto to allow flexibility for different association types let association_proto: Eip191AssociationProto = association.into(); diff --git a/xmtp_mls/src/storage/encrypted_store/models.rs b/xmtp_mls/src/storage/encrypted_store/models.rs index b5a877d01..06e42e15e 100644 --- a/xmtp_mls/src/storage/encrypted_store/models.rs +++ b/xmtp_mls/src/storage/encrypted_store/models.rs @@ -19,7 +19,7 @@ pub struct StoredKeyStoreEntry { #[diesel(table_name = identity)] pub struct StoredIdentity { pub account_address: String, - pub signature_keypair: Vec, + pub installation_keys: Vec, pub credential_bytes: Vec, rowid: Option, } @@ -27,12 +27,12 @@ pub struct StoredIdentity { impl StoredIdentity { pub fn new( account_address: String, - signature_keypair: Vec, + installation_keys: Vec, credential_bytes: Vec, ) -> Self { Self { account_address, - signature_keypair, + installation_keys, credential_bytes, rowid: None, } @@ -43,7 +43,7 @@ impl From for StoredIdentity { fn from(identity: Identity) -> Self { StoredIdentity { account_address: identity.account_address, - signature_keypair: db_serialize(&identity.signer).unwrap(), + installation_keys: db_serialize(&identity.installation_keys).unwrap(), credential_bytes: db_serialize(&identity.credential).unwrap(), rowid: None, } @@ -54,7 +54,7 @@ impl Into for StoredIdentity { fn into(self) -> Identity { Identity { account_address: self.account_address, - signer: db_deserialize(&self.signature_keypair).unwrap(), + installation_keys: db_deserialize(&self.installation_keys).unwrap(), credential: db_deserialize(&self.credential_bytes).unwrap(), } } diff --git a/xmtp_mls/src/storage/encrypted_store/schema.rs b/xmtp_mls/src/storage/encrypted_store/schema.rs index bb26d1082..d291b1378 100644 --- a/xmtp_mls/src/storage/encrypted_store/schema.rs +++ b/xmtp_mls/src/storage/encrypted_store/schema.rs @@ -3,7 +3,7 @@ diesel::table! { identity (rowid) { account_address -> Text, - signature_keypair -> Binary, + installation_keys -> Binary, credential_bytes -> Binary, rowid -> Nullable, } From 9dad3614ee32cb42669ee4b43ac974f6107cb501 Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Fri, 27 Oct 2023 14:53:09 -0700 Subject: [PATCH 10/24] Update Diesel --- Cargo.lock | 36 +++++++++++++++++++++++------------- xmtp_mls/Cargo.toml | 4 ++-- 2 files changed, 25 insertions(+), 15 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1b98f71d5..35d68f072 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1063,38 +1063,48 @@ dependencies = [ [[package]] name = "diesel" -version = "2.0.4" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72eb77396836a4505da85bae0712fa324b74acfe1876d7c2f7e694ef3d0ee373" +checksum = "2268a214a6f118fce1838edba3d1561cf0e78d8de785475957a580a7f8c69d33" dependencies = [ "diesel_derives", "libsqlite3-sys", "r2d2", + "time 0.3.21", ] [[package]] name = "diesel_derives" -version = "2.0.2" +version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ad74fdcf086be3d4fdd142f67937678fe60ed431c3b2f08599e7687269410c4" +checksum = "ef8337737574f55a468005a83499da720f20c65586241ffea339db9ecdfd2b44" dependencies = [ - "proc-macro-error", + "diesel_table_macro_syntax", "proc-macro2", "quote", - "syn 1.0.109", + "syn 2.0.38", ] [[package]] name = "diesel_migrations" -version = "2.0.0" +version = "2.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9ae22beef5e9d6fab9225ddb073c1c6c1a7a6ded5019d5da11d1e5c5adc34e2" +checksum = "6036b3f0120c5961381b570ee20a02432d7e2d27ea60de9578799cf9156914ac" dependencies = [ "diesel", "migrations_internals", "migrations_macros", ] +[[package]] +name = "diesel_table_macro_syntax" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc5557efc453706fed5e4fa85006fe9817c224c3f480a34c7e5959fd700921c5" +dependencies = [ + "syn 2.0.38", +] + [[package]] name = "diff" version = "0.1.13" @@ -2950,19 +2960,19 @@ dependencies = [ [[package]] name = "migrations_internals" -version = "2.0.0" +version = "2.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c493c09323068c01e54c685f7da41a9ccf9219735c3766fbfd6099806ea08fbc" +checksum = "0f23f71580015254b020e856feac3df5878c2c7a8812297edd6c0a485ac9dada" dependencies = [ "serde", - "toml 0.5.11", + "toml 0.7.4", ] [[package]] name = "migrations_macros" -version = "2.0.0" +version = "2.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a8ff27a350511de30cdabb77147501c36ef02e0451d957abea2f30caffb2b58" +checksum = "cce3325ac70e67bbab5bd837a31cae01f1a6db64e0e744a33cb03a543469ef08" dependencies = [ "migrations_internals", "proc-macro2", diff --git a/xmtp_mls/Cargo.toml b/xmtp_mls/Cargo.toml index f68cc5213..05cd8b56d 100644 --- a/xmtp_mls/Cargo.toml +++ b/xmtp_mls/Cargo.toml @@ -19,8 +19,8 @@ native = ["libsqlite3-sys/bundled-sqlcipher-vendored-openssl"] [dependencies] anyhow = "1.0.71" async-trait = "0.1.68" -diesel = { version = "2.0.4", features = ["sqlite", "r2d2", "returning_clauses_for_sqlite_3_35"] } -diesel_migrations = { version = "2.0.0", features = ["sqlite"] } +diesel = { version = "2.1.3", features = ["sqlite", "r2d2", "returning_clauses_for_sqlite_3_35"] } +diesel_migrations = { version = "2.1.0", features = ["sqlite"] } ethers = "2.0.4" ethers-core = "2.0.4" futures = "0.3.28" From b1f411b985ab4eff85f7815afe7ea71df1a14fda Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Fri, 27 Oct 2023 14:53:29 -0700 Subject: [PATCH 11/24] Hard code ciphersuite --- xmtp_mls/src/builder.rs | 2 +- xmtp_mls/src/identity.rs | 32 ++++++++++++++++++++++++-------- 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/xmtp_mls/src/builder.rs b/xmtp_mls/src/builder.rs index 2b1750a9c..0b9917f98 100644 --- a/xmtp_mls/src/builder.rs +++ b/xmtp_mls/src/builder.rs @@ -126,7 +126,7 @@ where } Ok(identity) } - None => Ok(Identity::new(CIPHERSUITE, &provider, &owner)?), + None => Ok(Identity::new(&provider, &owner)?), }, #[cfg(test)] IdentityStrategy::ExternalIdentity(identity) => Ok(identity), diff --git a/xmtp_mls/src/identity.rs b/xmtp_mls/src/identity.rs index 59199b859..3b901fda8 100644 --- a/xmtp_mls/src/identity.rs +++ b/xmtp_mls/src/identity.rs @@ -15,6 +15,7 @@ use xmtp_cryptography::signature::SignatureError; use crate::{ association::{AssociationError, AssociationText, Eip191Association}, + configuration::CIPHERSUITE, storage::StorageError, types::Address, xmtp_openmls_provider::XmtpOpenMlsProvider, @@ -44,11 +45,10 @@ pub struct Identity { impl Identity { pub(crate) fn new( - ciphersuite: Ciphersuite, provider: &XmtpOpenMlsProvider, owner: &impl InboxOwner, ) -> Result { - let signature_keys = SignatureKeyPair::new(ciphersuite.signature_algorithm())?; + let signature_keys = SignatureKeyPair::new(CIPHERSUITE.signature_algorithm())?; signature_keys.store(provider.key_store())?; let credential = Identity::create_credential(&signature_keys, owner)?; @@ -57,7 +57,7 @@ impl Identity { // TODO: Make OpenMLS not delete this once used let _last_resort_key_package = KeyPackage::builder().build( CryptoConfig { - ciphersuite, + ciphersuite: CIPHERSUITE, version: ProtocolVersion::default(), }, provider, @@ -78,6 +78,26 @@ impl Identity { }) } + pub(crate) fn new_key_package( + &self, + provider: &XmtpOpenMlsProvider, + ) -> Result { + let kp = KeyPackage::builder().build( + CryptoConfig { + ciphersuite: CIPHERSUITE, + version: ProtocolVersion::default(), + }, + provider, + &self.installation_keys, + CredentialWithKey { + credential: self.credential.clone(), + signature_key: self.installation_keys.to_public_vec().into(), + }, + )?; + + Ok(kp) + } + fn create_credential( installation_keys: &SignatureKeyPair, owner: &impl InboxOwner, @@ -102,17 +122,13 @@ impl Identity { mod tests { use xmtp_cryptography::utils::generate_local_wallet; - use crate::{ - configuration::CIPHERSUITE, storage::EncryptedMessageStore, - xmtp_openmls_provider::XmtpOpenMlsProvider, - }; + use crate::{storage::EncryptedMessageStore, xmtp_openmls_provider::XmtpOpenMlsProvider}; use super::Identity; #[test] fn does_not_error() { Identity::new( - CIPHERSUITE, &XmtpOpenMlsProvider::new(&EncryptedMessageStore::default()), &generate_local_wallet(), ) From 091afdba4ba8fd46019a770bec5fe4d4157febd2 Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Fri, 27 Oct 2023 16:43:20 -0700 Subject: [PATCH 12/24] Add tests and fix bugs --- xmtp_mls/src/builder.rs | 94 +++++++++++++++++-- xmtp_mls/src/client.rs | 14 ++- xmtp_mls/src/identity.rs | 19 ++-- xmtp_mls/src/mock_xmtp_api_client.rs | 2 + .../src/storage/encrypted_store/models.rs | 6 +- 5 files changed, 113 insertions(+), 22 deletions(-) diff --git a/xmtp_mls/src/builder.rs b/xmtp_mls/src/builder.rs index e1181f53a..a627c92c2 100644 --- a/xmtp_mls/src/builder.rs +++ b/xmtp_mls/src/builder.rs @@ -8,6 +8,10 @@ use crate::{ InboxOwner, }; use crate::{Fetch, StorageError}; +#[cfg(not(test))] +use log::debug; +#[cfg(test)] +use std::println as debug; use thiserror::Error; use xmtp_proto::api_client::XmtpApiClient; @@ -113,8 +117,9 @@ where store: &EncryptedMessageStore, provider: &XmtpOpenMlsProvider, ) -> Result { - let conn = &mut store.conn()?; - let identity_option: Option = conn.fetch(())?.map(|i: StoredIdentity| i.into()); + let identity_option: Option = + store.conn()?.fetch(())?.map(|i: StoredIdentity| i.into()); + debug!("Existing identity in store: {:?}", identity_option); match self.identity_strategy { IdentityStrategy::CachedOnly => { identity_option.ok_or(ClientBuilderError::RequiredIdentityNotFound) @@ -126,7 +131,7 @@ where } Ok(identity) } - None => Ok(Identity::new(CIPHERSUITE, &provider, &owner)?), + None => Ok(Identity::new(store, CIPHERSUITE, &provider, &owner)?), }, #[cfg(test)] IdentityStrategy::ExternalIdentity(identity) => Ok(identity), @@ -137,18 +142,87 @@ where #[cfg(test)] mod tests { - use ethers::signers::LocalWallet; + use ethers::signers::{LocalWallet, Signer}; + use tempfile::TempPath; use xmtp_cryptography::utils::generate_local_wallet; - use crate::mock_xmtp_api_client::MockXmtpApiClient; + use crate::{ + mock_xmtp_api_client::MockXmtpApiClient, + storage::{EncryptedMessageStore, StorageOption}, + }; use super::ClientBuilder; - impl ClientBuilder { - pub fn new_test() -> Self { - let wallet = generate_local_wallet(); + #[test] + fn builder_test() { + let wallet = generate_local_wallet(); + let address = wallet.address(); + let client = ClientBuilder::new(wallet.into()) + .api_client(MockXmtpApiClient::default()) + .build() + .unwrap(); + assert!(client.account_address() == format!("{address:#020x}")); + assert!(!client.installation_public_key().is_empty()); + } - Self::new(wallet.into()) - } + #[test] + fn identity_persistence_test() { + let tmpdb = TempPath::from_path("./db.db3") + .to_str() + .unwrap() + .to_string(); + let wallet = generate_local_wallet(); + + // Generate a new Wallet + Store + let store_a = + EncryptedMessageStore::new_unencrypted(StorageOption::Persistent(tmpdb.clone())) + .unwrap(); + + let client_a = ClientBuilder::new(wallet.clone().into()) + .store(store_a) + .api_client(MockXmtpApiClient::default()) + .build() + .unwrap(); + let keybytes_a = client_a.installation_public_key(); + drop(client_a); + + // Reload the existing store and wallet + let store_b = + EncryptedMessageStore::new_unencrypted(StorageOption::Persistent(tmpdb.clone())) + .unwrap(); + + let client_b = ClientBuilder::new(wallet.into()) + .store(store_b) + .api_client(MockXmtpApiClient::default()) + .build() + .unwrap(); + let keybytes_b = client_b.installation_public_key(); + drop(client_b); + + // Ensure the persistence was used to store the generated keys + assert_eq!(keybytes_a, keybytes_b); + + // Create a new wallet and store + let store_c = + EncryptedMessageStore::new_unencrypted(StorageOption::Persistent(tmpdb.clone())) + .unwrap(); + + ClientBuilder::::new(generate_local_wallet().into()) + .store(store_c) + .build() + .expect_err("Testing expected mismatch error"); + + // Use cached only strategy + let store_d = + EncryptedMessageStore::new_unencrypted(StorageOption::Persistent(tmpdb.clone())) + .unwrap(); + let client_d = ClientBuilder::::new( + crate::builder::IdentityStrategy::CachedOnly, + ) + .store(store_d) + .api_client(MockXmtpApiClient::default()) + .build() + .unwrap(); + assert_eq!(client_d.installation_public_key(), keybytes_a); } } diff --git a/xmtp_mls/src/client.rs b/xmtp_mls/src/client.rs index 79bf4ada1..0c363aba3 100644 --- a/xmtp_mls/src/client.rs +++ b/xmtp_mls/src/client.rs @@ -3,6 +3,7 @@ use thiserror::Error; use crate::{ identity::Identity, storage::{EncryptedMessageStore, StorageError}, + types::Address, }; #[derive(Clone, Copy, Default, Debug)] @@ -39,10 +40,11 @@ impl From<&str> for ClientError { } } +#[derive(Debug)] pub struct Client { pub api_client: ApiClient, pub(crate) _network: Network, - pub(crate) _identity: Identity, + pub(crate) identity: Identity, pub store: EncryptedMessageStore, // Temporarily exposed outside crate for CLI client } @@ -56,8 +58,16 @@ impl Client { Self { api_client, _network: network, - _identity: identity, + identity, store, } } + + pub fn account_address(&self) -> Address { + self.identity.account_address.clone() + } + + pub fn installation_public_key(&self) -> Vec { + self.identity.installation_keys.to_public_vec() + } } diff --git a/xmtp_mls/src/identity.rs b/xmtp_mls/src/identity.rs index 59199b859..317b191da 100644 --- a/xmtp_mls/src/identity.rs +++ b/xmtp_mls/src/identity.rs @@ -15,10 +15,10 @@ use xmtp_cryptography::signature::SignatureError; use crate::{ association::{AssociationError, AssociationText, Eip191Association}, - storage::StorageError, + storage::{EncryptedMessageStore, StorageError, StoredIdentity}, types::Address, xmtp_openmls_provider::XmtpOpenMlsProvider, - InboxOwner, + InboxOwner, Store, }; use xmtp_proto::xmtp::v3::message_contents::Eip191Association as Eip191AssociationProto; @@ -36,6 +36,7 @@ pub enum IdentityError { KeyPackageGenerationError(#[from] KeyPackageNewError), } +#[derive(Debug)] pub struct Identity { pub(crate) account_address: Address, pub(crate) installation_keys: SignatureKeyPair, @@ -44,6 +45,7 @@ pub struct Identity { impl Identity { pub(crate) fn new( + store: &EncryptedMessageStore, ciphersuite: Ciphersuite, provider: &XmtpOpenMlsProvider, owner: &impl InboxOwner, @@ -68,14 +70,16 @@ impl Identity { }, )?; - // TODO: persist identity - // TODO: upload credential_with_key and last_resort_key_package - - Ok(Self { + let identity = Self { account_address: owner.get_address(), installation_keys: signature_keys, credential, - }) + }; + StoredIdentity::from(&identity).store(&mut store.conn()?)?; + + // TODO: upload credential_with_key and last_resort_key_package + + Ok(identity) } fn create_credential( @@ -112,6 +116,7 @@ mod tests { #[test] fn does_not_error() { Identity::new( + &EncryptedMessageStore::default(), CIPHERSUITE, &XmtpOpenMlsProvider::new(&EncryptedMessageStore::default()), &generate_local_wallet(), diff --git a/xmtp_mls/src/mock_xmtp_api_client.rs b/xmtp_mls/src/mock_xmtp_api_client.rs index 2d49b630e..b70be389d 100644 --- a/xmtp_mls/src/mock_xmtp_api_client.rs +++ b/xmtp_mls/src/mock_xmtp_api_client.rs @@ -19,11 +19,13 @@ impl XmtpApiSubscription for MockXmtpApiSubscription { fn close_stream(&mut self) {} } +#[derive(Debug)] struct InnerMockXmtpApiClient { pub messages: HashMap>, pub app_version: String, } +#[derive(Debug)] pub struct MockXmtpApiClient { inner_client: Arc>, } diff --git a/xmtp_mls/src/storage/encrypted_store/models.rs b/xmtp_mls/src/storage/encrypted_store/models.rs index 06e42e15e..3c75379e8 100644 --- a/xmtp_mls/src/storage/encrypted_store/models.rs +++ b/xmtp_mls/src/storage/encrypted_store/models.rs @@ -39,10 +39,10 @@ impl StoredIdentity { } } -impl From for StoredIdentity { - fn from(identity: Identity) -> Self { +impl From<&Identity> for StoredIdentity { + fn from(identity: &Identity) -> Self { StoredIdentity { - account_address: identity.account_address, + account_address: identity.account_address.clone(), installation_keys: db_serialize(&identity.installation_keys).unwrap(), credential_bytes: db_serialize(&identity.credential).unwrap(), rowid: None, From 18c6aa4e66cba97cc2bfc422ef6f74af6cb34455 Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Fri, 27 Oct 2023 16:50:33 -0700 Subject: [PATCH 13/24] Fix lints --- xmtp_mls/src/builder.rs | 2 +- xmtp_mls/src/storage/encrypted_store/models.rs | 10 +++++----- xmtp_mls/src/storage/serialization.rs | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/xmtp_mls/src/builder.rs b/xmtp_mls/src/builder.rs index a627c92c2..22fd95227 100644 --- a/xmtp_mls/src/builder.rs +++ b/xmtp_mls/src/builder.rs @@ -131,7 +131,7 @@ where } Ok(identity) } - None => Ok(Identity::new(store, CIPHERSUITE, &provider, &owner)?), + None => Ok(Identity::new(store, CIPHERSUITE, provider, &owner)?), }, #[cfg(test)] IdentityStrategy::ExternalIdentity(identity) => Ok(identity), diff --git a/xmtp_mls/src/storage/encrypted_store/models.rs b/xmtp_mls/src/storage/encrypted_store/models.rs index 3c75379e8..b1b618515 100644 --- a/xmtp_mls/src/storage/encrypted_store/models.rs +++ b/xmtp_mls/src/storage/encrypted_store/models.rs @@ -50,12 +50,12 @@ impl From<&Identity> for StoredIdentity { } } -impl Into for StoredIdentity { - fn into(self) -> Identity { +impl From for Identity { + fn from(identity: StoredIdentity) -> Self { Identity { - account_address: self.account_address, - installation_keys: db_deserialize(&self.installation_keys).unwrap(), - credential: db_deserialize(&self.credential_bytes).unwrap(), + account_address: identity.account_address, + installation_keys: db_deserialize(&identity.installation_keys).unwrap(), + credential: db_deserialize(&identity.credential_bytes).unwrap(), } } } diff --git a/xmtp_mls/src/storage/serialization.rs b/xmtp_mls/src/storage/serialization.rs index 911a024be..f3a6333f1 100644 --- a/xmtp_mls/src/storage/serialization.rs +++ b/xmtp_mls/src/storage/serialization.rs @@ -13,5 +13,5 @@ pub fn db_deserialize(bytes: &[u8]) -> Result where T: serde::de::DeserializeOwned, { - serde_json::from_slice(&bytes).map_err(|_| StorageError::DeserializationError) + serde_json::from_slice(bytes).map_err(|_| StorageError::DeserializationError) } From dae1c7498a8f2dfa262c2fd969ef8450763e689d Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Fri, 27 Oct 2023 16:51:12 -0700 Subject: [PATCH 14/24] Update xmtp_mls/src/storage/encrypted_store/mod.rs Co-authored-by: Andrew Plaza --- xmtp_mls/src/storage/encrypted_store/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xmtp_mls/src/storage/encrypted_store/mod.rs b/xmtp_mls/src/storage/encrypted_store/mod.rs index c9367019f..cca453e8a 100644 --- a/xmtp_mls/src/storage/encrypted_store/mod.rs +++ b/xmtp_mls/src/storage/encrypted_store/mod.rs @@ -291,7 +291,7 @@ mod tests { enc_key[3] = 145; // Alter the enc_key let res = EncryptedMessageStore::new(StorageOption::Persistent(db_path.clone()), enc_key); // Ensure it fails - match res.err() { + assert!(matches!(res.err(), Some(StorageError::DbInitError(_)), "Expected DbInitError"); Some(StorageError::DbInitError(_)) => (), _ => panic!("Expected a DbInitError"), } From 77aa94b8f1650313600d992b1ca1c85f6a22cf09 Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Fri, 27 Oct 2023 16:53:03 -0700 Subject: [PATCH 15/24] Dont use the word Error in error names --- xmtp_mls/src/storage/encrypted_store/mod.rs | 20 +++++++++----------- xmtp_mls/src/storage/errors.rs | 14 ++++++-------- xmtp_mls/src/storage/serialization.rs | 4 ++-- 3 files changed, 17 insertions(+), 21 deletions(-) diff --git a/xmtp_mls/src/storage/encrypted_store/mod.rs b/xmtp_mls/src/storage/encrypted_store/mod.rs index cca453e8a..8a298b391 100644 --- a/xmtp_mls/src/storage/encrypted_store/mod.rs +++ b/xmtp_mls/src/storage/encrypted_store/mod.rs @@ -88,11 +88,11 @@ impl EncryptedMessageStore { StorageOption::Ephemeral => Pool::builder() .max_size(1) .build(ConnectionManager::::new(":memory:")) - .map_err(|e| StorageError::DbInitError(e.to_string()))?, + .map_err(|e| StorageError::DbInit(e.to_string()))?, StorageOption::Persistent(ref path) => Pool::builder() .max_size(10) .build(ConnectionManager::::new(path)) - .map_err(|e| StorageError::DbInitError(e.to_string()))?, + .map_err(|e| StorageError::DbInit(e.to_string()))?, }; // // Setup SqlCipherKey @@ -116,7 +116,7 @@ impl EncryptedMessageStore { let conn = &mut self.conn()?; conn.run_pending_migrations(MIGRATIONS) - .map_err(|e| StorageError::DbInitError(e.to_string()))?; + .map_err(|e| StorageError::DbInit(e.to_string()))?; Ok(()) } @@ -127,7 +127,7 @@ impl EncryptedMessageStore { let conn = self .pool .get() - .map_err(|e| StorageError::PoolError(e.to_string()))?; + .map_err(|e| StorageError::Pool(e.to_string()))?; Ok(conn) } @@ -136,9 +136,7 @@ impl EncryptedMessageStore { pool: Pool>, encryption_key: &[u8; 32], ) -> Result<(), StorageError> { - let conn = &mut pool - .get() - .map_err(|e| StorageError::PoolError(e.to_string()))?; + let conn = &mut pool.get().map_err(|e| StorageError::Pool(e.to_string()))?; conn.batch_execute(&format!( "PRAGMA key = \"x'{}'\";", @@ -291,10 +289,10 @@ mod tests { enc_key[3] = 145; // Alter the enc_key let res = EncryptedMessageStore::new(StorageOption::Persistent(db_path.clone()), enc_key); // Ensure it fails - assert!(matches!(res.err(), Some(StorageError::DbInitError(_)), "Expected DbInitError"); - Some(StorageError::DbInitError(_)) => (), - _ => panic!("Expected a DbInitError"), - } + assert!( + matches!(res.err(), Some(StorageError::DbInit(_))), + "Expected DbInitError" + ); fs::remove_file(db_path).unwrap(); } diff --git a/xmtp_mls/src/storage/errors.rs b/xmtp_mls/src/storage/errors.rs index c44ab4e24..13d431f43 100644 --- a/xmtp_mls/src/storage/errors.rs +++ b/xmtp_mls/src/storage/errors.rs @@ -3,19 +3,17 @@ use thiserror::Error; #[derive(Debug, Error, PartialEq)] pub enum StorageError { #[error("Diesel connection error")] - DieselConnectError(#[from] diesel::ConnectionError), + DieselConnect(#[from] diesel::ConnectionError), #[error("Diesel result error: {0}")] - DieselResultError(#[from] diesel::result::Error), + DieselResult(#[from] diesel::result::Error), #[error("Pool error {0}")] - PoolError(String), + Pool(String), #[error("Either incorrect encryptionkey or file is not a db {0}")] - DbInitError(String), + DbInit(String), #[error("Store Error")] Store(String), #[error("serialization error")] - SerializationError, + Serialization, #[error("deserialization error")] - DeserializationError, - #[error("unknown storage error: {0}")] - Unknown(String), + Deserialization, } diff --git a/xmtp_mls/src/storage/serialization.rs b/xmtp_mls/src/storage/serialization.rs index f3a6333f1..272143c3e 100644 --- a/xmtp_mls/src/storage/serialization.rs +++ b/xmtp_mls/src/storage/serialization.rs @@ -6,12 +6,12 @@ pub fn db_serialize(value: &T) -> Result, StorageError> where T: ?Sized + Serialize, { - serde_json::to_vec(value).map_err(|_| StorageError::SerializationError) + serde_json::to_vec(value).map_err(|_| StorageError::Serialization) } pub fn db_deserialize(bytes: &[u8]) -> Result where T: serde::de::DeserializeOwned, { - serde_json::from_slice(bytes).map_err(|_| StorageError::DeserializationError) + serde_json::from_slice(bytes).map_err(|_| StorageError::Deserialization) } From 5278d13c955bc89814a01dbc482ca0dea0767f0d Mon Sep 17 00:00:00 2001 From: Richard Hua Date: Fri, 27 Oct 2023 17:00:59 -0700 Subject: [PATCH 16/24] Move identity initialization to identity strategy --- xmtp_mls/src/builder.rs | 61 +++++++++++++++++++++++------------------ 1 file changed, 34 insertions(+), 27 deletions(-) diff --git a/xmtp_mls/src/builder.rs b/xmtp_mls/src/builder.rs index 22fd95227..f4847b7e3 100644 --- a/xmtp_mls/src/builder.rs +++ b/xmtp_mls/src/builder.rs @@ -46,6 +46,37 @@ pub enum IdentityStrategy { ExternalIdentity(Identity), } +impl IdentityStrategy +where + Owner: InboxOwner, +{ + fn initialize_identity( + self, + store: &EncryptedMessageStore, + provider: &XmtpOpenMlsProvider, + ) -> Result { + let identity_option: Option = + store.conn()?.fetch(())?.map(|i: StoredIdentity| i.into()); + debug!("Existing identity in store: {:?}", identity_option); + match self { + IdentityStrategy::CachedOnly => { + identity_option.ok_or(ClientBuilderError::RequiredIdentityNotFound) + } + IdentityStrategy::CreateIfNotFound(owner) => match identity_option { + Some(identity) => { + if identity.account_address != owner.get_address() { + return Err(ClientBuilderError::StoredIdentityMismatch); + } + Ok(identity) + } + None => Ok(Identity::new(store, CIPHERSUITE, provider, &owner)?), + }, + #[cfg(test)] + IdentityStrategy::ExternalIdentity(identity) => Ok(identity), + } + } +} + impl From for IdentityStrategy where Owner: InboxOwner, @@ -108,35 +139,11 @@ where let network = self.network; let store = self.store.take().unwrap_or_default(); let provider = XmtpOpenMlsProvider::new(&store); - let identity = self.initialize_identity(&store, &provider)?; + let identity = self + .identity_strategy + .initialize_identity(&store, &provider)?; Ok(Client::new(api_client, network, identity, store)) } - - fn initialize_identity( - self, - store: &EncryptedMessageStore, - provider: &XmtpOpenMlsProvider, - ) -> Result { - let identity_option: Option = - store.conn()?.fetch(())?.map(|i: StoredIdentity| i.into()); - debug!("Existing identity in store: {:?}", identity_option); - match self.identity_strategy { - IdentityStrategy::CachedOnly => { - identity_option.ok_or(ClientBuilderError::RequiredIdentityNotFound) - } - IdentityStrategy::CreateIfNotFound(owner) => match identity_option { - Some(identity) => { - if identity.account_address != owner.get_address() { - return Err(ClientBuilderError::StoredIdentityMismatch); - } - Ok(identity) - } - None => Ok(Identity::new(store, CIPHERSUITE, provider, &owner)?), - }, - #[cfg(test)] - IdentityStrategy::ExternalIdentity(identity) => Ok(identity), - } - } } #[cfg(test)] From 04f5dd979fde5ad4ccb1be0b0ebd2171aebf8580 Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Sat, 28 Oct 2023 13:11:45 -0700 Subject: [PATCH 17/24] Add create key packages methods --- Cargo.lock | 1 + xmtp_mls/Cargo.toml | 1 + xmtp_mls/src/api_client_wrapper.rs | 6 +-- xmtp_mls/src/builder.rs | 10 ++--- xmtp_mls/src/client.rs | 62 ++++++++++++++++++++++++++++++ xmtp_mls/src/configuration.rs | 2 + 6 files changed, 74 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 35d68f072..030ea398e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6402,6 +6402,7 @@ dependencies = [ "serde_json", "tempfile", "thiserror", + "tls_codec", "tokio", "toml 0.7.4", "xmtp_api_grpc", diff --git a/xmtp_mls/Cargo.toml b/xmtp_mls/Cargo.toml index 05cd8b56d..94fd4e2cc 100644 --- a/xmtp_mls/Cargo.toml +++ b/xmtp_mls/Cargo.toml @@ -40,6 +40,7 @@ tokio = { version = "1.28.1", features = ["macros"] } toml = "0.7.4" xmtp_cryptography = { path = "../xmtp_cryptography"} xmtp_proto = { path = "../xmtp_proto", features = ["proto_full"] } +tls_codec = "0.3.0" [dev-dependencies] rand = "0.8.5" diff --git a/xmtp_mls/src/api_client_wrapper.rs b/xmtp_mls/src/api_client_wrapper.rs index 074a0545a..58cbab775 100644 --- a/xmtp_mls/src/api_client_wrapper.rs +++ b/xmtp_mls/src/api_client_wrapper.rs @@ -38,7 +38,7 @@ where pub async fn register_installation( &self, - last_resort_key_package: &[u8], + last_resort_key_package: Vec, ) -> Result, ApiError> { let res = self .api_client @@ -52,13 +52,13 @@ where Ok(res.installation_id) } - pub async fn upload_key_packages(&self, key_packages: Vec<&[u8]>) -> Result<(), ApiError> { + pub async fn upload_key_packages(&self, key_packages: Vec>) -> Result<(), ApiError> { self.api_client .upload_key_packages(UploadKeyPackagesRequest { key_packages: key_packages .into_iter() .map(|kp| KeyPackageUpload { - key_package_tls_serialized: kp.to_vec(), + key_package_tls_serialized: kp, }) .collect(), }) diff --git a/xmtp_mls/src/builder.rs b/xmtp_mls/src/builder.rs index 0b9917f98..f78a22acf 100644 --- a/xmtp_mls/src/builder.rs +++ b/xmtp_mls/src/builder.rs @@ -59,7 +59,7 @@ pub struct ClientBuilder { identity_strategy: IdentityStrategy, } -impl ClientBuilder +impl<'a, ApiClient, Owner> ClientBuilder where ApiClient: XmtpApiClient + XmtpMlsClient, Owner: InboxOwner, @@ -102,10 +102,10 @@ where parameter: "api_client", })?; let network = self.network; - let store = self.store.take().unwrap_or_default(); - let provider = XmtpOpenMlsProvider::new(&store); + let store = &self.store.take().unwrap_or_default(); + let provider = XmtpOpenMlsProvider::new(store); let identity = self.initialize_identity(&store, &provider)?; - Ok(Client::new(api_client, network, identity, store)) + Ok(Client::new(api_client, network, identity, store.clone())) } fn initialize_identity( @@ -158,7 +158,7 @@ mod tests { async fn test_mls() { let client = ClientBuilder::new_test().await.build().unwrap(); - let result = client.api_client.register_installation(&[1, 2, 3]).await; + let result = client.api_client.register_installation(vec![1, 2, 3]).await; assert!(result.is_err()); let error_string = result.err().unwrap().to_string(); diff --git a/xmtp_mls/src/client.rs b/xmtp_mls/src/client.rs index d1f3b7cbe..5145eb478 100644 --- a/xmtp_mls/src/client.rs +++ b/xmtp_mls/src/client.rs @@ -1,10 +1,14 @@ +use openmls::prelude::TlsSerializeTrait; use thiserror::Error; +use tls_codec::Error as TlsSerializationError; use xmtp_proto::api_client::{XmtpApiClient, XmtpMlsClient}; use crate::{ api_client_wrapper::ApiClientWrapper, + configuration::KEY_PACKAGE_TOP_UP_AMOUNT, identity::Identity, storage::{EncryptedMessageStore, StorageError}, + xmtp_openmls_provider::XmtpOpenMlsProvider, }; #[derive(Clone, Copy, Default, Debug)] @@ -25,6 +29,10 @@ pub enum ClientError { Ddd(#[from] diesel::result::Error), #[error("Query failed: {0}")] QueryError(#[from] xmtp_proto::api_client::Error), + #[error("identity error: {0}")] + Identity(#[from] crate::identity::IdentityError), + #[error("serialization error: {0}")] + Serialization(#[from] TlsSerializationError), #[error("generic:{0}")] Generic(String), } @@ -65,4 +73,58 @@ where store, } } + + // TODO: Remove this and figure out the correct lifetimes to allow long lived provider + fn mls_provider(&self) -> XmtpOpenMlsProvider { + XmtpOpenMlsProvider::new(&self.store) + } + + pub(crate) async fn register_identity(&self) -> Result<(), ClientError> { + // TODO: Mark key package as last_resort in creation + let last_resort_kp = self._identity.new_key_package(&self.mls_provider())?; + let last_resort_kp_bytes = last_resort_kp.tls_serialize_detached()?; + + self.api_client + .register_installation(last_resort_kp_bytes) + .await?; + + Ok(()) + } + + pub async fn top_up_key_packages(&self) -> Result<(), ClientError> { + let key_packages: Result>, ClientError> = (0..KEY_PACKAGE_TOP_UP_AMOUNT) + .into_iter() + .map(|_| -> Result, ClientError> { + let kp = self._identity.new_key_package(&self.mls_provider())?; + let kp_bytes = kp.tls_serialize_detached()?; + + Ok(kp_bytes) + }) + .collect(); + + self.api_client.upload_key_packages(key_packages?).await?; + + Ok(()) + } +} + +#[cfg(test)] +mod test { + use crate::builder::ClientBuilder; + + #[tokio::test] + async fn test_register_installation() { + let client = ClientBuilder::new_test().await.build().unwrap(); + let res = client.register_identity().await; + assert!(res.is_ok()); + } + + #[tokio::test] + async fn top_up_key_packages() { + let client = ClientBuilder::new_test().await.build().unwrap(); + client.register_identity().await.unwrap(); + + let res = client.top_up_key_packages().await; + assert!(res.is_ok()); + } } diff --git a/xmtp_mls/src/configuration.rs b/xmtp_mls/src/configuration.rs index 2c5ce4b7d..761b62938 100644 --- a/xmtp_mls/src/configuration.rs +++ b/xmtp_mls/src/configuration.rs @@ -3,3 +3,5 @@ use openmls_traits::types::Ciphersuite; // TODO confirm ciphersuite choice pub const CIPHERSUITE: Ciphersuite = Ciphersuite::MLS_128_DHKEMX25519_CHACHA20POLY1305_SHA256_Ed25519; + +pub const KEY_PACKAGE_TOP_UP_AMOUNT: u16 = 100; From 34d42f0e7e47dbe829564646bb0e49264dc5c533 Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Sun, 29 Oct 2023 10:19:32 -0700 Subject: [PATCH 18/24] Add verified key package handling --- Cargo.lock | 97 ++++++++ xmtp_mls/Cargo.toml | 1 + xmtp_mls/src/api_client_wrapper.rs | 329 +++++++++++++++++++++++++-- xmtp_mls/src/builder.rs | 11 - xmtp_mls/src/client.rs | 125 +++++++++- xmtp_mls/src/configuration.rs | 3 + xmtp_mls/src/identity.rs | 5 +- xmtp_mls/src/lib.rs | 1 + xmtp_mls/src/verified_key_package.rs | 80 +++++++ 9 files changed, 613 insertions(+), 39 deletions(-) create mode 100644 xmtp_mls/src/verified_key_package.rs diff --git a/Cargo.lock b/Cargo.lock index 030ea398e..6ed71ff32 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1111,6 +1111,12 @@ version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56254986775e3233ffa9c4d7d3faaf6d36a2c09d30b20687e9f88bc8bafc16c8" +[[package]] +name = "difflib" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6184e33543162437515c2e2b48714794e37845ec9851711914eec9d308f6ebe8" + [[package]] name = "digest" version = "0.8.1" @@ -1203,6 +1209,12 @@ dependencies = [ "winapi 0.3.9", ] +[[package]] +name = "downcast" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1435fa1053d8b2fbbe9be7e97eca7f33d37b28409959813daefc1446a14247f1" + [[package]] name = "dunce" version = "1.0.4" @@ -1934,6 +1946,15 @@ dependencies = [ "miniz_oxide", ] +[[package]] +name = "float-cmp" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98de4bbd547a563b716d8dfa9aad1cb19bfab00f4fa09a6a4ed21dbcf44ce9c4" +dependencies = [ + "num-traits", +] + [[package]] name = "fnv" version = "1.0.7" @@ -1964,6 +1985,12 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fragile" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c2141d6d6c8512188a7891b4b01590a45f6dac67afb4f255c4124dbb86d4eaa" + [[package]] name = "fs2" version = "0.4.3" @@ -3071,6 +3098,33 @@ dependencies = [ "xmtp_proto", ] +[[package]] +name = "mockall" +version = "0.11.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c84490118f2ee2d74570d114f3d0493cbf02790df303d2707606c3e14e07c96" +dependencies = [ + "cfg-if 1.0.0", + "downcast", + "fragile", + "lazy_static", + "mockall_derive", + "predicates", + "predicates-tree", +] + +[[package]] +name = "mockall_derive" +version = "0.11.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22ce75669015c4f47b289fd4d4f56e894e4c96003ffdf3ac51313126f94c6cbb" +dependencies = [ + "cfg-if 1.0.0", + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "multimap" version = "0.8.3" @@ -3122,6 +3176,12 @@ dependencies = [ "version_check", ] +[[package]] +name = "normalize-line-endings" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61807f77802ff30975e01f4f071c8ba10c022052f98b3294119f3e615d13e5be" + [[package]] name = "num-bigint" version = "0.4.4" @@ -3879,6 +3939,36 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "925383efa346730478fb4838dbe9137d2a47675ad789c546d150a6e1dd4ab31c" +[[package]] +name = "predicates" +version = "2.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59230a63c37f3e18569bdb90e4a89cbf5bf8b06fea0b84e65ea10cc4df47addd" +dependencies = [ + "difflib", + "float-cmp", + "itertools 0.10.5", + "normalize-line-endings", + "predicates-core", + "regex", +] + +[[package]] +name = "predicates-core" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b794032607612e7abeb4db69adb4e33590fa6cf1149e95fd7cb00e634b92f174" + +[[package]] +name = "predicates-tree" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368ba315fb8c5052ab692e68a0eefec6ec57b23a36959c14496f0b0df2c0cecf" +dependencies = [ + "predicates-core", + "termtree", +] + [[package]] name = "prettyplease" version = "0.1.25" @@ -5242,6 +5332,12 @@ dependencies = [ "phf_codegen", ] +[[package]] +name = "termtree" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3369f5ac52d5eb6ab48c6b4ffdc8efbcad6b89c765749064ba298f2c68a16a76" + [[package]] name = "thiserror" version = "1.0.50" @@ -6392,6 +6488,7 @@ dependencies = [ "hex", "libsqlite3-sys", "log", + "mockall", "openmls 0.5.0 (git+https://github.com/xmtp/openmls)", "openmls_basic_credential 0.2.0 (git+https://github.com/xmtp/openmls)", "openmls_rust_crypto 0.2.0 (git+https://github.com/xmtp/openmls)", diff --git a/xmtp_mls/Cargo.toml b/xmtp_mls/Cargo.toml index 94fd4e2cc..ca2ccb97b 100644 --- a/xmtp_mls/Cargo.toml +++ b/xmtp_mls/Cargo.toml @@ -43,6 +43,7 @@ xmtp_proto = { path = "../xmtp_proto", features = ["proto_full"] } tls_codec = "0.3.0" [dev-dependencies] +mockall = "0.11.4" rand = "0.8.5" tempfile = "3.5.0" xmtp_api_grpc = { path = "../xmtp_api_grpc" } diff --git a/xmtp_mls/src/api_client_wrapper.rs b/xmtp_mls/src/api_client_wrapper.rs index f43fc74e3..e1b4d7745 100644 --- a/xmtp_mls/src/api_client_wrapper.rs +++ b/xmtp_mls/src/api_client_wrapper.rs @@ -2,13 +2,17 @@ use std::collections::HashMap; use xmtp_proto::{ api_client::{Error as ApiError, ErrorKind, XmtpApiClient, XmtpMlsClient}, - xmtp::mls::message_contents::{ - group_message::{Version as GroupMessageVersion, V1 as GroupMessageV1}, - welcome_message::{Version as WelcomeMessageVersion, V1 as WelcomeMessageV1}, - WelcomeMessage as WelcomeMessageProto, + xmtp::{ + message_api::v3::GetIdentityUpdatesRequest, + mls::message_contents::{ + group_message::{Version as GroupMessageVersion, V1 as GroupMessageV1}, + welcome_message::{Version as WelcomeMessageVersion, V1 as WelcomeMessageV1}, + WelcomeMessage as WelcomeMessageProto, + }, }, xmtp::{ message_api::v3::{ + get_identity_updates_response::update::Kind as UpdateKind, publish_welcomes_request::WelcomeMessageRequest, ConsumeKeyPackagesRequest, KeyPackageUpload, PublishToGroupRequest, PublishWelcomesRequest, RegisterInstallationRequest, UploadKeyPackagesRequest, @@ -17,13 +21,6 @@ use xmtp_proto::{ }, }; -pub struct WelcomeMessage { - pub(crate) ciphertext: Vec, - pub(crate) installation_id: Vec, -} - -type KeyPackageMap = HashMap, Vec>; - #[derive(Debug)] pub struct ApiClientWrapper { api_client: ApiClient, @@ -31,7 +28,7 @@ pub struct ApiClientWrapper { impl ApiClientWrapper where - ApiClient: XmtpMlsClient + XmtpApiClient, + ApiClient: XmtpMlsClient, { pub fn new(api_client: ApiClient) -> Self { Self { api_client } @@ -70,12 +67,12 @@ where pub async fn consume_key_packages( &self, - installation_ids: Vec<&[u8]>, + installation_ids: Vec>, ) -> Result { let res = self .api_client .consume_key_packages(ConsumeKeyPackagesRequest { - installation_ids: installation_ids.iter().map(|id| id.to_vec()).collect(), + installation_ids: installation_ids.clone(), }) .await?; @@ -124,6 +121,61 @@ where Ok(()) } + pub async fn get_identity_updates( + &self, + start_time_ns: u64, + wallet_addresses: Vec, + ) -> Result { + let result = self + .api_client + .get_identity_updates(GetIdentityUpdatesRequest { + start_time_ns, + wallet_addresses: wallet_addresses.clone(), + }) + .await?; + + if result.updates.len() != wallet_addresses.len() { + println!("mismatched number of results"); + return Err(ApiError::new(ErrorKind::MlsError)); + } + + let mapping: IdentityUpdatesMap = result + .updates + .into_iter() + .enumerate() + .map(|(idx, update)| { + ( + wallet_addresses[idx].clone(), + update + .updates + .into_iter() + .map(|update| match update.kind { + Some(UpdateKind::NewInstallation(new_installation)) => { + IdentityUpdate::NewInstallation(NewInstallation { + timestamp_ns: update.timestamp_ns, + installation_id: new_installation.installation_id, + credential_bytes: new_installation.credential_identity, + }) + } + Some(UpdateKind::RevokedInstallation(revoke_installation)) => { + IdentityUpdate::RevokeInstallation(RevokeInstallation { + timestamp_ns: update.timestamp_ns, + installation_id: revoke_installation.installation_id, + }) + } + None => { + println!("no update kind"); + IdentityUpdate::Invalid + } + }) + .collect(), + ) + }) + .collect(); + + Ok(mapping) + } + pub async fn publish_to_group(&self, group_messages: Vec<&[u8]>) -> Result<(), ApiError> { let to_send: Vec = group_messages .iter() @@ -141,3 +193,252 @@ where Ok(()) } } + +#[derive(Debug, PartialEq)] +pub struct WelcomeMessage { + pub(crate) ciphertext: Vec, + pub(crate) installation_id: Vec, +} + +#[derive(Debug, PartialEq)] +pub struct NewInstallation { + pub installation_id: Vec, + pub credential_bytes: Vec, + pub timestamp_ns: u64, +} + +#[derive(Debug, PartialEq)] +pub struct RevokeInstallation { + pub installation_id: Vec, // TODO: Add proof of revocation + pub timestamp_ns: u64, +} + +#[derive(Debug, PartialEq)] +pub enum IdentityUpdate { + NewInstallation(NewInstallation), + RevokeInstallation(RevokeInstallation), + Invalid, +} + +type KeyPackageMap = HashMap, Vec>; + +type IdentityUpdatesMap = HashMap>; + +#[cfg(test)] +mod tests { + use super::ApiClientWrapper; + use mockall::mock; + use xmtp_proto::api_client::{Error, XmtpApiClient, XmtpApiSubscription, XmtpMlsClient}; + use xmtp_proto::xmtp::message_api::v3::consume_key_packages_response::KeyPackage; + use xmtp_proto::xmtp::message_api::v3::get_identity_updates_response::update::Kind as UpdateKind; + use xmtp_proto::xmtp::message_api::v3::get_identity_updates_response::{ + NewInstallationUpdate, Update, WalletUpdates, + }; + use xmtp_proto::xmtp::message_api::v3::{ + ConsumeKeyPackagesRequest, ConsumeKeyPackagesResponse, GetIdentityUpdatesRequest, + GetIdentityUpdatesResponse, PublishToGroupRequest, PublishWelcomesRequest, + RegisterInstallationRequest, RegisterInstallationResponse, UploadKeyPackagesRequest, + }; + + use xmtp_proto::xmtp::message_api::v1::{ + BatchQueryRequest, BatchQueryResponse, Envelope, PublishRequest, PublishResponse, + QueryRequest, QueryResponse, SubscribeRequest, + }; + + use async_trait::async_trait; + + mock! { + pub Subscription {} + + impl XmtpApiSubscription for Subscription { + fn is_closed(&self) -> bool; + fn get_messages(&self) -> Vec; + fn close_stream(&mut self); + } + } + + // Create a mock XmtpClient for testing the client wrapper + mock! { + pub ApiClient {} + + #[async_trait] + impl XmtpMlsClient for ApiClient { + async fn register_installation( + &self, + request: RegisterInstallationRequest, + ) -> Result; + async fn upload_key_packages(&self, request: UploadKeyPackagesRequest) -> Result<(), Error>; + async fn consume_key_packages( + &self, + request: ConsumeKeyPackagesRequest, + ) -> Result; + async fn publish_to_group(&self, request: PublishToGroupRequest) -> Result<(), Error>; + async fn publish_welcomes(&self, request: PublishWelcomesRequest) -> Result<(), Error>; + async fn get_identity_updates( + &self, + request: GetIdentityUpdatesRequest, + ) -> Result; + } + + #[async_trait] + impl XmtpApiClient for ApiClient { + // Need to set an associated type and don't currently need streaming + // Can figure out a mocked stream type later + type Subscription = MockSubscription; + + fn set_app_version(&mut self, version: String); + + async fn publish( + &self, + token: String, + request: PublishRequest, + ) -> Result; + + async fn subscribe(&self, request: SubscribeRequest) -> Result<::Subscription, Error>; + + async fn query(&self, request: QueryRequest) -> Result; + + async fn batch_query(&self, request: BatchQueryRequest) -> Result; + } + + + } + + #[tokio::test] + async fn test_register_installation() { + let mut mock_api = MockApiClient::new(); + mock_api.expect_register_installation().returning(move |_| { + Ok(RegisterInstallationResponse { + installation_id: vec![1, 2, 3], + }) + }); + let wrapper = ApiClientWrapper::new(mock_api); + let result = wrapper.register_installation(vec![2, 3, 4]).await.unwrap(); + assert_eq!(result, vec![1, 2, 3]); + } + + #[tokio::test] + async fn test_upload_key_packages() { + let mut mock_api = MockApiClient::new(); + let key_package = vec![1, 2, 3]; + // key_package gets moved below but needs to be used for assertions later + let key_package_clone = key_package.clone(); + mock_api + .expect_upload_key_packages() + .withf(move |req| { + req.key_packages[0] + .key_package_tls_serialized + .eq(&key_package) + }) + .returning(move |_| Ok(())); + let wrapper = ApiClientWrapper::new(mock_api); + let result = wrapper.upload_key_packages(vec![key_package_clone]).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_consume_key_packages() { + let mut mock_api = MockApiClient::new(); + let installation_ids: Vec> = vec![vec![1, 2, 3], vec![4, 5, 6]]; + mock_api.expect_consume_key_packages().returning(move |_| { + Ok(ConsumeKeyPackagesResponse { + key_packages: vec![ + KeyPackage { + key_package_tls_serialized: vec![7, 8, 9], + }, + KeyPackage { + key_package_tls_serialized: vec![10, 11, 12], + }, + ], + }) + }); + let wrapper = ApiClientWrapper::new(mock_api); + let result = wrapper + .consume_key_packages(installation_ids.clone()) + .await + .unwrap(); + assert_eq!(result.len(), 2); + + for (k, v) in result { + if k.eq(&installation_ids[0]) { + assert_eq!(v, vec![7, 8, 9]); + } else { + assert_eq!(v, vec![10, 11, 12]); + } + } + } + + #[tokio::test] + async fn test_get_identity_updates() { + let mut mock_api = MockApiClient::new(); + let start_time_ns = 12; + let wallet_addresses = vec!["wallet1".to_string(), "wallet2".to_string()]; + let wallet_addresses_clone = wallet_addresses.clone(); + mock_api + .expect_get_identity_updates() + .withf(move |req| { + req.start_time_ns.eq(&start_time_ns) && req.wallet_addresses.eq(&wallet_addresses) + }) + .returning(move |_| { + Ok(GetIdentityUpdatesResponse { + updates: { + vec![ + WalletUpdates { + updates: vec![Update { + timestamp_ns: 1, + kind: Some(UpdateKind::NewInstallation( + NewInstallationUpdate { + installation_id: vec![1, 2, 3], + credential_identity: vec![4, 5, 6], + }, + )), + }], + }, + WalletUpdates { + updates: vec![Update { + timestamp_ns: 2, + kind: Some(UpdateKind::NewInstallation( + NewInstallationUpdate { + installation_id: vec![7, 8, 9], + credential_identity: vec![10, 11, 12], + }, + )), + }], + }, + ] + }, + }) + }); + + let wrapper = ApiClientWrapper::new(mock_api); + let result = wrapper + .get_identity_updates(start_time_ns, wallet_addresses_clone.clone()) + .await + .unwrap(); + assert_eq!(result.len(), 2); + + for (k, v) in result { + if k.eq(&wallet_addresses_clone[0]) { + assert_eq!(v.len(), 1); + assert_eq!( + v[0], + super::IdentityUpdate::NewInstallation(super::NewInstallation { + installation_id: vec![1, 2, 3], + credential_bytes: vec![4, 5, 6], + timestamp_ns: 1, + }) + ); + } else { + assert_eq!(v.len(), 1); + assert_eq!( + v[0], + super::IdentityUpdate::NewInstallation(super::NewInstallation { + installation_id: vec![7, 8, 9], + credential_bytes: vec![10, 11, 12], + timestamp_ns: 2, + }) + ); + } + } + } +} diff --git a/xmtp_mls/src/builder.rs b/xmtp_mls/src/builder.rs index 01f172d52..453c178a2 100644 --- a/xmtp_mls/src/builder.rs +++ b/xmtp_mls/src/builder.rs @@ -1,4 +1,3 @@ -use crate::configuration::CIPHERSUITE; use crate::storage::StoredIdentity; use crate::xmtp_openmls_provider::XmtpOpenMlsProvider; use crate::{ @@ -189,16 +188,6 @@ mod tests { assert!(!client.installation_public_key().is_empty()); } - #[tokio::test] - async fn test_mls() { - let client = ClientBuilder::new_test_client(generate_local_wallet().into()).await; - let result = client.api_client.register_installation(vec![1, 2, 3]).await; - - assert!(result.is_err()); - let error_string = result.err().unwrap().to_string(); - assert!(error_string.contains("invalid identity")); - } - #[tokio::test] async fn identity_persistence_test() { let tmpdb = TempPath::from_path("./db.db3") diff --git a/xmtp_mls/src/client.rs b/xmtp_mls/src/client.rs index 2920b30de..2841df981 100644 --- a/xmtp_mls/src/client.rs +++ b/xmtp_mls/src/client.rs @@ -1,14 +1,17 @@ +use std::collections::HashSet; + use openmls::prelude::TlsSerializeTrait; use thiserror::Error; use tls_codec::Error as TlsSerializationError; use xmtp_proto::api_client::{XmtpApiClient, XmtpMlsClient}; use crate::{ - api_client_wrapper::ApiClientWrapper, + api_client_wrapper::{ApiClientWrapper, IdentityUpdate}, configuration::KEY_PACKAGE_TOP_UP_AMOUNT, identity::Identity, storage::{EncryptedMessageStore, StorageError}, types::Address, + verified_key_package::{KeyPackageVerificationError, VerifiedKeyPackage}, xmtp_openmls_provider::XmtpOpenMlsProvider, }; @@ -34,6 +37,8 @@ pub enum ClientError { Identity(#[from] crate::identity::IdentityError), #[error("serialization error: {0}")] Serialization(#[from] TlsSerializationError), + #[error("key package verification: {0}")] + KeyPackageVerification(#[from] KeyPackageVerificationError), #[error("generic:{0}")] Generic(String), } @@ -81,7 +86,7 @@ where XmtpOpenMlsProvider::new(&self.store) } - pub(crate) async fn register_identity(&self) -> Result<(), ClientError> { + pub async fn register_identity(&self) -> Result<(), ClientError> { // TODO: Mark key package as last_resort in creation let last_resort_kp = self.identity.new_key_package(&self.mls_provider())?; let last_resort_kp_bytes = last_resort_kp.tls_serialize_detached()?; @@ -109,6 +114,67 @@ where Ok(()) } + async fn get_all_active_installation_ids( + &self, + wallet_addresses: Vec, + ) -> Result>, ClientError> { + let update_mapping = self + .api_client + .get_identity_updates(0, wallet_addresses) + .await?; + + let mut installation_ids: Vec> = vec![]; + + for (_, updates) in update_mapping { + let mut tmp: HashSet> = HashSet::new(); + for update in updates { + match update { + IdentityUpdate::Invalid => {} + IdentityUpdate::NewInstallation(new_installation) => { + // TODO: Validate credential + tmp.insert(new_installation.installation_id); + } + IdentityUpdate::RevokeInstallation(revoke_installation) => { + tmp.remove(&revoke_installation.installation_id); + } + } + } + installation_ids.extend(tmp); + } + + Ok(installation_ids) + } + + // Get a flat list of one key package per installation for all the wallet addresses provided. + // Revoked installations will be omitted from the list + pub async fn get_key_packages_for_wallet_addresses( + &self, + wallet_addresses: Vec, + ) -> Result, ClientError> { + let installation_ids = self + .get_all_active_installation_ids(wallet_addresses) + .await?; + + let key_package_results = self + .api_client + .consume_key_packages(installation_ids) + .await?; + + let mut out: Vec = vec![]; + let mls_provider = self.mls_provider(); + + for kp_bytes in key_package_results.values() { + // Do we actually want to return errors here or should we just log them and move on? + // These key packages are validated by the node, so all should be valid + out.push(VerifiedKeyPackage::from_bytes( + &mls_provider, + kp_bytes.as_slice(), + )?) + } + + Ok(out) + } + pub fn account_address(&self) -> Address { self.identity.account_address.clone() } @@ -122,23 +188,62 @@ where mod tests { use xmtp_cryptography::utils::generate_local_wallet; - use crate::builder::ClientBuilder; + use crate::{builder::ClientBuilder, InboxOwner}; + + #[tokio::test] + async fn test_mls_error() { + let client = ClientBuilder::new_test_client(generate_local_wallet().into()).await; + let result = client.api_client.register_installation(vec![1, 2, 3]).await; + + assert!(result.is_err()); + let error_string = result.err().unwrap().to_string(); + assert!(error_string.contains("invalid identity")); + } #[tokio::test] async fn test_register_installation() { let wallet = generate_local_wallet(); - let client = ClientBuilder::new_test_client(wallet.into()).await; - let res = client.register_identity().await; - assert!(res.is_ok()); + let client = ClientBuilder::new_test_client(wallet.clone().into()).await; + client.register_identity().await.unwrap(); + + // Make sure the installation is actually on the network + let installation_ids = client + .get_all_active_installation_ids(vec![wallet.get_address()]) + .await + .unwrap(); + assert_eq!(installation_ids.len(), 1); } #[tokio::test] - async fn top_up_key_packages() { + async fn test_top_up_key_packages() { let wallet = generate_local_wallet(); - let client = ClientBuilder::new_test_client(wallet.into()).await; + let wallet_address = wallet.get_address(); + + println!("Wallet address {:?}", wallet_address); + let client = ClientBuilder::new_test_client(wallet.clone().into()).await; client.register_identity().await.unwrap(); + client.top_up_key_packages().await.unwrap(); + + let key_packages = client + .get_key_packages_for_wallet_addresses(vec![wallet_address.clone()]) + .await + .unwrap(); + + assert_eq!(key_packages.len(), 1); + + let key_package = key_packages.first().unwrap(); + assert_eq!(key_package.wallet_address, wallet_address); + + let key_packages_2 = client + .get_key_packages_for_wallet_addresses(vec![wallet_address.clone()]) + .await + .unwrap(); + + assert_eq!(key_packages_2.len(), 1); - let res = client.top_up_key_packages().await; - assert!(res.is_ok()); + // Ensure we got back different key packages + let key_package_2 = key_packages_2.first().unwrap(); + assert_eq!(key_package_2.wallet_address, wallet_address); + assert!(!(key_package_2.eq(&key_package))); } } diff --git a/xmtp_mls/src/configuration.rs b/xmtp_mls/src/configuration.rs index 761b62938..577acc965 100644 --- a/xmtp_mls/src/configuration.rs +++ b/xmtp_mls/src/configuration.rs @@ -1,7 +1,10 @@ +use openmls::versions::ProtocolVersion; use openmls_traits::types::Ciphersuite; // TODO confirm ciphersuite choice pub const CIPHERSUITE: Ciphersuite = Ciphersuite::MLS_128_DHKEMX25519_CHACHA20POLY1305_SHA256_Ed25519; +pub const MLS_PROTOCOL_VERSION: ProtocolVersion = ProtocolVersion::Mls10; + pub const KEY_PACKAGE_TOP_UP_AMOUNT: u16 = 100; diff --git a/xmtp_mls/src/identity.rs b/xmtp_mls/src/identity.rs index 5ed52cfc1..f08632eab 100644 --- a/xmtp_mls/src/identity.rs +++ b/xmtp_mls/src/identity.rs @@ -5,10 +5,7 @@ use openmls::{ versions::ProtocolVersion, }; use openmls_basic_credential::SignatureKeyPair; -use openmls_traits::{ - types::{Ciphersuite, CryptoError}, - OpenMlsProvider, -}; +use openmls_traits::{types::CryptoError, OpenMlsProvider}; use prost::Message; use thiserror::Error; use xmtp_cryptography::signature::SignatureError; diff --git a/xmtp_mls/src/lib.rs b/xmtp_mls/src/lib.rs index 378a69f9e..b55ea7da3 100644 --- a/xmtp_mls/src/lib.rs +++ b/xmtp_mls/src/lib.rs @@ -9,6 +9,7 @@ pub mod owner; mod proto_wrapper; pub mod storage; pub mod types; +pub mod verified_key_package; mod xmtp_openmls_provider; pub use client::{Client, Network}; diff --git a/xmtp_mls/src/verified_key_package.rs b/xmtp_mls/src/verified_key_package.rs new file mode 100644 index 000000000..ab7f557be --- /dev/null +++ b/xmtp_mls/src/verified_key_package.rs @@ -0,0 +1,80 @@ +use openmls::prelude::{KeyPackage, KeyPackageIn, KeyPackageVerifyError}; +use openmls_traits::OpenMlsProvider; +use prost::{DecodeError, Message}; +use thiserror::Error; +use tls_codec::{Deserialize, Error as TlsSerializationError}; +use xmtp_proto::xmtp::v3::message_contents::Eip191Association as Eip191AssociationProto; + +use crate::{ + association::{AssociationError, Eip191Association}, + configuration::MLS_PROTOCOL_VERSION, + xmtp_openmls_provider::XmtpOpenMlsProvider, +}; + +#[derive(Debug, Error)] +pub enum KeyPackageVerificationError { + #[error("serialization error: {0}")] + Serialization(#[from] TlsSerializationError), + #[error("mls validation: {0}")] + MlsValidation(#[from] KeyPackageVerifyError), + #[error("association: {0}")] + Association(#[from] AssociationError), + #[error("decode: {0}")] + Decode(#[from] DecodeError), + #[error("generic: {0}")] + Generic(String), +} + +#[derive(Debug, Clone, PartialEq)] +pub struct VerifiedKeyPackage { + pub inner: KeyPackage, + pub wallet_address: String, +} + +impl VerifiedKeyPackage { + pub fn new(inner: KeyPackage, wallet_address: String) -> Self { + Self { + inner, + wallet_address, + } + } + + // Validates starting with a KeyPackage (which is already validated by OpenMLS) + pub fn from_key_package(kp: KeyPackage) -> Result { + let leaf_node = kp.leaf_node(); + let identity_bytes = leaf_node.credential().identity(); + let pub_key_bytes = leaf_node.signature_key().as_slice(); + let wallet_address = identity_to_wallet_address(identity_bytes, pub_key_bytes)?; + + Ok(Self::new(kp, wallet_address)) + } + + // Validates starting with a KeyPackageIn as bytes (which is not validated by OpenMLS) + pub fn from_bytes( + mls_provider: &XmtpOpenMlsProvider, + data: &[u8], + ) -> Result { + let kp_in: KeyPackageIn = KeyPackageIn::tls_deserialize_bytes(data)?; + let kp = kp_in.validate(mls_provider.crypto(), MLS_PROTOCOL_VERSION)?; + + Self::from_key_package(kp) + } + + pub fn installation_id(&self) -> Vec { + self.inner.leaf_node().signature_key().as_slice().to_vec() + } +} + +fn identity_to_wallet_address( + identity_bytes: &[u8], + pub_key_bytes: &[u8], +) -> Result { + let proto_value = Eip191AssociationProto::decode(identity_bytes)?; + let association = Eip191Association::from_proto_with_expected_address( + pub_key_bytes, + proto_value.clone(), + proto_value.wallet_address, + )?; + + Ok(association.address()) +} From 249dabaa1edc54fe051f6ca41c8f6b8322d18373 Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Sun, 29 Oct 2023 10:21:00 -0700 Subject: [PATCH 19/24] Add comment --- xmtp_mls/src/api_client_wrapper.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/xmtp_mls/src/api_client_wrapper.rs b/xmtp_mls/src/api_client_wrapper.rs index e1b4d7745..d1dd85bf6 100644 --- a/xmtp_mls/src/api_client_wrapper.rs +++ b/xmtp_mls/src/api_client_wrapper.rs @@ -373,6 +373,7 @@ mod tests { let mut mock_api = MockApiClient::new(); let start_time_ns = 12; let wallet_addresses = vec!["wallet1".to_string(), "wallet2".to_string()]; + // wallet_addresses gets moved below but needs to be used for assertions later let wallet_addresses_clone = wallet_addresses.clone(); mock_api .expect_get_identity_updates() From 49f8332db75161aa40eeb9b41d54a63dad8be6b9 Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Sun, 29 Oct 2023 13:01:43 -0700 Subject: [PATCH 20/24] Add read_topic API --- xmtp_mls/src/api_client_wrapper.rs | 201 +++++++++++++++++++++++++++-- 1 file changed, 190 insertions(+), 11 deletions(-) diff --git a/xmtp_mls/src/api_client_wrapper.rs b/xmtp_mls/src/api_client_wrapper.rs index d1dd85bf6..9b287fa3a 100644 --- a/xmtp_mls/src/api_client_wrapper.rs +++ b/xmtp_mls/src/api_client_wrapper.rs @@ -1,9 +1,12 @@ use std::collections::HashMap; use xmtp_proto::{ - api_client::{Error as ApiError, ErrorKind, XmtpApiClient, XmtpMlsClient}, + api_client::{ + Envelope, Error as ApiError, ErrorKind, PagingInfo, QueryRequest, XmtpApiClient, + XmtpMlsClient, + }, xmtp::{ - message_api::v3::GetIdentityUpdatesRequest, + message_api::{v1::Cursor, v3::GetIdentityUpdatesRequest}, mls::message_contents::{ group_message::{Version as GroupMessageVersion, V1 as GroupMessageV1}, welcome_message::{Version as WelcomeMessageVersion, V1 as WelcomeMessageV1}, @@ -11,11 +14,14 @@ use xmtp_proto::{ }, }, xmtp::{ - message_api::v3::{ - get_identity_updates_response::update::Kind as UpdateKind, - publish_welcomes_request::WelcomeMessageRequest, ConsumeKeyPackagesRequest, - KeyPackageUpload, PublishToGroupRequest, PublishWelcomesRequest, - RegisterInstallationRequest, UploadKeyPackagesRequest, + message_api::{ + v1::SortDirection, + v3::{ + get_identity_updates_response::update::Kind as UpdateKind, + publish_welcomes_request::WelcomeMessageRequest, ConsumeKeyPackagesRequest, + KeyPackageUpload, PublishToGroupRequest, PublishWelcomesRequest, + RegisterInstallationRequest, UploadKeyPackagesRequest, + }, }, mls::message_contents::GroupMessage, }, @@ -28,12 +34,59 @@ pub struct ApiClientWrapper { impl ApiClientWrapper where - ApiClient: XmtpMlsClient, + ApiClient: XmtpMlsClient + XmtpApiClient, { pub fn new(api_client: ApiClient) -> Self { Self { api_client } } + pub async fn read_topic( + &self, + topic: &str, + start_time_ns: u64, + ) -> Result, ApiError> { + let mut cursor: Option = None; + let mut out: Vec = vec![]; + let page_size = 100; + loop { + let result = self + .api_client + .query(QueryRequest { + content_topics: vec![topic.to_string()], + start_time_ns, + end_time_ns: 0, + paging_info: Some(PagingInfo { + cursor, + limit: page_size, + direction: SortDirection::Ascending as i32, + }), + }) + .await?; + + for envelope in &result.envelopes { + out.push(envelope.clone()); + } + + if result.envelopes.len() < page_size as usize || result.paging_info.is_none() { + break; + } + + cursor = match result.paging_info.unwrap().cursor { + Some(cursor_wrapper) => match cursor_wrapper.cursor { + Some(_) => Some(cursor_wrapper), + None => None, + }, + None => None, + }; + + if cursor.is_none() { + break; + } + } + + Ok(out) + } + pub async fn register_installation( &self, last_resort_key_package: Vec, @@ -228,7 +281,10 @@ type IdentityUpdatesMap = HashMap>; mod tests { use super::ApiClientWrapper; use mockall::mock; - use xmtp_proto::api_client::{Error, XmtpApiClient, XmtpApiSubscription, XmtpMlsClient}; + use xmtp_proto::api_client::{ + Error, PagingInfo, XmtpApiClient, XmtpApiSubscription, XmtpMlsClient, + }; + use xmtp_proto::xmtp::message_api::v1::IndexCursor; use xmtp_proto::xmtp::message_api::v3::consume_key_packages_response::KeyPackage; use xmtp_proto::xmtp::message_api::v3::get_identity_updates_response::update::Kind as UpdateKind; use xmtp_proto::xmtp::message_api::v3::get_identity_updates_response::{ @@ -241,12 +297,24 @@ mod tests { }; use xmtp_proto::xmtp::message_api::v1::{ - BatchQueryRequest, BatchQueryResponse, Envelope, PublishRequest, PublishResponse, - QueryRequest, QueryResponse, SubscribeRequest, + cursor::Cursor as InnerCursor, BatchQueryRequest, BatchQueryResponse, Cursor, Envelope, + PublishRequest, PublishResponse, QueryRequest, QueryResponse, SubscribeRequest, }; use async_trait::async_trait; + fn build_envelopes(num_envelopes: usize, topic: &str) -> Vec { + let mut out: Vec = vec![]; + for i in 0..num_envelopes { + out.push(Envelope { + content_topic: topic.to_string(), + message: vec![i as u8], + timestamp_ns: i as u64, + }) + } + out + } + mock! { pub Subscription {} @@ -442,4 +510,115 @@ mod tests { } } } + + #[tokio::test] + async fn test_read_topic_single_page() { + let mut mock_api = MockApiClient::new(); + let topic = "topic"; + let start_time_ns = 10; + // Set expectation for first request with no cursor + mock_api.expect_query().returning(move |req| { + assert_eq!(req.content_topics[0], topic); + + Ok(QueryResponse { + paging_info: Some(PagingInfo { + cursor: None, + limit: 100, + direction: 0, + }), + envelopes: build_envelopes(10, topic), + }) + }); + + let wrapper = ApiClientWrapper::new(mock_api); + + let result = wrapper.read_topic(topic, start_time_ns).await.unwrap(); + assert_eq!(result.len(), 10); + } + + #[tokio::test] + async fn test_read_topic_single_page_exactly_100_results() { + let mut mock_api = MockApiClient::new(); + let topic = "topic"; + let start_time_ns = 10; + // Set expectation for first request with no cursor + mock_api.expect_query().returning(move |req| { + assert_eq!(req.content_topics[0], topic); + + Ok(QueryResponse { + paging_info: Some(PagingInfo { + cursor: None, + limit: 100, + direction: 0, + }), + envelopes: build_envelopes(100, topic), + }) + }); + + let wrapper = ApiClientWrapper::new(mock_api); + + let result = wrapper.read_topic(topic, start_time_ns).await.unwrap(); + assert_eq!(result.len(), 100); + } + + #[tokio::test] + async fn test_read_topic_multi_page() { + let mut mock_api = MockApiClient::new(); + let topic = "topic"; + let start_time_ns = 10; + // Set expectation for first request with no cursor + mock_api + .expect_query() + .withf(move |req| match req.paging_info.clone() { + Some(paging_info) => match paging_info.cursor { + Some(_) => false, + None => true, + }, + None => true, + } && req.start_time_ns == 10) + .returning(move |req| { + assert_eq!(req.content_topics[0], topic); + + Ok(QueryResponse { + paging_info: Some(PagingInfo { + cursor: Some(Cursor { + cursor: Some(InnerCursor::Index(IndexCursor { + digest: vec![], + sender_time_ns: 0, + })), + }), + limit: 100, + direction: 0, + }), + envelopes: build_envelopes(100, topic), + }) + }); + // Set expectation for requests with a cursor + mock_api + .expect_query() + .withf(|req| match req.paging_info.clone() { + Some(paging_info) => match paging_info.cursor { + Some(_) => true, + None => false, + }, + None => false, + }) + .returning(move |req| { + assert_eq!(req.content_topics[0], topic); + + Ok(QueryResponse { + paging_info: Some(PagingInfo { + cursor: None, + limit: 100, + direction: 0, + }), + envelopes: build_envelopes(100, topic), + }) + }); + + let wrapper = ApiClientWrapper::new(mock_api); + + let result = wrapper.read_topic(topic, start_time_ns).await.unwrap(); + assert_eq!(result.len(), 200); + } } From 23fc060946b840693d0b43f30c732cf72db7ef7d Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Sun, 29 Oct 2023 13:04:49 -0700 Subject: [PATCH 21/24] Remove unused lifetime --- xmtp_mls/src/builder.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xmtp_mls/src/builder.rs b/xmtp_mls/src/builder.rs index 453c178a2..22770e21f 100644 --- a/xmtp_mls/src/builder.rs +++ b/xmtp_mls/src/builder.rs @@ -93,7 +93,7 @@ pub struct ClientBuilder { identity_strategy: IdentityStrategy, } -impl<'a, ApiClient, Owner> ClientBuilder +impl ClientBuilder where ApiClient: XmtpApiClient + XmtpMlsClient, Owner: InboxOwner, From 77f1829463cb9d5d62619390873203fd0a70e72f Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Sun, 29 Oct 2023 13:23:42 -0700 Subject: [PATCH 22/24] Add times check --- xmtp_mls/src/api_client_wrapper.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xmtp_mls/src/api_client_wrapper.rs b/xmtp_mls/src/api_client_wrapper.rs index 9b287fa3a..1d2066fa5 100644 --- a/xmtp_mls/src/api_client_wrapper.rs +++ b/xmtp_mls/src/api_client_wrapper.rs @@ -542,7 +542,7 @@ mod tests { let topic = "topic"; let start_time_ns = 10; // Set expectation for first request with no cursor - mock_api.expect_query().returning(move |req| { + mock_api.expect_query().times(1).returning(move |req| { assert_eq!(req.content_topics[0], topic); Ok(QueryResponse { From c58ed2836ebb90dcb9188b8bff3b08ae7d75b9be Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Tue, 31 Oct 2023 15:11:03 -0700 Subject: [PATCH 23/24] Tidy up cursor extraction --- xmtp_mls/src/api_client_wrapper.rs | 12 +++++------- xmtp_mls/src/client.rs | 1 + 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/xmtp_mls/src/api_client_wrapper.rs b/xmtp_mls/src/api_client_wrapper.rs index 997618bdc..663744821 100644 --- a/xmtp_mls/src/api_client_wrapper.rs +++ b/xmtp_mls/src/api_client_wrapper.rs @@ -71,13 +71,11 @@ where break; } - cursor = match result.paging_info.unwrap().cursor { - Some(cursor_wrapper) => match cursor_wrapper.cursor { - Some(_) => Some(cursor_wrapper), - None => None, - }, - None => None, - }; + cursor = result + .paging_info + .expect("Empty paging info") + .cursor + .map(|wrapper| wrapper); if cursor.is_none() { break; diff --git a/xmtp_mls/src/client.rs b/xmtp_mls/src/client.rs index befc47337..ebf81c04c 100644 --- a/xmtp_mls/src/client.rs +++ b/xmtp_mls/src/client.rs @@ -211,6 +211,7 @@ mod tests { let wallet = generate_local_wallet(); let wallet_address = wallet.get_address(); let client = ClientBuilder::new_test_client(wallet.clone().into()).await; + client.register_identity().await.unwrap(); client.top_up_key_packages().await.unwrap(); From 292a55d4f20e6072aad2f2a24af027b16f5a77ea Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Wed, 1 Nov 2023 14:07:57 -0700 Subject: [PATCH 24/24] Avoid unnecessary clones --- xmtp_mls/src/api_client_wrapper.rs | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/xmtp_mls/src/api_client_wrapper.rs b/xmtp_mls/src/api_client_wrapper.rs index 663744821..49f0c5511 100644 --- a/xmtp_mls/src/api_client_wrapper.rs +++ b/xmtp_mls/src/api_client_wrapper.rs @@ -49,7 +49,7 @@ where let mut out: Vec = vec![]; let page_size = 100; loop { - let result = self + let mut result = self .api_client .query(QueryRequest { content_topics: vec![topic.to_string()], @@ -63,19 +63,14 @@ where }) .await?; - for envelope in &result.envelopes { - out.push(envelope.clone()); - } + let num_envelopes = result.envelopes.len(); + out.append(&mut result.envelopes); - if result.envelopes.len() < page_size as usize || result.paging_info.is_none() { + if num_envelopes < page_size as usize || result.paging_info.is_none() { break; } - cursor = result - .paging_info - .expect("Empty paging info") - .cursor - .map(|wrapper| wrapper); + cursor = result.paging_info.expect("Empty paging info").cursor; if cursor.is_none() { break;