diff --git a/Cargo.lock b/Cargo.lock index fe7d560ac..d049ab82c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5795,7 +5795,9 @@ dependencies = [ "openmls_rust_crypto", "openmls_traits", "prost 0.12.3", + "rand", "serde", + "sha2 0.10.8", "thiserror", "tracing", "xmtp_cryptography", diff --git a/bindings_ffi/Cargo.lock b/bindings_ffi/Cargo.lock index ce3dade1a..f835d820b 100644 --- a/bindings_ffi/Cargo.lock +++ b/bindings_ffi/Cargo.lock @@ -5383,7 +5383,6 @@ dependencies = [ "pbjson-types 0.5.1", "prost 0.12.3", "serde", - "serde_json", "tokio", "tonic", "tower", @@ -5474,9 +5473,7 @@ name = "xmtp_v2" version = "0.1.0" dependencies = [ "aes-gcm", - "chrono", "ecdsa 0.15.1", - "ethers-core", "generic-array", "getrandom", "hex", @@ -5484,11 +5481,8 @@ dependencies = [ "k256 0.12.0", "rand", "rand_chacha", - "rlp", - "serde", "sha2", "sha3", - "thiserror", ] [[package]] diff --git a/xmtp_id/Cargo.toml b/xmtp_id/Cargo.toml index 76532705a..e79db7af2 100644 --- a/xmtp_id/Cargo.toml +++ b/xmtp_id/Cargo.toml @@ -1,24 +1,25 @@ [package] +edition = "2021" name = "xmtp_id" version = "0.1.0" -edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +async-trait.workspace = true +chrono.workspace = true +futures.workspace = true log.workspace = true -tracing.workspace = true -thiserror.workspace = true -xmtp_cryptography.workspace = true -xmtp_mls.workspace = true -xmtp_proto.workspace = true -openmls_traits.workspace = true openmls.workspace = true openmls_basic_credential.workspace = true openmls_rust_crypto.workspace = true +openmls_traits.workspace = true prost.workspace = true -chrono.workspace = true +rand.workspace = true serde.workspace = true -async-trait.workspace = true -futures.workspace = true - +sha2 = "0.10.8" +thiserror.workspace = true +tracing.workspace = true +xmtp_cryptography.workspace = true +xmtp_mls.workspace = true +xmtp_proto.workspace = true diff --git a/xmtp_id/src/associations/association_log.rs b/xmtp_id/src/associations/association_log.rs new file mode 100644 index 000000000..098917af3 --- /dev/null +++ b/xmtp_id/src/associations/association_log.rs @@ -0,0 +1,278 @@ +use super::entity::{Entity, EntityRole}; +use super::hashes::{generate_xid, sha256_string}; +use super::signature::{Signature, SignatureError, SignatureKind}; +use super::state::{AssociationState, StateError}; + +use thiserror::Error; + +// const ALLOWED_CREATE_ENTITY_ROLES: [EntityRole; 2] = [EntityRole::LegacyKey, EntityRole::Address]; + +#[derive(Debug, Error, PartialEq)] +pub enum AssociationError { + #[error("Error creating association {0}")] + Generic(String), + #[error("Multiple create operations detected")] + MultipleCreate, + #[error("XID not yet created")] + NotCreated, + #[error("Signature validation failed {0}")] + Signature(#[from] SignatureError), + #[error("State update failed")] + StateError(#[from] StateError), + #[error("Missing existing member")] + MissingExistingMember, + #[error("Legacy key is only allowed to be associated using a legacy signature with nonce 0")] + LegacySignatureReuse, + #[error("Signature not allowed for role {0:?} {1:?}")] + SignatureNotAllowed(EntityRole, SignatureKind), + #[error("Replay detected")] + Replay, +} + +pub trait LogEntry { + fn update_state( + &self, + existing_state: Option, + ) -> Result; + fn hash(&self) -> String; +} + +pub struct CreateXid { + pub nonce: u32, + pub account_address: String, + pub initial_association: AddAssociation, +} + +impl LogEntry for CreateXid { + fn update_state( + &self, + existing_state: Option, + ) -> Result { + if existing_state.is_some() { + return Err(AssociationError::MultipleCreate); + } + + let account_address = self.account_address.clone(); + + let initial_state = AssociationState::new(account_address, self.nonce); + let new_state = self.initial_association.update_state(Some(initial_state))?; + + Ok(new_state.mark_event_seen(self.hash())) + } + + fn hash(&self) -> String { + // Once we have real signatures the nonce and the recovery address should become part of the text + let inputs = format!( + "{}{}{}", + self.nonce, + self.account_address, + self.initial_association.hash() + ); + + sha256_string(inputs) + } +} + +pub struct AddAssociation { + pub client_timestamp_ns: u32, + pub new_member_role: EntityRole, + pub new_member_signature: Box, + pub existing_member_signature: Box, +} + +impl AddAssociation { + pub fn new_member_address(&self) -> String { + self.new_member_signature.recover_signer().unwrap() + } +} + +impl LogEntry for AddAssociation { + fn update_state( + &self, + maybe_existing_state: Option, + ) -> Result { + let existing_state = maybe_existing_state.ok_or(AssociationError::NotCreated)?; + + // Catch replays per-association + // The real hash function should probably just be the signature text, but since that's stubbed out I have some more inputs + let association_hash = self.hash(); + if existing_state.has_seen(&association_hash) { + return Err(AssociationError::Replay); + } + + let new_member_address = self.new_member_signature.recover_signer()?; + let existing_member_address = self.existing_member_signature.recover_signer()?; + let recovery_address = &existing_state.recovery_address; + + // You cannot add yourself + if new_member_address == existing_member_address { + return Err(AssociationError::Generic("tried to add self".to_string())); + } + + // Only allow LegacyDelegated signatures on XIDs with a nonce of 0 + // Otherwise the client should use the regular wallet signature to create + if self.new_member_signature.signature_kind() == SignatureKind::LegacyDelegated { + if existing_state.xid != generate_xid(&existing_member_address, &0) { + return Err(AssociationError::LegacySignatureReuse); + } + } + + // Make sure that the signature type lines up with the role + if !allowed_signature_for_role( + &self.new_member_role, + &self.new_member_signature.signature_kind(), + ) { + return Err(AssociationError::SignatureNotAllowed( + self.new_member_role.clone(), + self.new_member_signature.signature_kind(), + )); + } + + let existing_member = existing_state.get(&existing_member_address); + + let existing_entity_id = match existing_member { + // If there is an existing member of the XID, use that member's ID + Some(member) => member.id.clone(), + None => { + // Check if it is a signature from the recovery address, which is allowed to add members + if existing_member_address.ne(recovery_address) { + return Err(AssociationError::MissingExistingMember); + } + // BUT, the recovery address has to be used with a real wallet signature, can't be delegated + if self.existing_member_signature.signature_kind() == SignatureKind::LegacyDelegated + { + return Err(AssociationError::LegacySignatureReuse); + } + // If it is a real wallet signature, then it is allowed to add members + recovery_address.clone() + } + }; + + let new_member = Entity::new( + self.new_member_role.clone(), + new_member_address, + Some(existing_entity_id), + ); + + println!( + "Adding new entity to state {:?} with hash {}", + &new_member, &association_hash + ); + + Ok(existing_state.add(new_member).mark_event_seen(self.hash())) + } + + fn hash(&self) -> String { + let inputs = format!( + "{}{:?}{}{}", + self.client_timestamp_ns, + self.new_member_role, + self.existing_member_signature.text(), + self.new_member_signature.text() + ); + sha256_string(inputs) + } +} + +pub struct RevokeAssociation { + pub client_timestamp_ns: u32, + pub recovery_address_signature: Box, + pub revoked_member: String, +} + +impl LogEntry for RevokeAssociation { + fn update_state( + &self, + maybe_existing_state: Option, + ) -> Result { + let existing_state = maybe_existing_state.ok_or(AssociationError::NotCreated)?; + + if self.recovery_address_signature.signature_kind() == SignatureKind::LegacyDelegated { + return Err(AssociationError::SignatureNotAllowed( + EntityRole::Address, + SignatureKind::LegacyDelegated, + )); + } + // Don't need to check for replay here since revocation is idempotent + let recovery_signer = self.recovery_address_signature.recover_signer()?; + // Make sure there is a recovery address set on the state + let state_recovery_address = existing_state.recovery_address.clone(); + + // Ensure this message is signed by the recovery address + if recovery_signer != state_recovery_address { + return Err(AssociationError::MissingExistingMember); + } + + let installations_to_remove: Vec = existing_state + .entities_by_parent(&self.revoked_member) + .into_iter() + // Only remove children if they are installations + .filter(|child| child.role == EntityRole::Installation) + .collect(); + + // Actually apply the revocation to the parent + let new_state = existing_state.remove(self.revoked_member.clone()); + + Ok(installations_to_remove + .iter() + .fold(new_state, |state, installation| { + state.remove(installation.id.clone()) + }) + .mark_event_seen(self.hash())) + } + + fn hash(&self) -> String { + let inputs = format!( + "{}{}{}", + self.client_timestamp_ns, + self.recovery_address_signature.text(), + self.revoked_member, + ); + sha256_string(inputs) + } +} + +pub enum AssociationEvent { + CreateXid(CreateXid), + AddAssociation(AddAssociation), + RevokeAssociation(RevokeAssociation), +} + +impl LogEntry for AssociationEvent { + fn update_state( + &self, + existing_state: Option, + ) -> Result { + match self { + AssociationEvent::CreateXid(event) => event.update_state(existing_state), + AssociationEvent::AddAssociation(event) => event.update_state(existing_state), + AssociationEvent::RevokeAssociation(event) => event.update_state(existing_state), + } + } + + fn hash(&self) -> String { + match self { + AssociationEvent::CreateXid(event) => event.hash(), + AssociationEvent::AddAssociation(event) => event.hash(), + AssociationEvent::RevokeAssociation(event) => event.hash(), + } + } +} + +// Ensure that the type of signature matches the new entity's role. +pub fn allowed_signature_for_role(role: &EntityRole, signature_kind: &SignatureKind) -> bool { + match role { + EntityRole::Address => match signature_kind { + SignatureKind::Erc191 => true, + SignatureKind::Erc1271 => true, + SignatureKind::InstallationKey => false, + SignatureKind::LegacyDelegated => true, + }, + EntityRole::Installation => match signature_kind { + SignatureKind::Erc191 => false, + SignatureKind::Erc1271 => false, + SignatureKind::InstallationKey => true, + SignatureKind::LegacyDelegated => false, + }, + } +} diff --git a/xmtp_id/src/associations/entity.rs b/xmtp_id/src/associations/entity.rs new file mode 100644 index 000000000..9862caf01 --- /dev/null +++ b/xmtp_id/src/associations/entity.rs @@ -0,0 +1,41 @@ +#[derive(Clone, Debug, PartialEq)] +pub enum EntityRole { + Installation, + Address, +} + +#[derive(Clone, Debug)] +pub struct Entity { + pub role: EntityRole, + pub id: String, + pub added_by_entity: Option, +} + +impl Entity { + pub fn new(role: EntityRole, id: String, added_by_entity: Option) -> Self { + Self { + role, + id, + added_by_entity, + } + } +} + +#[cfg(test)] +mod tests { + use crate::associations::test_utils; + + use super::*; + + use test_utils::rand_string; + + impl Default for Entity { + fn default() -> Self { + Self { + role: EntityRole::Address, + id: rand_string(), + added_by_entity: None, + } + } + } +} diff --git a/xmtp_id/src/associations/hashes.rs b/xmtp_id/src/associations/hashes.rs new file mode 100644 index 000000000..7aa582ba7 --- /dev/null +++ b/xmtp_id/src/associations/hashes.rs @@ -0,0 +1,12 @@ +use sha2::{Digest, Sha256}; + +pub fn sha256_string(input: String) -> String { + let mut hasher = Sha256::new(); + hasher.update(input.as_bytes()); + let result = hasher.finalize(); + format!("{:x}", result) +} + +pub fn generate_xid(account_address: &String, nonce: &u32) -> String { + sha256_string(format!("{}{}", account_address, nonce)) +} diff --git a/xmtp_id/src/associations/mod.rs b/xmtp_id/src/associations/mod.rs new file mode 100644 index 000000000..616aa3ae2 --- /dev/null +++ b/xmtp_id/src/associations/mod.rs @@ -0,0 +1,379 @@ +mod association_log; +mod entity; +mod hashes; +mod signature; +mod state; +#[cfg(test)] +mod test_utils; + +pub use self::association_log::*; +pub use self::entity::{Entity, EntityRole}; +pub use self::signature::{Signature, SignatureError, SignatureKind}; +pub use self::state::{AssociationState, StateError}; + +pub fn apply_updates( + initial_state: AssociationState, + association_events: Vec, +) -> AssociationState { + association_events + .iter() + .fold(initial_state, |state, update| { + match update.update_state(Some(state.clone())) { + Ok(new_state) => new_state, + Err(err) => { + println!("invalid entry {}", err); + state + } + } + }) +} + +pub fn get_state( + association_updates: Vec, +) -> Result { + association_updates + .iter() + .try_fold(None, |existing_state, update| { + match update.update_state(existing_state.clone()) { + Ok(new_state) => Ok(Some(new_state)), + Err(err) => { + println!("Invalid state update {}", err); + Err(err) + } + } + }) + .map(|v| v.unwrap()) +} + +#[cfg(test)] +mod tests { + use self::test_utils::{rand_string, rand_u32}; + + use super::*; + + struct MockSignature { + is_valid: bool, + signer_identity: String, + signature_kind: SignatureKind, + } + + impl MockSignature { + pub fn new_boxed( + is_valid: bool, + signer_identity: String, + signature_kind: SignatureKind, + ) -> Box { + Box::new(Self { + is_valid, + signer_identity, + signature_kind, + }) + } + } + + impl Signature for MockSignature { + fn signature_kind(&self) -> SignatureKind { + self.signature_kind.clone() + } + + fn recover_signer(&self) -> Result { + match self.is_valid { + true => Ok(self.signer_identity.clone()), + false => Err(SignatureError::Invalid), + } + } + + fn text(&self) -> String { + self.signer_identity.clone() + } + } + + impl Default for AddAssociation { + fn default() -> Self { + return Self { + client_timestamp_ns: rand_u32(), + new_member_role: EntityRole::Installation, + existing_member_signature: MockSignature::new_boxed( + true, + rand_string(), + SignatureKind::Erc191, + ), + new_member_signature: MockSignature::new_boxed( + true, + rand_string(), + SignatureKind::InstallationKey, + ), + }; + } + } + + impl Default for CreateXid { + fn default() -> Self { + let signer = rand_string(); + return Self { + nonce: rand_u32(), + account_address: signer.clone(), + initial_association: AddAssociation { + existing_member_signature: MockSignature::new_boxed( + true, + signer, + SignatureKind::Erc191, + ), + ..Default::default() + }, + }; + } + } + + impl Default for RevokeAssociation { + fn default() -> Self { + let signer = rand_string(); + return Self { + client_timestamp_ns: rand_u32(), + recovery_address_signature: MockSignature::new_boxed( + true, + signer, + SignatureKind::Erc191, + ), + revoked_member: rand_string(), + }; + } + } + + #[test] + fn test_create_and_add() { + let create_request = CreateXid::default(); + let account_address = create_request.account_address.clone(); + let state = get_state(vec![AssociationEvent::CreateXid(create_request)]).unwrap(); + assert_eq!(state.entities().len(), 2); + + let existing_entity = state.get(&account_address).unwrap(); + assert_eq!(existing_entity.id, account_address); + } + + #[test] + fn create_and_add_chained() { + let create_request = CreateXid::default(); + let state = get_state(vec![AssociationEvent::CreateXid(create_request)]).unwrap(); + assert_eq!(state.entities().len(), 2); + + let all_installations = state.entities_by_role(EntityRole::Installation); + let new_installation = all_installations.first().unwrap(); + + let update = AssociationEvent::AddAssociation(AddAssociation { + client_timestamp_ns: rand_u32(), + new_member_role: EntityRole::Address, + new_member_signature: MockSignature::new_boxed( + true, + rand_string(), + SignatureKind::Erc191, + ), + existing_member_signature: MockSignature::new_boxed( + true, + new_installation.id.clone(), + SignatureKind::InstallationKey, + ), + }); + + let new_state = apply_updates(state, vec![update]); + assert_eq!(new_state.entities().len(), 3); + } + + #[test] + fn create_from_legacy_key() { + let account_address = rand_string(); + let create_request = CreateXid { + nonce: 0, + account_address: account_address.clone(), + initial_association: AddAssociation { + client_timestamp_ns: rand_u32(), + existing_member_signature: MockSignature::new_boxed( + true, + account_address.clone(), + SignatureKind::LegacyDelegated, + ), + new_member_signature: MockSignature::new_boxed( + true, + rand_string(), + SignatureKind::InstallationKey, + ), + new_member_role: EntityRole::Installation, + }, + }; + + let initial_state = get_state(vec![AssociationEvent::CreateXid(create_request)]).unwrap(); + assert_eq!(initial_state.entities().len(), 2); + + let installation_keys = initial_state.entities_by_role(EntityRole::Installation); + assert_eq!(installation_keys.len(), 1); + } + + #[test] + fn reject_invalid_signature() { + let create_request = CreateXid { + initial_association: AddAssociation { + existing_member_signature: MockSignature::new_boxed( + false, + rand_string(), + SignatureKind::Erc191, + ), + ..Default::default() + }, + ..Default::default() + }; + + let state_result = get_state(vec![AssociationEvent::CreateXid(create_request)]); + assert_eq!(state_result.is_err(), true); + assert_eq!( + state_result.err().unwrap(), + AssociationError::Signature(SignatureError::Invalid) + ); + } + + #[test] + fn reject_if_signer_not_existing_member() { + let create_request = CreateXid::default(); + let state = get_state(vec![AssociationEvent::CreateXid(create_request)]).unwrap(); + + // The default here will create an AddAssociation from a random wallet + let update = AssociationEvent::AddAssociation(AddAssociation { + ..Default::default() + }); + + let new_state_result = update.update_state(Some(state)); + assert_eq!( + new_state_result.err(), + Some(AssociationError::MissingExistingMember) + ) + } + + #[test] + fn test_revoke_wallet() { + let create_request = CreateXid::default(); + let initial_wallet = create_request.account_address.clone(); + let state = get_state(vec![AssociationEvent::CreateXid(create_request)]).unwrap(); + + // This update should revoke the initial wallet and the one installation that is associated to it + let update = AssociationEvent::RevokeAssociation(RevokeAssociation { + recovery_address_signature: MockSignature::new_boxed( + true, + initial_wallet.clone(), + SignatureKind::Erc191, + ), + revoked_member: initial_wallet.clone(), + ..Default::default() + }); + + let new_state = update.update_state(Some(state)).unwrap(); + assert_eq!(new_state.entities().len(), 0); + } + + #[test] + fn test_revoke_installation() { + let create_request = CreateXid::default(); + let state = get_state(vec![AssociationEvent::CreateXid(create_request)]).unwrap(); + let recovery_address = state.recovery_address.clone(); + let all_installations = state.entities_by_role(EntityRole::Installation); + let installation_to_revoke = all_installations.first().unwrap(); + + let update = RevokeAssociation { + recovery_address_signature: MockSignature::new_boxed( + true, + recovery_address, + SignatureKind::Erc191, + ), + revoked_member: installation_to_revoke.id.clone(), + ..Default::default() + }; + + let new_state = update.update_state(Some(state)).unwrap(); + assert_eq!(new_state.entities().len(), 1); + } + + #[test] + fn test_replay_detection() { + let create_request = CreateXid::default(); + let original_nonce = create_request.initial_association.client_timestamp_ns; + let original_installation_id = create_request + .initial_association + .new_member_signature + .recover_signer() + .unwrap(); + + let state = get_state(vec![AssociationEvent::CreateXid(create_request)]).unwrap(); + + let recovery_address = state.recovery_address.clone(); + let all_installations = state.entities_by_role(EntityRole::Installation); + let installation_to_revoke = all_installations.first().unwrap(); + + let update = RevokeAssociation { + recovery_address_signature: MockSignature::new_boxed( + true, + recovery_address.clone(), + SignatureKind::Erc191, + ), + revoked_member: installation_to_revoke.id.clone(), + ..Default::default() + }; + + let new_state = update.update_state(Some(state)).unwrap(); + assert_eq!(new_state.entities().len(), 1); + + let attempt_to_replay = AddAssociation { + client_timestamp_ns: original_nonce, + existing_member_signature: MockSignature::new_boxed( + true, + recovery_address, + SignatureKind::Erc191, + ), + new_member_signature: MockSignature::new_boxed( + true, + original_installation_id, + SignatureKind::InstallationKey, + ), + ..Default::default() + }; + + let replay_result = attempt_to_replay.update_state(Some(new_state)); + assert_eq!(replay_result.is_err(), true); + assert_eq!(replay_result.err().unwrap(), AssociationError::Replay) + } + + #[test] + fn test_add_from_recovery_address() { + let create_request = CreateXid::default(); + let initial_wallet_address = create_request.account_address.clone(); + let revocation = RevokeAssociation { + recovery_address_signature: MockSignature::new_boxed( + true, + initial_wallet_address.clone(), + SignatureKind::Erc191, + ), + revoked_member: initial_wallet_address.clone(), + ..Default::default() + }; + + let state = get_state(vec![ + AssociationEvent::CreateXid(create_request), + AssociationEvent::RevokeAssociation(revocation), + ]) + .unwrap(); + + // Ensure that the initial wallet isn't in there. The recovery address has not changed, however. + assert_eq!(state.get(&initial_wallet_address).is_none(), true); + + let add_from_recovery = AddAssociation { + client_timestamp_ns: rand_u32(), + existing_member_signature: MockSignature::new_boxed( + true, + initial_wallet_address, + SignatureKind::Erc191, + ), + ..Default::default() + }; + + add_from_recovery + .update_state(Some(state)) + .expect("should be allowed because we are using the recovery address"); + } +} diff --git a/xmtp_id/src/associations/signature.rs b/xmtp_id/src/associations/signature.rs new file mode 100644 index 000000000..3600ccfc6 --- /dev/null +++ b/xmtp_id/src/associations/signature.rs @@ -0,0 +1,22 @@ +use thiserror::Error; + +#[derive(Debug, Error, PartialEq)] +pub enum SignatureError { + #[error("Signature validation failed")] + Invalid, +} + +#[derive(Clone, Debug, PartialEq)] +pub enum SignatureKind { + // We might want to have some sort of LegacyErc191 Signature Kind for the `CreateIdentity` signatures only + Erc191, + Erc1271, + InstallationKey, + LegacyDelegated, +} + +pub trait Signature { + fn recover_signer(&self) -> Result; + fn signature_kind(&self) -> SignatureKind; + fn text(&self) -> String; +} diff --git a/xmtp_id/src/associations/state.rs b/xmtp_id/src/associations/state.rs new file mode 100644 index 000000000..efbc5c94c --- /dev/null +++ b/xmtp_id/src/associations/state.rs @@ -0,0 +1,110 @@ +use std::collections::{HashMap, HashSet}; + +use thiserror::Error; + +use super::{entity::Entity, hashes::generate_xid, EntityRole}; + +#[derive(Debug, Error, PartialEq)] +pub enum StateError { + #[error("Not found")] + NotFound, + #[error("Replay detected")] + ReplayDetected, +} + +#[derive(Clone, Debug)] +pub struct AssociationState { + pub xid: String, + pub current_entities: HashMap, + pub recovery_address: String, + pub seen_events: HashSet, +} + +impl AssociationState { + pub fn add(&self, entity: Entity) -> Self { + let mut new_state = self.clone(); + let _ = new_state.current_entities.insert(entity.id.clone(), entity); + + new_state + } + + pub fn set_recovery_address(&self, recovery_address: String) -> Self { + let mut new_state = self.clone(); + new_state.recovery_address = recovery_address; + + new_state + } + + pub fn get(&self, id: &String) -> Option { + self.current_entities.get(id).map(|e| e.clone()) + } + + pub fn mark_event_seen(&self, event_hash: String) -> Self { + let mut new_state = self.clone(); + new_state.seen_events.insert(event_hash); + + new_state + } + + pub fn has_seen(&self, event_hash: &String) -> bool { + self.seen_events.contains(event_hash) + } + + pub fn remove(&self, entity_id: String) -> Self { + let mut new_state = self.clone(); + let _ = new_state.current_entities.remove(&entity_id); + + new_state + } + + pub fn entities(&self) -> Vec { + self.current_entities.values().cloned().collect() + } + + pub fn entities_by_parent(&self, parent_id: &String) -> Vec { + self.current_entities + .values() + .filter(|e| e.added_by_entity == Some(parent_id.clone())) + .cloned() + .collect() + } + + pub fn entities_by_role(&self, role: EntityRole) -> Vec { + self.current_entities + .values() + .filter(|e| e.role == role) + .cloned() + .collect() + } + + pub fn new(account_address: String, nonce: u32) -> Self { + let xid = generate_xid(&account_address, &nonce); + let new_entity = Entity::new(EntityRole::Address, account_address.clone(), None); + Self { + current_entities: { + let mut entities = HashMap::new(); + entities.insert(account_address.clone(), new_entity); + entities + }, + seen_events: HashSet::new(), + recovery_address: account_address, + xid, + } + } +} + +#[cfg(test)] +mod tests { + use crate::associations::test_utils::rand_string; + + use super::*; + + #[test] + fn can_add_remove() { + let starting_state = AssociationState::new(rand_string(), 0); + let new_entity = Entity::default(); + let with_add = starting_state.add(new_entity.clone()); + assert!(with_add.get(&new_entity.id).is_some()); + assert!(starting_state.get(&new_entity.id).is_none()); + } +} diff --git a/xmtp_id/src/associations/test_utils.rs b/xmtp_id/src/associations/test_utils.rs new file mode 100644 index 000000000..98d94292a --- /dev/null +++ b/xmtp_id/src/associations/test_utils.rs @@ -0,0 +1,15 @@ +use rand::{distributions::Alphanumeric, Rng}; + +pub fn rand_string() -> String { + let v: String = rand::thread_rng() + .sample_iter(&Alphanumeric) + .take(32) + .map(char::from) + .collect(); + + v +} + +pub fn rand_u32() -> u32 { + rand::thread_rng().gen() +} diff --git a/xmtp_id/src/lib.rs b/xmtp_id/src/lib.rs index 5eee16f19..6d03c5270 100644 --- a/xmtp_id/src/lib.rs +++ b/xmtp_id/src/lib.rs @@ -1,3 +1,4 @@ +pub mod associations; pub mod credential_verifier; pub mod verified_key_package;