From 4d403632f5d8f56b62247eb8f82174789bcdeb08 Mon Sep 17 00:00:00 2001 From: Naomi Plasterer Date: Wed, 4 Dec 2024 11:58:07 -0800 Subject: [PATCH 1/2] add migration for hmac keys database --- .../down.sql | 1 + .../2024-12-04-194513_hmac-keys-records-table/up.sql | 11 +++++++++++ xmtp_mls/src/storage/encrypted_store/schema.rs | 10 ++++++++++ 3 files changed, 22 insertions(+) create mode 100644 xmtp_mls/migrations/2024-12-04-194513_hmac-keys-records-table/down.sql create mode 100644 xmtp_mls/migrations/2024-12-04-194513_hmac-keys-records-table/up.sql diff --git a/xmtp_mls/migrations/2024-12-04-194513_hmac-keys-records-table/down.sql b/xmtp_mls/migrations/2024-12-04-194513_hmac-keys-records-table/down.sql new file mode 100644 index 000000000..992e800e0 --- /dev/null +++ b/xmtp_mls/migrations/2024-12-04-194513_hmac-keys-records-table/down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "hmac_key_records"; diff --git a/xmtp_mls/migrations/2024-12-04-194513_hmac-keys-records-table/up.sql b/xmtp_mls/migrations/2024-12-04-194513_hmac-keys-records-table/up.sql new file mode 100644 index 000000000..b4a679a4d --- /dev/null +++ b/xmtp_mls/migrations/2024-12-04-194513_hmac-keys-records-table/up.sql @@ -0,0 +1,11 @@ +CREATE TABLE "hmac_key_records"( + -- Group ID that the Hmac keys are associated with + "group_id" BLOB NOT NULL, + -- Dm ID that the Hmac keys are associated with + "dm_id" TEXT, + -- The hmac key + "hmac_key" BLOB NOT NULL, + -- The number of 30 day periods since epoch + "thirty_day_periods_since_epoch" INT NOT NULL, + PRIMARY KEY ("group_id", "hmac_key") +); diff --git a/xmtp_mls/src/storage/encrypted_store/schema.rs b/xmtp_mls/src/storage/encrypted_store/schema.rs index 5a6e3385b..6201e80be 100644 --- a/xmtp_mls/src/storage/encrypted_store/schema.rs +++ b/xmtp_mls/src/storage/encrypted_store/schema.rs @@ -58,6 +58,15 @@ diesel::table! { } } +diesel::table! { + hmac_key_records (group_id, hmac_key) { + group_id -> Binary, + dm_id -> Nullable, + hmac_key -> Binary, + thirty_day_periods_since_epoch -> Integer, + } +} + diesel::table! { identity (rowid) { inbox_id -> Text, @@ -123,6 +132,7 @@ diesel::allow_tables_to_appear_in_same_query!( group_intents, group_messages, groups, + hmac_key_records, identity, identity_updates, key_package_history, From 1e1cf1553f6cfcbf7ea78053e0e887db49e74e46 Mon Sep 17 00:00:00 2001 From: Naomi Plasterer Date: Wed, 4 Dec 2024 12:15:13 -0800 Subject: [PATCH 2/2] add the database managament --- .../encrypted_store/hmac_key_record.rs | 159 ++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 xmtp_mls/src/storage/encrypted_store/hmac_key_record.rs diff --git a/xmtp_mls/src/storage/encrypted_store/hmac_key_record.rs b/xmtp_mls/src/storage/encrypted_store/hmac_key_record.rs new file mode 100644 index 000000000..96e931e5c --- /dev/null +++ b/xmtp_mls/src/storage/encrypted_store/hmac_key_record.rs @@ -0,0 +1,159 @@ +use crate::{impl_store, storage::StorageError}; + +use super::Sqlite; +use super::{ + db_connection::DbConnection, + schema::consent_records::{self, dsl}, +}; +use diesel::{ + backend::Backend, + deserialize::{self, FromSql, FromSqlRow}, + expression::AsExpression, + prelude::*, + serialize::{self, IsNull, Output, ToSql}, + sql_types::Integer, + upsert::excluded, +}; +use serde::{Deserialize, Serialize}; +use xmtp_id::associations::DeserializationError; +use xmtp_proto::xmtp::mls::message_contents::{ + ConsentEntityType, ConsentState as ConsentStateProto, ConsentUpdate as ConsentUpdateProto, +}; + +/// StoredConsentRecord holds a serialized ConsentRecord +#[derive(Insertable, Queryable, Debug, Clone, PartialEq, Eq, Deserialize, Serialize)] +#[diesel(table_name = hmac_key_records)] +#[diesel(primary_key(group_id, hmac_key))] +pub struct StoredHmacKeyRecord { + /// The group id associate with these hmac keys + pub group_id: Vec, + /// The dm id for stitching + pub dm_id: Option, + /// The hmac key + pub hmac_key: Vec, + /// The number of 30 day periods since epoch + pub thirty_day_periods_since_epoch: i32, +} + +impl StoredHmacKeyRecord { + pub fn new(group_id: Vec, dm_id: Option, hmac_key: Vec, thirty_day_periods_since_epoch: i32) -> Self { + Self { + group_id, + dm_id, + hmac_key, + thirty_day_periods_since_epoch + } + } +} + +impl_store!(StoredHmacKeyRecord, hmac_key_records); + +impl DbConnection { + /// Returns all hmac_key_records for the given group_id + pub fn get_hmac_key_records( + &self, + group_id: Vec, + ) -> Result, StorageError> { + Ok(self.raw_query(|conn| -> diesel::QueryResult<_> { + dsl::hmac_key_records + .filter(dsl::group_id.eq(group_id)) + .load::(conn) + })?) + } + + /// Insert hmac_key_records without replacing existing ones + pub fn insert_hmac_key_records( + &self, + records: &[StoredHmacKeyRecord], + ) -> Result<(), StorageError> { + self.raw_query(|conn| -> diesel::QueryResult<_> { + conn.transaction::<_, diesel::result::Error, _>(|conn| { + for record in records.iter() { + diesel::insert_into(dsl::hmac_key_records) + .values(record) + .on_conflict_do_nothing() + .execute(conn)?; + } + Ok(()) + }) + })?; + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use crate::storage::encrypted_store::tests::with_connection; + #[cfg(target_arch = "wasm32")] + wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_dedicated_worker); + + use super::*; + + fn generate_hmac_key_record( + group_id: Vec, + hmac_key: Vec, + thirty_day_periods_since_epoch: i32, + ) -> StoredHmacKeyRecord { + StoredHmacKeyRecord { + group_id, + None, + hmac_key, + thirty_day_periods_since_epoch, + } + } + + #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)] + #[cfg_attr(not(target_arch = "wasm32"), tokio::test)] + async fn insert_and_read_hmac_key_records() { + with_connection(|conn| { + // Prepare test data + let group_id = b"group_id".to_vec(); + let hmac_key = b"hmac_key".to_vec(); + let thirty_day_periods_since_epoch = 123; + + let hmac_record = generate_hmac_key_record( + group_id.clone(), + hmac_key.clone(), + thirty_day_periods_since_epoch, + ); + + // Insert the record + conn.insert_hmac_key_records(&[hmac_record.clone()]) + .expect("should insert hmac_key_record without error"); + + // Read back the inserted record + let records = conn + .get_hmac_key_records(group_id.clone()) + .expect("query should work"); + + // Ensure the records match + assert_eq!(records.len(), 1, "There should be exactly one record"); + let retrieved_record = &records[0]; + + assert_eq!(retrieved_record.group_id, hmac_record.group_id); + assert_eq!(retrieved_record.hmac_key, hmac_record.hmac_key); + assert_eq!( + retrieved_record.thirty_day_periods_since_epoch, + hmac_record.thirty_day_periods_since_epoch + ); + + // Insert a second record (same group_id and thirty_day_periods_since_epoch) + let conflict_record = generate_hmac_key_record( + group_id.clone(), + b"new_hmac_key".to_vec(), // Different hmac_key + thirty_day_periods_since_epoch, + ); + + conn.insert_hmac_key_records(&[conflict_record]) + .expect("should insert second record without error"); + + let records = conn + .get_hmac_key_records(group_id.clone()) + .expect("query should work"); + + assert_eq!(records.len(), 2, "Both records should exist"); + }) + .await; + } +}