From b54f94f1f3c3d58ded45d0486576362e47a8a702 Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Wed, 21 Aug 2024 10:16:26 -0700 Subject: [PATCH] Add ability to revoke installations (#977) ## tl;dr - Exposes the ability to revoke installations - Updates the internal API for revoking wallets to not require the inbox_id. We can safely assume you're always doing this with your own `inbox_id`. - Adds a bunch more tests around revocation --- bindings_ffi/src/mls.rs | 12 +-- xmtp_mls/src/groups/mod.rs | 49 +++++++++ xmtp_mls/src/identity_updates.rs | 174 ++++++++++++++++++++++--------- 3 files changed, 177 insertions(+), 58 deletions(-) diff --git a/bindings_ffi/src/mls.rs b/bindings_ffi/src/mls.rs index ed2568b53..f01d6102e 100644 --- a/bindings_ffi/src/mls.rs +++ b/bindings_ffi/src/mls.rs @@ -333,12 +333,9 @@ impl FfiXmtpClient { existing_wallet_address: &str, new_wallet_address: &str, ) -> Result, GenericError> { - let inbox_id = self.inner_client.inbox_id(); - let signature_request = self.inner_client.associate_wallet( - inbox_id, - existing_wallet_address.into(), - new_wallet_address.into(), - )?; + let signature_request = self + .inner_client + .associate_wallet(existing_wallet_address.into(), new_wallet_address.into())?; let request = Arc::new(FfiSignatureRequest { inner: Arc::new(tokio::sync::Mutex::new(signature_request)), @@ -364,10 +361,9 @@ impl FfiXmtpClient { &self, wallet_address: &str, ) -> Result, GenericError> { - let inbox_id = self.inner_client.inbox_id(); let signature_request = self .inner_client - .revoke_wallet(inbox_id, wallet_address.into()) + .revoke_wallets(vec![wallet_address.into()]) .await?; let request = Arc::new(FfiSignatureRequest { diff --git a/xmtp_mls/src/groups/mod.rs b/xmtp_mls/src/groups/mod.rs index 3160f98b5..b842aa2ca 100644 --- a/xmtp_mls/src/groups/mod.rs +++ b/xmtp_mls/src/groups/mod.rs @@ -1254,6 +1254,7 @@ mod tests { members::{GroupMember, PermissionLevel}, DeliveryStatus, GroupMetadataOptions, PreconfiguredPolicies, UpdateAdminListType, }, + identity_updates::tests::sign_with_wallet, storage::{ group_intent::IntentState, group_message::{GroupMessageKind, StoredGroupMessage}, @@ -2914,4 +2915,52 @@ mod tests { panic!("Expected error") } } + + #[tokio::test(flavor = "multi_thread")] + async fn ensure_removed_after_revoke() { + let alix_wallet = generate_local_wallet(); + let bo_wallet = generate_local_wallet(); + let alix1 = ClientBuilder::new_test_client(&alix_wallet).await; + let alix2 = ClientBuilder::new_test_client(&alix_wallet).await; + let bo = ClientBuilder::new_test_client(&bo_wallet).await; + + let alix_group = alix1 + .create_group(None, GroupMetadataOptions::default()) + .unwrap(); + alix_group + .add_members(&alix1, vec![bo_wallet.get_address()]) + .await + .unwrap(); + let bo_group = receive_group_invite(&bo).await; + + // Check the MLS group for the number of members + let bo_provider = bo.mls_provider(bo.store().conn().unwrap()); + let bo_mls_group = bo_group.load_mls_group(&bo_provider).unwrap(); + let members = bo_mls_group.members().collect::>(); + assert_eq!(members.len(), 3); + + let mut revoke_installation_request = alix1 + .revoke_installations(vec![alix2.installation_public_key()]) + .await + .unwrap(); + + sign_with_wallet(&alix_wallet, &mut revoke_installation_request).await; + alix1 + .apply_signature_request(revoke_installation_request) + .await + .unwrap(); + + bo_group.sync(&bo).await.unwrap(); + // Check the MLS group for the number of members after alix2 has been removed + let bo_mls_group = bo_group.load_mls_group(&bo_provider).unwrap(); + let members = bo_mls_group.members().collect::>(); + assert_eq!(members.len(), 2); + + let members = bo_group.members().unwrap(); + let alix_member = members + .iter() + .find(|m| m.inbox_id == alix1.inbox_id()) + .unwrap(); + assert_eq!(alix_member.installation_ids.len(), 1); + } } diff --git a/xmtp_mls/src/identity_updates.rs b/xmtp_mls/src/identity_updates.rs index 14a4cd44a..817f40143 100644 --- a/xmtp_mls/src/identity_updates.rs +++ b/xmtp_mls/src/identity_updates.rs @@ -222,11 +222,11 @@ where pub fn associate_wallet( &self, - // TODO: Replace this argument with a value stored on the client - inbox_id: String, existing_wallet_address: String, new_wallet_address: String, ) -> Result { + log::info!("Associating new wallet with inbox_id {}", self.inbox_id()); + let inbox_id = self.inbox_id(); let builder = SignatureRequestBuilder::new(inbox_id); Ok(builder @@ -234,22 +234,45 @@ where .build()) } - pub async fn revoke_wallet( + pub async fn revoke_wallets( &self, - inbox_id: String, - wallet_to_revoke: String, + wallets_to_revoke: Vec, ) -> Result { + let inbox_id = self.inbox_id(); let current_state = self .get_association_state(&self.store().conn()?, &inbox_id, None) .await?; - let builder = SignatureRequestBuilder::new(inbox_id); + let mut builder = SignatureRequestBuilder::new(inbox_id); - Ok(builder - .revoke_association( + for wallet in wallets_to_revoke { + builder = builder.revoke_association( current_state.recovery_address().clone().into(), - wallet_to_revoke.into(), + wallet.into(), ) - .build()) + } + + Ok(builder.build()) + } + + pub async fn revoke_installations( + &self, + installation_ids: Vec>, + ) -> Result { + let inbox_id = self.inbox_id(); + let current_state = self + .get_association_state(&self.store().conn()?, &inbox_id, None) + .await?; + + let mut builder = SignatureRequestBuilder::new(inbox_id); + + for installation_id in installation_ids { + builder = builder.revoke_association( + current_state.recovery_address().clone().into(), + installation_id.into(), + ) + } + + Ok(builder.build()) } pub async fn apply_signature_request( @@ -398,7 +421,7 @@ pub async fn load_identity_updates( } #[cfg(test)] -mod tests { +pub(crate) mod tests { use ethers::signers::LocalWallet; use tracing_test::traced_test; use xmtp_cryptography::utils::generate_local_wallet; @@ -418,7 +441,10 @@ mod tests { use super::load_identity_updates; - async fn sign_with_wallet(wallet: &LocalWallet, signature_request: &mut SignatureRequest) { + pub(crate) async fn sign_with_wallet( + wallet: &LocalWallet, + signature_request: &mut SignatureRequest, + ) { let wallet_signature: Vec = wallet .sign(signature_request.signature_text().as_str()) .unwrap() @@ -493,25 +519,8 @@ mod tests { let wallet_2_address = wallet_2.get_address(); let client = ClientBuilder::new_test_client(&wallet).await; - let mut signature_request: SignatureRequest = client - .create_inbox(wallet_address.clone(), None) - .await - .unwrap(); - let inbox_id = signature_request.inbox_id(); - - sign_with_wallet(&wallet, &mut signature_request).await; - - client - .apply_signature_request(signature_request) - .await - .unwrap(); - let mut add_association_request = client - .associate_wallet( - inbox_id.clone(), - wallet_address.clone(), - wallet_2_address.clone(), - ) + .associate_wallet(wallet_address.clone(), wallet_2_address.clone()) .unwrap(); sign_with_wallet(&wallet, &mut add_association_request).await; @@ -522,7 +531,7 @@ mod tests { .await .unwrap(); - let association_state = get_association_state(&client, inbox_id.clone()).await; + let association_state = get_association_state(&client, client.inbox_id()).await; assert_eq!(association_state.members().len(), 3); assert_eq!(association_state.recovery_address(), &wallet_address); @@ -537,19 +546,7 @@ mod tests { let wallet_address = wallet.get_address(); let wallet_2_address = wallet_2.get_address(); let client = ClientBuilder::new_test_client(&wallet).await; - - let mut signature_request: SignatureRequest = client - .create_inbox(wallet_address.clone(), None) - .await - .unwrap(); - let inbox_id = signature_request.inbox_id(); - - sign_with_wallet(&wallet, &mut signature_request).await; - - client - .apply_signature_request(signature_request) - .await - .unwrap(); + let inbox_id = client.inbox_id(); get_association_state(&client, inbox_id.clone()).await; @@ -568,11 +565,7 @@ mod tests { assert_logged!("Wrote association", 1); let mut add_association_request = client - .associate_wallet( - inbox_id.clone(), - wallet_address.clone(), - wallet_2_address.clone(), - ) + .associate_wallet(wallet_address.clone(), wallet_2_address.clone()) .unwrap(); sign_with_wallet(&wallet, &mut add_association_request).await; @@ -650,7 +643,7 @@ mod tests { .unwrap(); let new_wallet = generate_local_wallet(); let mut add_association_request = client - .associate_wallet(inbox_id, wallet.get_address(), new_wallet.get_address()) + .associate_wallet(wallet.get_address(), new_wallet.get_address()) .unwrap(); sign_with_wallet(&wallet, &mut add_association_request).await; @@ -722,4 +715,85 @@ mod tests { .removed_installations .contains(&client_2_installation_key)); } + + #[tokio::test] + pub async fn revoke_wallet() { + let recovery_wallet = generate_local_wallet(); + let second_wallet = generate_local_wallet(); + let client = ClientBuilder::new_test_client(&recovery_wallet).await; + + let mut add_wallet_signature_request = client + .associate_wallet(recovery_wallet.get_address(), second_wallet.get_address()) + .unwrap(); + + sign_with_wallet(&recovery_wallet, &mut add_wallet_signature_request).await; + sign_with_wallet(&second_wallet, &mut add_wallet_signature_request).await; + + client + .apply_signature_request(add_wallet_signature_request) + .await + .unwrap(); + + let association_state_after_add = get_association_state(&client, client.inbox_id()).await; + assert_eq!(association_state_after_add.account_addresses().len(), 2); + + // Make sure the inbox ID is correctly registered + let inbox_ids = client + .api_client + .get_inbox_ids(vec![second_wallet.get_address()]) + .await + .unwrap(); + assert_eq!(inbox_ids.len(), 1); + + // Now revoke the second wallet + + let mut revoke_signature_request = client + .revoke_wallets(vec![second_wallet.get_address()]) + .await + .unwrap(); + sign_with_wallet(&recovery_wallet, &mut revoke_signature_request).await; + client + .apply_signature_request(revoke_signature_request) + .await + .unwrap(); + + // Make sure that the association state has removed the second wallet + let association_state_after_revoke = + get_association_state(&client, client.inbox_id()).await; + assert_eq!(association_state_after_revoke.account_addresses().len(), 1); + + // Make sure the inbox ID is correctly unregistered + let inbox_ids = client + .api_client + .get_inbox_ids(vec![second_wallet.get_address()]) + .await + .unwrap(); + assert_eq!(inbox_ids.len(), 0); + } + + #[tokio::test] + pub async fn revoke_installation() { + let wallet = generate_local_wallet(); + let client1 = ClientBuilder::new_test_client(&wallet).await; + let client2 = ClientBuilder::new_test_client(&wallet).await; + + let association_state = get_association_state(&client1, client1.inbox_id()).await; + // Ensure there are two installations on the inbox + assert_eq!(association_state.installation_ids().len(), 2); + + // Now revoke the second client + let mut revoke_installation_request = client1 + .revoke_installations(vec![client2.installation_public_key()]) + .await + .unwrap(); + sign_with_wallet(&wallet, &mut revoke_installation_request).await; + client1 + .apply_signature_request(revoke_installation_request) + .await + .unwrap(); + + // Make sure there is only one installation on the inbox + let association_state = get_association_state(&client1, client1.inbox_id()).await; + assert_eq!(association_state.installation_ids().len(), 1); + } }