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.

Sign up for GitHub

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

Add ability to revoke installations #977

Merged
merged 2 commits into from
Aug 21, 2024
Merged

Add ability to revoke installations #977

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
09b59d1
Fix revocation
neekolas Aug 21, 2024
e3f387b
Fix broken tests
neekolas Aug 21, 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
12 changes: 4 additions & 8 deletions bindings_ffi/src/mls.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)),
Expand All @@ -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 {
Expand Down
49 changes: 49 additions & 0 deletions xmtp_mls/src/groups/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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);
}
}
174 changes: 124 additions & 50 deletions xmtp_mls/src/identity_updates.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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;
Expand All @@ -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()
Expand Down Expand Up @@ -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;
Expand All @@ -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);
Expand All @@ -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;

Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
}