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

@@ -333,12 +333,9 @@ impl FfiXmtpClient {
         existing_wallet_address: &str,
         new_wallet_address: &str,
     ) -> Result<Arc<FfiSignatureRequest>, 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<Arc<FfiSignatureRequest>, 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 {

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::<Vec<_>>();
+        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::<Vec<_>>();
+        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);
+    }
 }

xmtp_mls/src/identity_updates.rs

@@ -222,34 +222,57 @@ 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<SignatureRequest, ClientError> {
+        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
             .add_association(new_wallet_address.into(), existing_wallet_address.into())
             .build())
     }
 
-    pub async fn revoke_wallet(
+    pub async fn revoke_wallets(
         &self,
-        inbox_id: String,
-        wallet_to_revoke: String,
+        wallets_to_revoke: Vec<String>,
     ) -> Result<SignatureRequest, ClientError> {
+        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<Vec<u8>>,
+    ) -> Result<SignatureRequest, ClientError> {
+        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<ApiClient: XmtpApi>(
 }
 
 #[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<u8> = 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);
+    }
 }