From 998623d7d3a9fd1a1ecdc08779a144f6b9c53105 Mon Sep 17 00:00:00 2001 From: neekolas Date: Fri, 20 Sep 2024 14:39:14 +0000 Subject: [PATCH] Create DB for key package rotation (#1068) ## tl;dr - Creates database tables to track key packages for rotation --- .../down.sql | 2 + .../up.sql | 6 + .../encrypted_store/key_package_history.rs | 107 ++++++++++++++++++ xmtp_mls/src/storage/encrypted_store/mod.rs | 1 + .../src/storage/encrypted_store/schema.rs | 9 ++ 5 files changed, 125 insertions(+) create mode 100644 xmtp_mls/migrations/2024-09-16-221459_key_package_history/down.sql create mode 100644 xmtp_mls/migrations/2024-09-16-221459_key_package_history/up.sql create mode 100644 xmtp_mls/src/storage/encrypted_store/key_package_history.rs diff --git a/xmtp_mls/migrations/2024-09-16-221459_key_package_history/down.sql b/xmtp_mls/migrations/2024-09-16-221459_key_package_history/down.sql new file mode 100644 index 000000000..2f4a654c3 --- /dev/null +++ b/xmtp_mls/migrations/2024-09-16-221459_key_package_history/down.sql @@ -0,0 +1,2 @@ +DROP TABLE IF EXISTS key_package_history; + diff --git a/xmtp_mls/migrations/2024-09-16-221459_key_package_history/up.sql b/xmtp_mls/migrations/2024-09-16-221459_key_package_history/up.sql new file mode 100644 index 000000000..f5a0a78a1 --- /dev/null +++ b/xmtp_mls/migrations/2024-09-16-221459_key_package_history/up.sql @@ -0,0 +1,6 @@ +CREATE TABLE key_package_history( + "id" integer PRIMARY KEY AUTOINCREMENT NOT NULL, + "key_package_hash_ref" BLOB UNIQUE NOT NULL, + "created_at_ns" bigint NOT NULL +); + diff --git a/xmtp_mls/src/storage/encrypted_store/key_package_history.rs b/xmtp_mls/src/storage/encrypted_store/key_package_history.rs new file mode 100644 index 000000000..6fa4ebecb --- /dev/null +++ b/xmtp_mls/src/storage/encrypted_store/key_package_history.rs @@ -0,0 +1,107 @@ +use diesel::prelude::*; + +use super::{db_connection::DbConnection, schema::key_package_history, StorageError}; +use crate::{impl_store_or_ignore, utils::time::now_ns, StoreOrIgnore}; + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = key_package_history)] +pub struct NewKeyPackageHistoryEntry { + pub key_package_hash_ref: Vec, + pub created_at_ns: i64, +} + +#[derive(Queryable, Selectable, Debug, Clone)] +#[diesel(table_name = key_package_history)] +pub struct StoredKeyPackageHistoryEntry { + pub id: i32, + pub key_package_hash_ref: Vec, + pub created_at_ns: i64, +} + +impl_store_or_ignore!(NewKeyPackageHistoryEntry, key_package_history); + +impl DbConnection { + pub fn store_key_package_history_entry( + &self, + key_package_hash_ref: Vec, + ) -> Result { + let entry = NewKeyPackageHistoryEntry { + key_package_hash_ref: key_package_hash_ref.clone(), + created_at_ns: now_ns(), + }; + entry.store_or_ignore(self)?; + + self.find_key_package_history_entry_by_hash_ref(key_package_hash_ref) + } + + pub fn find_key_package_history_entry_by_hash_ref( + &self, + hash_ref: Vec, + ) -> Result { + let result = self.raw_query(|conn| { + key_package_history::dsl::key_package_history + .filter(key_package_history::dsl::key_package_hash_ref.eq(hash_ref)) + .first::(conn) + })?; + + Ok(result) + } + + pub fn find_key_package_history_entries_before_id( + &self, + id: i32, + ) -> Result, StorageError> { + let result = self.raw_query(|conn| { + key_package_history::dsl::key_package_history + .filter(key_package_history::dsl::id.lt(id)) + .load::(conn) + })?; + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use crate::{storage::encrypted_store::tests::with_connection, utils::test::rand_vec}; + + #[test] + fn test_store_key_package_history_entry() { + with_connection(|conn| { + let hash_ref = rand_vec(); + let new_entry = conn + .store_key_package_history_entry(hash_ref.clone()) + .unwrap(); + assert_eq!(new_entry.key_package_hash_ref, hash_ref); + assert_eq!(new_entry.id, 1); + }) + } + + #[test] + fn test_store_multiple() { + with_connection(|conn| { + let hash_ref1 = rand_vec(); + let hash_ref2 = rand_vec(); + let hash_ref3 = rand_vec(); + + conn.store_key_package_history_entry(hash_ref1.clone()) + .unwrap(); + conn.store_key_package_history_entry(hash_ref2.clone()) + .unwrap(); + let entry_3 = conn + .store_key_package_history_entry(hash_ref3.clone()) + .unwrap(); + + let all_entries = conn + .find_key_package_history_entries_before_id(100) + .unwrap(); + + assert_eq!(all_entries.len(), 3); + + let earlier_entries = conn + .find_key_package_history_entries_before_id(entry_3.id) + .unwrap(); + assert_eq!(earlier_entries.len(), 2); + }) + } +} diff --git a/xmtp_mls/src/storage/encrypted_store/mod.rs b/xmtp_mls/src/storage/encrypted_store/mod.rs index c117d80d8..687ba7631 100644 --- a/xmtp_mls/src/storage/encrypted_store/mod.rs +++ b/xmtp_mls/src/storage/encrypted_store/mod.rs @@ -18,6 +18,7 @@ pub mod group_intent; pub mod group_message; pub mod identity; pub mod identity_update; +pub mod key_package_history; pub mod key_store_entry; pub mod refresh_state; pub mod schema; diff --git a/xmtp_mls/src/storage/encrypted_store/schema.rs b/xmtp_mls/src/storage/encrypted_store/schema.rs index 026969582..7d835a5b9 100644 --- a/xmtp_mls/src/storage/encrypted_store/schema.rs +++ b/xmtp_mls/src/storage/encrypted_store/schema.rs @@ -74,6 +74,14 @@ diesel::table! { } } +diesel::table! { + key_package_history (id) { + id -> Integer, + key_package_hash_ref -> Binary, + created_at_ns -> BigInt, + } +} + diesel::table! { openmls_key_store (key_bytes) { key_bytes -> Binary, @@ -108,6 +116,7 @@ diesel::allow_tables_to_appear_in_same_query!( groups, identity, identity_updates, + key_package_history, openmls_key_store, openmls_key_value, refresh_state,