Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding Key Update Before Every New GroupIntent #658

Closed
wants to merge 34 commits into from
Closed
Show file tree
Hide file tree
Changes from 14 commits
Commits
Show all changes
34 commits
Select commit Hold shift + click to select a range
8e87ba7
queue_key_update
Apr 15, 2024
9df8308
lint
Apr 15, 2024
60e8adc
lint
Apr 15, 2024
cb88c9e
merge main
Apr 15, 2024
3438591
queue key update test
Apr 16, 2024
0e9d812
key-update
Apr 22, 2024
1624a05
Fix paths in update-schema script
richardhuaaa Apr 22, 2024
7d0dade
merge master and add rotated_at_ns item
Apr 30, 2024
e62cdfc
add get_rotated_at_ns
Apr 30, 2024
5f35add
fetch rotated_at_ns
May 1, 2024
928bca4
Merge remote-tracking branch 'origin' into dz/key-update
May 1, 2024
b319f56
add pre_intent_hook
May 6, 2024
ed6e1dd
Revert unneeded fix, and specify the correct directory in the README …
richardhuaaa May 13, 2024
cd28efd
apply pre_intent_hook
May 16, 2024
c160c14
Merge remote-tracking branch 'origin' into dz/key-update
May 16, 2024
e1b5f56
merge master and resolve comments
May 16, 2024
c6ed2c0
fix tests
May 17, 2024
79d714e
test 1
May 30, 2024
4a317b4
merge master
May 31, 2024
ceb8d34
modify test
May 31, 2024
ecc3dc0
merge main
Jun 20, 2024
48ad52a
modify
Jun 20, 2024
059df1c
solve db connection and update
Jun 21, 2024
e8a3308
more key updates and test updates
Jun 24, 2024
a4482c6
add
Jun 24, 2024
dc2ba36
test
Jun 28, 2024
b4d676e
finish key update test
Jul 2, 2024
979dcea
Merge remote-tracking branch 'origin/main' into dz/key-update
richardhuaaa Oct 1, 2024
41a7618
Small fixes
richardhuaaa Oct 1, 2024
3d143eb
Fix group initialization error - no need to rotate keys if you are th…
richardhuaaa Oct 1, 2024
6377d85
Undo unit test changes
richardhuaaa Oct 1, 2024
71478f1
Make it impossible to miss a pre_intent_hook, add missing locations
richardhuaaa Oct 2, 2024
bb75ed4
Reduce duplication of maybe_update_installations
richardhuaaa Oct 2, 2024
8750c10
Progress on key rotation for optimistic message sends
richardhuaaa Oct 2, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions xmtp_mls/migrations/2024-04-22-061128_key_update/down.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
-- This file should undo anything in `up.sql`
2 changes: 2 additions & 0 deletions xmtp_mls/migrations/2024-04-22-061128_key_update/up.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
ALTER TABLE groups
ADD COLUMN rotated_at_ns BIGINT NOT NULL DEFAULT 0
2 changes: 2 additions & 0 deletions xmtp_mls/migrations/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
cargo install diesel_cli --no-default-features --features sqlite
```

### Change directory to libxmtp/xmtp_mls/

### Create your migration SQL

In this example the migration is called `create_key_store`:
Expand Down
Binary file not shown.
Empty file.
103 changes: 98 additions & 5 deletions xmtp_mls/src/groups/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -262,12 +262,13 @@ where
added_by_address.clone(),
);
let stored_group = provider.conn().insert_or_ignore_group(to_store)?;

Ok(Self::new(
client,
stored_group.id,
let xmtp_group = Self::new(
client,
stored_group.id,
stored_group.created_at_ns,
))
);
let _ = xmtp_group.queue_key_update();
richardhuaaa marked this conversation as resolved.
Show resolved Hide resolved
Ok(xmtp_group)
}

// Decrypt a welcome message using HPKE and then create and save a group from the stored message
Expand Down Expand Up @@ -344,6 +345,9 @@ where
.map_err(GroupError::EncodeError)?;

let intent_data: Vec<u8> = SendMessageIntentData::new(encoded_envelope).into();

self.pre_intent_hook().await.unwrap();

let intent =
NewGroupIntent::new(IntentKind::SendMessage, self.group_id.clone(), intent_data);
intent.store(conn)?;
Expand Down Expand Up @@ -420,6 +424,9 @@ where
let conn = &mut self.client.store.conn()?;
let intent_data: Vec<u8> =
AddMembersIntentData::new(account_addresses.into()).try_into()?;

self.pre_intent_hook().await.unwrap();

let intent = conn.insert_group_intent(NewGroupIntent::new(
IntentKind::AddMembers,
self.group_id.clone(),
Expand All @@ -436,6 +443,9 @@ where
validate_ed25519_keys(&installation_ids)?;
let conn = &mut self.client.store.conn()?;
let intent_data: Vec<u8> = AddMembersIntentData::new(installation_ids.into()).try_into()?;

Box::pin(self.pre_intent_hook()).await.unwrap();

let intent = conn.insert_group_intent(NewGroupIntent::new(
IntentKind::AddMembers,
self.group_id.clone(),
Expand All @@ -452,6 +462,9 @@ where
let account_addresses = sanitize_evm_addresses(account_addresses_to_remove)?;
let conn = &mut self.client.store.conn()?;
let intent_data: Vec<u8> = RemoveMembersIntentData::new(account_addresses.into()).into();

self.pre_intent_hook().await.unwrap();

let intent = conn.insert_group_intent(NewGroupIntent::new(
IntentKind::RemoveMembers,
self.group_id.clone(),
Expand All @@ -465,6 +478,9 @@ where
let conn = &mut self.client.store.conn()?;
let intent_data: Vec<u8> =
UpdateMetadataIntentData::new_update_group_name(group_name).into();

self.pre_intent_hook().await.unwrap();

let intent = conn.insert_group_intent(NewGroupIntent::new(
IntentKind::MetadataUpdate,
self.group_id.clone(),
Expand Down Expand Up @@ -510,6 +526,9 @@ where
validate_ed25519_keys(&installation_ids)?;
let conn = &mut self.client.store.conn()?;
let intent_data: Vec<u8> = RemoveMembersIntentData::new(installation_ids.into()).into();

self.pre_intent_hook().await.unwrap();

let intent = conn.insert_group_intent(NewGroupIntent::new(
IntentKind::RemoveMembers,
self.group_id.clone(),
Expand All @@ -522,8 +541,44 @@ where
// Update this installation's leaf key in the group by creating a key update commit
pub async fn key_update(&self) -> Result<(), GroupError> {
let conn = &mut self.client.store.conn()?;

self.pre_intent_hook().await.unwrap();

let intent = NewGroupIntent::new(IntentKind::KeyUpdate, self.group_id.clone(), vec![]);
intent.store(conn)?;

self.sync_with_conn(conn).await
}

// Update the installation's leaf key in a synchronous way.
pub fn queue_key_update(&self) -> Result<(), GroupError> {
let conn = &mut self.client.store.conn()?;

self.pre_intent_hook();
richardhuaaa marked this conversation as resolved.
Show resolved Hide resolved

let intent = NewGroupIntent::new(IntentKind::KeyUpdate, self.group_id.clone(), vec![]);
intent.store(conn)?;
Ok(())
}

/// Checking the last key rotation time before rotating the key.
pub async fn pre_intent_hook(&self) -> Result<(), GroupError> {
let conn = &self.client.store.conn()?;
let last_rotated_time = conn.get_rotated_time_checked(self.group_id.clone());

match last_rotated_time {
Ok(rotated_time) => {
if rotated_time == 0 {
let _ = self.queue_key_update();
lfzkoala marked this conversation as resolved.
Show resolved Hide resolved
let update_interval = Some(5_000_000);
self.maybe_update_installations(conn, update_interval).await?;
lfzkoala marked this conversation as resolved.
Show resolved Hide resolved
}
}
Err(err) => {
// Handle the error appropriately
log::error!("database error fetching rotated time {:?}", err);
lfzkoala marked this conversation as resolved.
Show resolved Hide resolved
}
}

self.sync_with_conn(conn).await
}
Expand Down Expand Up @@ -936,6 +991,44 @@ mod tests {
assert_eq!(bola_messages.len(), 1);
}

#[tokio::test]
async fn test_queue_key_update() {
let client = ClientBuilder::new_test_client(&generate_local_wallet()).await;
let bola_client = ClientBuilder::new_test_client(&generate_local_wallet()).await;

let group = client.create_group(None).expect("create group");
group
.add_members(vec![bola_client.account_address()])
.await
.unwrap();

let _ = group.queue_key_update();
group.sync().await.unwrap();
let messages = client
.api_client
.query_group_messages(group.group_id.clone(), None)
.await
.unwrap();
assert_eq!(messages.len(), 2);

let conn = &client.store.conn().unwrap();
let provider = super::XmtpOpenMlsProvider::new(conn);
let mls_group = group.load_mls_group(&provider).unwrap();
let pending_commit = mls_group.pending_commit();
assert!(pending_commit.is_none());

group.send_message(b"hello").await.expect("send message");

bola_client.sync_welcomes().await.unwrap();
let bola_groups = bola_client.find_groups(None, None, None, None).unwrap();
let bola_group = bola_groups.first().unwrap();
bola_group.sync().await.unwrap();
let bola_messages = bola_group
.find_messages(None, None, None, None, None)
.unwrap();
assert_eq!(bola_messages.len(), 1);
}

#[tokio::test]
async fn test_post_commit() {
let client = ClientBuilder::new_test_client(&generate_local_wallet()).await;
Expand Down
16 changes: 16 additions & 0 deletions xmtp_mls/src/storage/encrypted_store/group.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ pub struct StoredGroup {
pub purpose: Purpose,
/// The wallet address of who added the user to a group.
pub added_by_address: String,
pub rotated_at_ns: i64,
}

impl_fetch!(StoredGroup, groups, Vec<u8>);
Expand All @@ -56,6 +57,7 @@ impl StoredGroup {
installations_last_checked: 0,
purpose: Purpose::Conversation,
added_by_address,
rotated_at_ns: 0,
}
}

Expand All @@ -72,6 +74,7 @@ impl StoredGroup {
installations_last_checked: 0,
purpose: Purpose::Sync,
added_by_address: "".into(),
rotated_at_ns: 0,
}
}
}
Expand Down Expand Up @@ -155,6 +158,19 @@ impl DbConnection<'_> {
last_ts.ok_or(StorageError::NotFound)
}

pub fn get_rotated_time_checked(&self, group_id: Vec<u8>) -> Result<i64, StorageError> {
let last_ts = self.raw_query(|conn| {
let ts = dsl::groups
.find(&group_id)
.select(dsl::rotated_at_ns)
.first(conn)
.optional()?;
Ok(ts)
})?;

last_ts.ok_or(StorageError::NotFound)
}

/// Updates the 'last time checked' we checked for new installations.
pub fn update_installations_time_checked(&self, group_id: Vec<u8>) -> Result<(), StorageError> {
self.raw_query(|conn| {
Expand Down
1 change: 1 addition & 0 deletions xmtp_mls/src/storage/encrypted_store/schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ diesel::table! {
installations_last_checked -> BigInt,
purpose -> Integer,
added_by_address -> Text,
rotated_at_ns -> BigInt,
}
}

Expand Down