From dfe9edf33311676c4932acf4465c97b71a10c26f Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Fri, 29 Mar 2024 15:34:07 -0700 Subject: [PATCH 1/5] Add toy implementation of association rollup --- Cargo.lock | 2 + bindings_ffi/Cargo.lock | 4 +- xmtp_id/Cargo.toml | 2 + xmtp_id/src/associations/entity.rs | 42 ++ xmtp_id/src/associations/mod.rs | 704 +++++++++++++++++++++++++ xmtp_id/src/associations/state.rs | 119 +++++ xmtp_id/src/associations/test_utils.rs | 15 + xmtp_id/src/lib.rs | 1 + 8 files changed, 886 insertions(+), 3 deletions(-) create mode 100644 xmtp_id/src/associations/entity.rs create mode 100644 xmtp_id/src/associations/mod.rs create mode 100644 xmtp_id/src/associations/state.rs create mode 100644 xmtp_id/src/associations/test_utils.rs diff --git a/Cargo.lock b/Cargo.lock index 13cc1e670..89c6831e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5773,7 +5773,9 @@ dependencies = [ "openmls_basic_credential", "openmls_traits", "prost 0.12.3", + "rand", "serde", + "sha2 0.10.8", "thiserror", "tls_codec", "tracing", diff --git a/bindings_ffi/Cargo.lock b/bindings_ffi/Cargo.lock index f5f08b14c..9ef50517a 100644 --- a/bindings_ffi/Cargo.lock +++ b/bindings_ffi/Cargo.lock @@ -5353,7 +5353,6 @@ dependencies = [ name = "xmtp_mls" version = "0.1.0" dependencies = [ - "anyhow", "async-trait", "chrono", "diesel", @@ -5376,8 +5375,7 @@ dependencies = [ "thiserror", "tls_codec", "tokio", - "toml 0.7.8", - "tracing", + "toml 0.8.8", "xmtp_cryptography", "xmtp_proto", "xmtp_v2", diff --git a/xmtp_id/Cargo.toml b/xmtp_id/Cargo.toml index bebdac9a2..7f073bafa 100644 --- a/xmtp_id/Cargo.toml +++ b/xmtp_id/Cargo.toml @@ -19,3 +19,5 @@ prost.workspace = true tls_codec.workspace = true chrono.workspace = true serde.workspace = true +sha2 = "0.10.8" +rand.workspace = true diff --git a/xmtp_id/src/associations/entity.rs b/xmtp_id/src/associations/entity.rs new file mode 100644 index 000000000..8c9c02ba3 --- /dev/null +++ b/xmtp_id/src/associations/entity.rs @@ -0,0 +1,42 @@ +#[derive(Clone, Debug, PartialEq)] +pub enum EntityRole { + Installation, + Address, + LegacyKey, +} + +#[derive(Clone, Debug)] +pub struct Entity { + pub role: EntityRole, + pub id: String, + pub is_revoked: bool, +} + +impl Entity { + pub fn new(role: EntityRole, id: String, is_revoked: bool) -> Self { + Self { + role, + id, + is_revoked, + } + } +} + +#[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(), + is_revoked: false, + } + } + } +} diff --git a/xmtp_id/src/associations/mod.rs b/xmtp_id/src/associations/mod.rs new file mode 100644 index 000000000..1e0fd8736 --- /dev/null +++ b/xmtp_id/src/associations/mod.rs @@ -0,0 +1,704 @@ +mod entity; +mod state; +#[cfg(test)] +mod test_utils; + +pub use self::entity::{Entity, EntityRole}; +pub use self::state::{AssociationState, StateError}; +use sha2::{Digest, Sha256}; + +use thiserror::Error; + +const ALLOWED_CREATE_ENTITY_ROLES: [EntityRole; 2] = [EntityRole::LegacyKey, EntityRole::Address]; + +#[derive(Debug, Error)] +pub enum SignatureError { + #[error("Signature validation failed")] + Invalid, +} + +#[derive(Debug, Error)] +pub enum AssociationError { + #[error("Error creating association {0}")] + Generic(String), + #[error("Multiple create operations detect")] + MultipleCreate, + #[error("Signature validation failed {0}")] + Signature(#[from] SignatureError), + #[error("State update failed")] + StateError(#[from] StateError), + #[error("Missing existing member")] + MissingExistingMember, + #[error("Signature not allowed for role {0:?} {1:?}")] + SignatureNotAllowed(EntityRole, SignatureKind), + #[error("Added by revoked member")] + AddedByRevokedMember, + #[error("Replay detected")] + Replay, + #[error("No recovery address")] + NoRecoveryAddress, +} + +#[derive(Clone, Debug)] +pub enum SignatureKind { + Erc191, + Erc1271, + InstallationKey, + LegacyKey, +} + +pub trait Signature { + fn recover_signer(&self) -> Result; + fn signature_kind(&self) -> SignatureKind; + fn text(&self) -> String; +} + +pub trait LogEntry { + fn update_state( + &self, + existing_state: AssociationState, + ) -> Result; + fn hash(&self) -> String; +} + +pub struct CreateXidEntry { + pub nonce: u32, + pub signature: Box, + pub recovery_address: String, + pub entity_role: EntityRole, +} + +impl LogEntry for CreateXidEntry { + fn update_state( + &self, + existing_state: AssociationState, + ) -> Result { + // Verify that the existing state is empty + if !existing_state.entities().is_empty() { + return Err(AssociationError::MultipleCreate); + } + + // This verifies that the signature is valid + let signer_address = self.signature.recover_signer()?; + if !ALLOWED_CREATE_ENTITY_ROLES.contains(&self.entity_role) { + return Err(AssociationError::Generic("invalid entity role".to_string())); + } + + let signature_kind = self.signature.signature_kind(); + if !allowed_signature_for_role(&self.entity_role, &signature_kind) { + return Err(AssociationError::SignatureNotAllowed( + self.entity_role.clone(), + signature_kind, + )); + } + + let entity = Entity::new(self.entity_role.clone(), signer_address, false); + Ok(existing_state + .set_recovery_address(self.recovery_address.clone()) + .add(entity, 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.signature.text(), + self.recovery_address + ); + + sha256_string(inputs) + } +} + +pub struct AddAssociationEntry { + pub nonce: u32, + pub new_member_role: EntityRole, + pub existing_member_signature: Box, + pub new_member_signature: Box, +} + +impl AddAssociationEntry { + pub fn new_member_address(&self) -> String { + self.new_member_signature.recover_signer().unwrap() + } +} + +impl LogEntry for AddAssociationEntry { + fn update_state( + &self, + existing_state: AssociationState, + ) -> Result { + let association_hash = self.hash(); + if existing_state.has_seen(&association_hash) { + return Err(AssociationError::Replay); + } + + // Recovery address has to be set + if existing_state.recovery_address.is_none() { + return Err(AssociationError::NoRecoveryAddress); + } + + let new_member_address = self.new_member_signature.recover_signer()?; + let existing_member_address = self.existing_member_signature.recover_signer()?; + if new_member_address == existing_member_address { + return Err(AssociationError::Generic("tried to add self".to_string())); + } + + // Get the current version of the entity that added this new entry. If it has been revoked and added back, it will now be unrevoked + let existing_entity = existing_state + .get(&existing_member_address) + .ok_or(AssociationError::MissingExistingMember)?; + + if existing_entity.is_revoked { + // The entity that added this member is currently revoked. Check if this particular association is allowlisted + if !existing_state + .allowlisted_association_hashes + .contains(&association_hash) + { + return Err(AssociationError::AddedByRevokedMember); + } + } + + // 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(), + )); + } + + // Check to see if the new member was revoked + let is_new_member_revoked = existing_state.was_association_revoked(&association_hash); + let new_member = Entity::new( + self.new_member_role.clone(), + new_member_address, + is_new_member_revoked, + ); + + println!( + "Adding new entity to state {:?} with hash {}", + &new_member, &association_hash + ); + + Ok(existing_state.add(new_member, association_hash)?) + } + + fn hash(&self) -> String { + let inputs = format!( + "{}{:?}{}{}", + self.nonce, + self.new_member_role, + self.existing_member_signature.text(), + self.new_member_signature.text() + ); + sha256_string(inputs) + } +} + +pub struct RevokeAssociationEntry { + pub nonce: u32, + pub recovery_address_signature: Box, + pub revoked_association_hash: String, + pub allowed_child_hashes: Vec, +} + +impl LogEntry for RevokeAssociationEntry { + fn update_state( + &self, + existing_state: AssociationState, + ) -> Result { + // 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() + .ok_or(AssociationError::NoRecoveryAddress)?; + + // Ensure this message is signed by the recovery address + if recovery_signer != state_recovery_address { + return Err(AssociationError::MissingExistingMember); + } + + // Actually apply the revocation + Ok(existing_state.apply_revocation( + self.revoked_association_hash.clone(), + self.allowed_child_hashes.clone(), + )) + } + + fn hash(&self) -> String { + let inputs = format!( + "{}{}{}{}", + self.nonce, + self.recovery_address_signature.text(), + self.revoked_association_hash, + self.allowed_child_hashes.join(",") + ); + sha256_string(inputs) + } +} + +pub struct ChangeRecoveryAddressEntry { + pub nonce: u32, + pub recovery_address_signature: Box, + pub new_recovery_address: String, +} + +pub enum RecoveryLogEntry { + CreateXid(CreateXidEntry), + RevokeAssociation(RevokeAssociationEntry), +} + +impl LogEntry for RecoveryLogEntry { + fn update_state( + &self, + existing_state: AssociationState, + ) -> Result { + match self { + RecoveryLogEntry::CreateXid(create_xid) => create_xid.update_state(existing_state), + RecoveryLogEntry::RevokeAssociation(revoke_association) => { + revoke_association.update_state(existing_state) + } + } + } + + fn hash(&self) -> String { + match self { + RecoveryLogEntry::CreateXid(create_xid) => create_xid.hash(), + RecoveryLogEntry::RevokeAssociation(revoke_association) => revoke_association.hash(), + } + } +} + +pub fn apply_updates( + initial_state: AssociationState, + associations: Vec, +) -> AssociationState { + associations.iter().fold(initial_state, |state, update| { + match update.update_state(state.clone()) { + Ok(new_state) => new_state, + Err(err) => { + println!("invalid entry {}", err); + state + } + } + }) +} + +pub fn get_initial_state(recovery_log: Vec) -> AssociationState { + recovery_log + .iter() + .fold(AssociationState::new(), |state, update| { + match update.update_state(state.clone()) { + Ok(new_state) => new_state, + Err(err) => { + println!("invalid entry {}", err); + state + } + } + }) +} + +pub fn get_state( + recovery_log: Vec, + association_updates: Vec, +) -> AssociationState { + let state = get_initial_state(recovery_log); + println!("Initial state {:?}", state); + apply_updates(state, association_updates) +} + +fn sha256_string(input: String) -> String { + let mut hasher = Sha256::new(); + hasher.update(input.as_bytes()); + let result = hasher.finalize(); + format!("{:x}", result) +} + +// 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::LegacyKey => false, + }, + EntityRole::LegacyKey => match signature_kind { + SignatureKind::Erc191 => false, + SignatureKind::Erc1271 => false, + SignatureKind::InstallationKey => false, + SignatureKind::LegacyKey => true, + }, + EntityRole::Installation => match signature_kind { + SignatureKind::Erc191 => false, + SignatureKind::Erc1271 => false, + SignatureKind::InstallationKey => true, + SignatureKind::LegacyKey => false, + }, + } +} + +/** + * Revocation properties + * 1. Revoking an association will mark the entity added as revoked + * 2. Revoking an association will prevent new associations from being created with an `existing_entity_signature` from the revoked entity + * 3. Entities created with an `existing_entity_signature` of a revoked entity can be protected from revocation if they were added before the revocation + * 4. Revoked entities can be re-added with a new signature, so long as a new nonce is included in the signature + * 5. A revocation payload can be added to a subset of the association log. When this happens, all entities present in the subset will have the same revocation status that they have in the full log. + */ + +#[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 Default for AddAssociationEntry { + fn default() -> Self { + return Self { + nonce: rand_u32(), + new_member_role: EntityRole::Address, + existing_member_signature: MockSignature::new_boxed( + true, + rand_string(), + SignatureKind::Erc191, + ), + new_member_signature: MockSignature::new_boxed( + true, + rand_string(), + SignatureKind::Erc191, + ), + }; + } + } + + impl Default for CreateXidEntry { + fn default() -> Self { + let signer = rand_string(); + return Self { + nonce: rand_u32(), + signature: MockSignature::new_boxed(true, signer.clone(), SignatureKind::Erc191), + recovery_address: signer, + entity_role: EntityRole::Address, + }; + } + } + + impl Default for RevokeAssociationEntry { + fn default() -> Self { + let signer = rand_string(); + return Self { + nonce: rand_u32(), + recovery_address_signature: MockSignature::new_boxed( + true, + signer, + SignatureKind::Erc191, + ), + revoked_association_hash: rand_string(), + allowed_child_hashes: vec![], + }; + } + } + + 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() + } + } + + fn init_recovery_log() -> (Vec, String) { + let create_request = CreateXidEntry::default(); + let creator_address = create_request.signature.recover_signer().unwrap(); + let entries = vec![RecoveryLogEntry::CreateXid(create_request)]; + + (entries, creator_address) + } + + #[test] + fn test_create_and_add() { + let create_request = CreateXidEntry::default(); + let creator_address = create_request.signature.recover_signer().unwrap(); + let recovery_log = vec![RecoveryLogEntry::CreateXid(create_request)]; + let mut state = get_state(recovery_log, vec![]); + assert_eq!(state.entities().len(), 1); + + let add_installation_entry = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + creator_address, + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + state = apply_updates(state, vec![add_installation_entry]); + assert_eq!(state.entities().len(), 2); + } + + #[test] + fn create_and_add_chained() { + let (recovery_log, creator_address) = init_recovery_log(); + let add_first_association = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + creator_address, + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + let first_association_address = add_first_association + .new_member_signature + .recover_signer() + .unwrap(); + + let add_second_association = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + first_association_address.clone(), + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + + let state = get_state( + recovery_log, + vec![add_first_association, add_second_association], + ); + + assert_eq!(state.entities().len(), 3); + assert_eq!( + state.get(&first_association_address).unwrap().is_revoked, + false + ); + assert_eq!( + state.get(&first_association_address).unwrap().id, + first_association_address + ); + } + + #[test] + fn add_from_revoked() { + let (mut recovery_log, creator_address) = init_recovery_log(); + let add_association = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + creator_address.clone(), + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + + recovery_log.push(RecoveryLogEntry::RevokeAssociation( + RevokeAssociationEntry { + recovery_address_signature: MockSignature::new_boxed( + true, + // Creator address is the recovery address, so this is valid + creator_address, + SignatureKind::Erc191, + ), + revoked_association_hash: add_association.hash(), + // Not setting any allowed children here, since this doesn't have any + ..Default::default() + }, + )); + + let add_another_association = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + add_association.new_member_address(), + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + let second_new_member_address = add_another_association.new_member_address(); + + let state = get_state(recovery_log, vec![add_association, add_another_association]); + assert_eq!(state.entities().len(), 2); + assert!(state.get(&second_new_member_address).is_none()) + } + + #[test] + fn add_from_re_added() { + let (mut recovery_log, creator_address) = init_recovery_log(); + + let add_association = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + creator_address.clone(), + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + + let new_member_address = add_association.new_member_address(); + + recovery_log.push(RecoveryLogEntry::RevokeAssociation( + RevokeAssociationEntry { + recovery_address_signature: MockSignature::new_boxed( + true, + // Creator address is the recovery address, so this is valid + creator_address.clone(), + SignatureKind::Erc191, + ), + revoked_association_hash: add_association.hash(), + // Not setting any allowed children here, since this doesn't have any + ..Default::default() + }, + )); + + let add_same_member_back = AddAssociationEntry { + nonce: rand_u32(), + existing_member_signature: MockSignature::new_boxed( + true, + creator_address.clone(), + SignatureKind::Erc191, + ), + new_member_signature: MockSignature::new_boxed( + true, + new_member_address.clone(), + SignatureKind::Erc191, + ), + ..Default::default() + }; + + let state = get_state(recovery_log, vec![add_association, add_same_member_back]); + assert_eq!(state.get(&new_member_address).unwrap().is_revoked, false) + } + + #[test] + fn protect_children_from_revocation() { + let (mut recovery_log, creator_address) = init_recovery_log(); + + let add_association = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + creator_address.clone(), + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + + let add_child = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + add_association.new_member_address(), + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + + let add_grandchild = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + add_child.new_member_address(), + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + + recovery_log.push(RecoveryLogEntry::RevokeAssociation( + RevokeAssociationEntry { + recovery_address_signature: MockSignature::new_boxed( + true, + // Creator address is the recovery address, so this is valid + creator_address.clone(), + SignatureKind::Erc191, + ), + revoked_association_hash: add_association.hash(), + allowed_child_hashes: vec![add_child.hash()], + // Not setting any allowed children here, since this doesn't have any + ..Default::default() + }, + )); + + let first_member_address = add_association.new_member_address(); + let first_child_address = add_child.new_member_address(); + let grandchild_address = add_grandchild.new_member_address(); + + let state = get_state( + recovery_log, + vec![add_association, add_child, add_grandchild], + ); + assert_eq!(state.get(&first_member_address).unwrap().is_revoked, true); + assert_eq!(state.get(&first_child_address).unwrap().is_revoked, false); + assert_eq!(state.get(&grandchild_address).unwrap().is_revoked, false); + } + + #[test] + fn fail_if_ancestor_missing() { + let (recovery_log, creator_address) = init_recovery_log(); + + let add_association = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + creator_address.clone(), + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + + let add_child = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + add_association.new_member_address(), + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + + let add_grandchild = AddAssociationEntry { + existing_member_signature: MockSignature::new_boxed( + true, + add_child.new_member_address(), + SignatureKind::Erc191, + ), + ..AddAssociationEntry::default() + }; + + let grandchild_address = add_grandchild.new_member_address(); + + let state = get_state( + recovery_log, + // Deliberately omitting the add_child, which is necessary here + vec![add_association, add_grandchild], + ); + + assert_eq!(state.get(&grandchild_address).is_none(), true); + } +} diff --git a/xmtp_id/src/associations/state.rs b/xmtp_id/src/associations/state.rs new file mode 100644 index 000000000..30fbbfa4a --- /dev/null +++ b/xmtp_id/src/associations/state.rs @@ -0,0 +1,119 @@ +use std::collections::{HashMap, HashSet}; + +use thiserror::Error; + +use super::{entity::Entity, EntityRole}; + +#[derive(Debug, Error)] +pub enum StateError { + #[error("Not found")] + NotFound, + #[error("Replay detected")] + ReplayDetected, +} + +#[derive(Clone, Debug)] +pub struct AssociationState { + pub current_entities: HashMap, + // Stores the entity as it was at the time it was added + pub entities_by_event: HashMap, + pub revoked_association_hashes: HashSet, + pub allowlisted_association_hashes: HashSet, + pub recovery_address: Option, +} + +impl AssociationState { + pub fn add(&self, entity: Entity, event_hash: String) -> Result { + self.replay_check(&event_hash)?; + let mut new_state = self.clone(); + let _ = new_state + .entities_by_event + .insert(event_hash, entity.clone()); + let _ = new_state.current_entities.insert(entity.id.clone(), entity); + + Ok(new_state) + } + + pub fn set_recovery_address(&self, recovery_address: String) -> Self { + let mut new_state = self.clone(); + new_state.recovery_address = Some(recovery_address); + + new_state + } + + pub fn get(&self, id: &String) -> Option { + self.current_entities.get(id).map(|e| e.clone()) + } + + pub fn has_seen(&self, event_hash: &String) -> bool { + self.entities_by_event.contains_key(event_hash) + } + + fn replay_check(&self, event_hash: &String) -> Result<(), StateError> { + if self.has_seen(event_hash) { + return Err(StateError::ReplayDetected); + } + + Ok(()) + } + + pub fn apply_revocation( + &self, + revoked_association_hash: String, + allowlisted_association_hashes: Vec, + ) -> Self { + let mut new_state = self.clone(); + let _ = new_state + .revoked_association_hashes + .insert(revoked_association_hash); + new_state + .allowlisted_association_hashes + .extend(allowlisted_association_hashes); + + new_state + } + + pub fn was_association_revoked(&self, association_hash: &String) -> bool { + self.revoked_association_hashes.contains(association_hash) + } + + pub fn entities(&self) -> Vec { + self.current_entities.values().map(|e| e.clone()).collect() + } + + pub fn entities_by_role(&self, role: EntityRole) -> Vec { + self.current_entities + .values() + .filter(|e| e.role == role) + .map(|e| e.clone()) + .collect() + } + + pub fn new() -> Self { + Self { + current_entities: HashMap::new(), + entities_by_event: HashMap::new(), + revoked_association_hashes: HashSet::new(), + allowlisted_association_hashes: HashSet::new(), + recovery_address: None, + } + } +} + +#[cfg(test)] +mod tests { + use crate::associations::test_utils::rand_string; + + use super::*; + + #[test] + fn can_add_remove() { + let starting_state = AssociationState::new(); + let new_entity = Entity::default(); + let with_add = starting_state + .add(new_entity.clone(), rand_string()) + .unwrap(); + 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 be029b8a2..3313d4c62 100644 --- a/xmtp_id/src/lib.rs +++ b/xmtp_id/src/lib.rs @@ -1,3 +1,4 @@ +pub mod associations; pub mod error; use std::sync::RwLock; From 7b2c947b83638302e047644d991d90dee111190f Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Fri, 29 Mar 2024 16:29:08 -0700 Subject: [PATCH 2/5] Lint --- xmtp_id/src/associations/state.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/xmtp_id/src/associations/state.rs b/xmtp_id/src/associations/state.rs index 30fbbfa4a..7c5ad4914 100644 --- a/xmtp_id/src/associations/state.rs +++ b/xmtp_id/src/associations/state.rs @@ -78,14 +78,14 @@ impl AssociationState { } pub fn entities(&self) -> Vec { - self.current_entities.values().map(|e| e.clone()).collect() + self.current_entities.values().cloned().collect() } pub fn entities_by_role(&self, role: EntityRole) -> Vec { self.current_entities .values() .filter(|e| e.role == role) - .map(|e| e.clone()) + .cloned() .collect() } @@ -100,6 +100,12 @@ impl AssociationState { } } +impl Default for AssociationState { + fn default() -> Self { + Self::new() + } +} + #[cfg(test)] mod tests { use crate::associations::test_utils::rand_string; From aa47584e4ec15de8d94b1ea64be6bd943d7a6076 Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Tue, 2 Apr 2024 16:34:29 -0700 Subject: [PATCH 3/5] Use new association state model --- xmtp_id/src/associations/association_log.rs | 255 +++++++ xmtp_id/src/associations/entity.rs | 8 +- xmtp_id/src/associations/hashes.rs | 12 + xmtp_id/src/associations/mod.rs | 743 +++++--------------- xmtp_id/src/associations/signature.rs | 21 + xmtp_id/src/associations/state.rs | 91 +-- 6 files changed, 520 insertions(+), 610 deletions(-) create mode 100644 xmtp_id/src/associations/association_log.rs create mode 100644 xmtp_id/src/associations/hashes.rs create mode 100644 xmtp_id/src/associations/signature.rs diff --git a/xmtp_id/src/associations/association_log.rs b/xmtp_id/src/associations/association_log.rs new file mode 100644 index 000000000..1852c222b --- /dev/null +++ b/xmtp_id/src/associations/association_log.rs @@ -0,0 +1,255 @@ +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)?; + + 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()?; + if new_member_address == existing_member_address { + return Err(AssociationError::Generic("tried to add self".to_string())); + } + + if self.new_member_role == EntityRole::LegacyKey { + if existing_state.xid != generate_xid(&existing_member_address, &0) { + return Err(AssociationError::LegacySignatureReuse); + } + } + + // Get the current version of the entity that added this new entry. If it has been revoked and added back, it will now be unrevoked + let existing_entity = existing_state + .get(&existing_member_address) + .ok_or(AssociationError::MissingExistingMember)?; + + // 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 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)?; + // 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::LegacyKey => false, + }, + EntityRole::LegacyKey => match signature_kind { + SignatureKind::Erc191 => false, + SignatureKind::Erc1271 => false, + SignatureKind::InstallationKey => false, + SignatureKind::LegacyKey => true, + }, + EntityRole::Installation => match signature_kind { + SignatureKind::Erc191 => false, + SignatureKind::Erc1271 => false, + SignatureKind::InstallationKey => true, + SignatureKind::LegacyKey => false, + }, + } +} diff --git a/xmtp_id/src/associations/entity.rs b/xmtp_id/src/associations/entity.rs index 8c9c02ba3..8e687f253 100644 --- a/xmtp_id/src/associations/entity.rs +++ b/xmtp_id/src/associations/entity.rs @@ -9,15 +9,15 @@ pub enum EntityRole { pub struct Entity { pub role: EntityRole, pub id: String, - pub is_revoked: bool, + pub added_by_entity: Option, } impl Entity { - pub fn new(role: EntityRole, id: String, is_revoked: bool) -> Self { + pub fn new(role: EntityRole, id: String, added_by_entity: Option) -> Self { Self { role, id, - is_revoked, + added_by_entity, } } } @@ -35,7 +35,7 @@ mod tests { Self { role: EntityRole::Address, id: rand_string(), - is_revoked: false, + 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 index 1e0fd8736..be0a3e8fd 100644 --- a/xmtp_id/src/associations/mod.rs +++ b/xmtp_id/src/associations/mod.rs @@ -1,300 +1,24 @@ +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}; -use sha2::{Digest, Sha256}; - -use thiserror::Error; - -const ALLOWED_CREATE_ENTITY_ROLES: [EntityRole; 2] = [EntityRole::LegacyKey, EntityRole::Address]; - -#[derive(Debug, Error)] -pub enum SignatureError { - #[error("Signature validation failed")] - Invalid, -} - -#[derive(Debug, Error)] -pub enum AssociationError { - #[error("Error creating association {0}")] - Generic(String), - #[error("Multiple create operations detect")] - MultipleCreate, - #[error("Signature validation failed {0}")] - Signature(#[from] SignatureError), - #[error("State update failed")] - StateError(#[from] StateError), - #[error("Missing existing member")] - MissingExistingMember, - #[error("Signature not allowed for role {0:?} {1:?}")] - SignatureNotAllowed(EntityRole, SignatureKind), - #[error("Added by revoked member")] - AddedByRevokedMember, - #[error("Replay detected")] - Replay, - #[error("No recovery address")] - NoRecoveryAddress, -} - -#[derive(Clone, Debug)] -pub enum SignatureKind { - Erc191, - Erc1271, - InstallationKey, - LegacyKey, -} - -pub trait Signature { - fn recover_signer(&self) -> Result; - fn signature_kind(&self) -> SignatureKind; - fn text(&self) -> String; -} - -pub trait LogEntry { - fn update_state( - &self, - existing_state: AssociationState, - ) -> Result; - fn hash(&self) -> String; -} - -pub struct CreateXidEntry { - pub nonce: u32, - pub signature: Box, - pub recovery_address: String, - pub entity_role: EntityRole, -} - -impl LogEntry for CreateXidEntry { - fn update_state( - &self, - existing_state: AssociationState, - ) -> Result { - // Verify that the existing state is empty - if !existing_state.entities().is_empty() { - return Err(AssociationError::MultipleCreate); - } - - // This verifies that the signature is valid - let signer_address = self.signature.recover_signer()?; - if !ALLOWED_CREATE_ENTITY_ROLES.contains(&self.entity_role) { - return Err(AssociationError::Generic("invalid entity role".to_string())); - } - - let signature_kind = self.signature.signature_kind(); - if !allowed_signature_for_role(&self.entity_role, &signature_kind) { - return Err(AssociationError::SignatureNotAllowed( - self.entity_role.clone(), - signature_kind, - )); - } - - let entity = Entity::new(self.entity_role.clone(), signer_address, false); - Ok(existing_state - .set_recovery_address(self.recovery_address.clone()) - .add(entity, 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.signature.text(), - self.recovery_address - ); - - sha256_string(inputs) - } -} - -pub struct AddAssociationEntry { - pub nonce: u32, - pub new_member_role: EntityRole, - pub existing_member_signature: Box, - pub new_member_signature: Box, -} - -impl AddAssociationEntry { - pub fn new_member_address(&self) -> String { - self.new_member_signature.recover_signer().unwrap() - } -} - -impl LogEntry for AddAssociationEntry { - fn update_state( - &self, - existing_state: AssociationState, - ) -> Result { - let association_hash = self.hash(); - if existing_state.has_seen(&association_hash) { - return Err(AssociationError::Replay); - } - - // Recovery address has to be set - if existing_state.recovery_address.is_none() { - return Err(AssociationError::NoRecoveryAddress); - } - - let new_member_address = self.new_member_signature.recover_signer()?; - let existing_member_address = self.existing_member_signature.recover_signer()?; - if new_member_address == existing_member_address { - return Err(AssociationError::Generic("tried to add self".to_string())); - } - - // Get the current version of the entity that added this new entry. If it has been revoked and added back, it will now be unrevoked - let existing_entity = existing_state - .get(&existing_member_address) - .ok_or(AssociationError::MissingExistingMember)?; - - if existing_entity.is_revoked { - // The entity that added this member is currently revoked. Check if this particular association is allowlisted - if !existing_state - .allowlisted_association_hashes - .contains(&association_hash) - { - return Err(AssociationError::AddedByRevokedMember); - } - } - - // 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(), - )); - } - - // Check to see if the new member was revoked - let is_new_member_revoked = existing_state.was_association_revoked(&association_hash); - let new_member = Entity::new( - self.new_member_role.clone(), - new_member_address, - is_new_member_revoked, - ); - - println!( - "Adding new entity to state {:?} with hash {}", - &new_member, &association_hash - ); - - Ok(existing_state.add(new_member, association_hash)?) - } - - fn hash(&self) -> String { - let inputs = format!( - "{}{:?}{}{}", - self.nonce, - self.new_member_role, - self.existing_member_signature.text(), - self.new_member_signature.text() - ); - sha256_string(inputs) - } -} - -pub struct RevokeAssociationEntry { - pub nonce: u32, - pub recovery_address_signature: Box, - pub revoked_association_hash: String, - pub allowed_child_hashes: Vec, -} - -impl LogEntry for RevokeAssociationEntry { - fn update_state( - &self, - existing_state: AssociationState, - ) -> Result { - // 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() - .ok_or(AssociationError::NoRecoveryAddress)?; - - // Ensure this message is signed by the recovery address - if recovery_signer != state_recovery_address { - return Err(AssociationError::MissingExistingMember); - } - - // Actually apply the revocation - Ok(existing_state.apply_revocation( - self.revoked_association_hash.clone(), - self.allowed_child_hashes.clone(), - )) - } - - fn hash(&self) -> String { - let inputs = format!( - "{}{}{}{}", - self.nonce, - self.recovery_address_signature.text(), - self.revoked_association_hash, - self.allowed_child_hashes.join(",") - ); - sha256_string(inputs) - } -} - -pub struct ChangeRecoveryAddressEntry { - pub nonce: u32, - pub recovery_address_signature: Box, - pub new_recovery_address: String, -} - -pub enum RecoveryLogEntry { - CreateXid(CreateXidEntry), - RevokeAssociation(RevokeAssociationEntry), -} - -impl LogEntry for RecoveryLogEntry { - fn update_state( - &self, - existing_state: AssociationState, - ) -> Result { - match self { - RecoveryLogEntry::CreateXid(create_xid) => create_xid.update_state(existing_state), - RecoveryLogEntry::RevokeAssociation(revoke_association) => { - revoke_association.update_state(existing_state) - } - } - } - - fn hash(&self) -> String { - match self { - RecoveryLogEntry::CreateXid(create_xid) => create_xid.hash(), - RecoveryLogEntry::RevokeAssociation(revoke_association) => revoke_association.hash(), - } - } -} pub fn apply_updates( initial_state: AssociationState, - associations: Vec, + association_events: Vec, ) -> AssociationState { - associations.iter().fold(initial_state, |state, update| { - match update.update_state(state.clone()) { - Ok(new_state) => new_state, - Err(err) => { - println!("invalid entry {}", err); - state - } - } - }) -} - -pub fn get_initial_state(recovery_log: Vec) -> AssociationState { - recovery_log + association_events .iter() - .fold(AssociationState::new(), |state, update| { - match update.update_state(state.clone()) { + .fold(initial_state, |state, update| { + match update.update_state(Some(state.clone())) { Ok(new_state) => new_state, Err(err) => { println!("invalid entry {}", err); @@ -305,54 +29,22 @@ pub fn get_initial_state(recovery_log: Vec) -> AssociationStat } pub fn get_state( - recovery_log: Vec, - association_updates: Vec, -) -> AssociationState { - let state = get_initial_state(recovery_log); - println!("Initial state {:?}", state); - apply_updates(state, association_updates) -} - -fn sha256_string(input: String) -> String { - let mut hasher = Sha256::new(); - hasher.update(input.as_bytes()); - let result = hasher.finalize(); - format!("{:x}", result) -} - -// 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::LegacyKey => false, - }, - EntityRole::LegacyKey => match signature_kind { - SignatureKind::Erc191 => false, - SignatureKind::Erc1271 => false, - SignatureKind::InstallationKey => false, - SignatureKind::LegacyKey => true, - }, - EntityRole::Installation => match signature_kind { - SignatureKind::Erc191 => false, - SignatureKind::Erc1271 => false, - SignatureKind::InstallationKey => true, - SignatureKind::LegacyKey => false, - }, - } + 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()) } -/** - * Revocation properties - * 1. Revoking an association will mark the entity added as revoked - * 2. Revoking an association will prevent new associations from being created with an `existing_entity_signature` from the revoked entity - * 3. Entities created with an `existing_entity_signature` of a revoked entity can be protected from revocation if they were added before the revocation - * 4. Revoked entities can be re-added with a new signature, so long as a new nonce is included in the signature - * 5. A revocation payload can be added to a subset of the association log. When this happens, all entities present in the subset will have the same revocation status that they have in the full log. - */ - #[cfg(test)] mod tests { use self::test_utils::{rand_string, rand_u32}; @@ -379,11 +71,28 @@ mod tests { } } - impl Default for AddAssociationEntry { + 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 { - nonce: rand_u32(), - new_member_role: EntityRole::Address, + client_timestamp_ns: rand_u32(), + new_member_role: EntityRole::Installation, existing_member_signature: MockSignature::new_boxed( true, rand_string(), @@ -392,313 +101,241 @@ mod tests { new_member_signature: MockSignature::new_boxed( true, rand_string(), - SignatureKind::Erc191, + SignatureKind::InstallationKey, ), }; } } - impl Default for CreateXidEntry { + impl Default for CreateXid { fn default() -> Self { let signer = rand_string(); return Self { nonce: rand_u32(), - signature: MockSignature::new_boxed(true, signer.clone(), SignatureKind::Erc191), - recovery_address: signer, - entity_role: EntityRole::Address, + account_address: signer.clone(), + initial_association: AddAssociation { + existing_member_signature: MockSignature::new_boxed( + true, + signer, + SignatureKind::Erc191, + ), + ..Default::default() + }, }; } } - impl Default for RevokeAssociationEntry { + impl Default for RevokeAssociation { fn default() -> Self { let signer = rand_string(); return Self { - nonce: rand_u32(), + client_timestamp_ns: rand_u32(), recovery_address_signature: MockSignature::new_boxed( true, signer, SignatureKind::Erc191, ), - revoked_association_hash: rand_string(), - allowed_child_hashes: vec![], + revoked_member: rand_string(), }; } } - 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() - } - } - - fn init_recovery_log() -> (Vec, String) { - let create_request = CreateXidEntry::default(); - let creator_address = create_request.signature.recover_signer().unwrap(); - let entries = vec![RecoveryLogEntry::CreateXid(create_request)]; - - (entries, creator_address) - } - #[test] fn test_create_and_add() { - let create_request = CreateXidEntry::default(); - let creator_address = create_request.signature.recover_signer().unwrap(); - let recovery_log = vec![RecoveryLogEntry::CreateXid(create_request)]; - let mut state = get_state(recovery_log, vec![]); - assert_eq!(state.entities().len(), 1); - - let add_installation_entry = AddAssociationEntry { - existing_member_signature: MockSignature::new_boxed( - true, - creator_address, - SignatureKind::Erc191, - ), - ..AddAssociationEntry::default() - }; - state = apply_updates(state, vec![add_installation_entry]); + 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 (recovery_log, creator_address) = init_recovery_log(); - let add_first_association = AddAssociationEntry { - existing_member_signature: MockSignature::new_boxed( + 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, - creator_address, + rand_string(), SignatureKind::Erc191, ), - ..AddAssociationEntry::default() - }; - let first_association_address = add_first_association - .new_member_signature - .recover_signer() - .unwrap(); - - let add_second_association = AddAssociationEntry { existing_member_signature: MockSignature::new_boxed( true, - first_association_address.clone(), - SignatureKind::Erc191, + new_installation.id.clone(), + SignatureKind::InstallationKey, ), - ..AddAssociationEntry::default() - }; - - let state = get_state( - recovery_log, - vec![add_first_association, add_second_association], - ); + }); - assert_eq!(state.entities().len(), 3); - assert_eq!( - state.get(&first_association_address).unwrap().is_revoked, - false - ); - assert_eq!( - state.get(&first_association_address).unwrap().id, - first_association_address - ); + let new_state = apply_updates(state, vec![update]); + assert_eq!(new_state.entities().len(), 3); } #[test] - fn add_from_revoked() { - let (mut recovery_log, creator_address) = init_recovery_log(); - let add_association = AddAssociationEntry { - existing_member_signature: MockSignature::new_boxed( - true, - creator_address.clone(), - SignatureKind::Erc191, - ), - ..AddAssociationEntry::default() - }; - - recovery_log.push(RecoveryLogEntry::RevokeAssociation( - RevokeAssociationEntry { - recovery_address_signature: MockSignature::new_boxed( + 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, - // Creator address is the recovery address, so this is valid - creator_address, + account_address.clone(), SignatureKind::Erc191, ), - revoked_association_hash: add_association.hash(), - // Not setting any allowed children here, since this doesn't have any - ..Default::default() + new_member_signature: MockSignature::new_boxed( + true, + rand_string(), + SignatureKind::LegacyKey, + ), + new_member_role: EntityRole::LegacyKey, }, - )); - - let add_another_association = AddAssociationEntry { - existing_member_signature: MockSignature::new_boxed( - true, - add_association.new_member_address(), - SignatureKind::Erc191, - ), - ..AddAssociationEntry::default() }; - let second_new_member_address = add_another_association.new_member_address(); - let state = get_state(recovery_log, vec![add_association, add_another_association]); - assert_eq!(state.entities().len(), 2); - assert!(state.get(&second_new_member_address).is_none()) + let initial_state = get_state(vec![AssociationEvent::CreateXid(create_request)]).unwrap(); + assert_eq!(initial_state.entities().len(), 2); + + let legacy_keys = initial_state.entities_by_role(EntityRole::LegacyKey); + assert_eq!(legacy_keys.len(), 1); } #[test] - fn add_from_re_added() { - let (mut recovery_log, creator_address) = init_recovery_log(); - - let add_association = AddAssociationEntry { - existing_member_signature: MockSignature::new_boxed( - true, - creator_address.clone(), - SignatureKind::Erc191, - ), - ..AddAssociationEntry::default() - }; - - let new_member_address = add_association.new_member_address(); - - recovery_log.push(RecoveryLogEntry::RevokeAssociation( - RevokeAssociationEntry { - recovery_address_signature: MockSignature::new_boxed( - true, - // Creator address is the recovery address, so this is valid - creator_address.clone(), + fn reject_invalid_signature() { + let create_request = CreateXid { + initial_association: AddAssociation { + existing_member_signature: MockSignature::new_boxed( + false, + rand_string(), SignatureKind::Erc191, ), - revoked_association_hash: add_association.hash(), - // Not setting any allowed children here, since this doesn't have any ..Default::default() }, - )); - - let add_same_member_back = AddAssociationEntry { - nonce: rand_u32(), - existing_member_signature: MockSignature::new_boxed( - true, - creator_address.clone(), - SignatureKind::Erc191, - ), - new_member_signature: MockSignature::new_boxed( - true, - new_member_address.clone(), - SignatureKind::Erc191, - ), ..Default::default() }; - let state = get_state(recovery_log, vec![add_association, add_same_member_back]); - assert_eq!(state.get(&new_member_address).unwrap().is_revoked, false) + 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 protect_children_from_revocation() { - let (mut recovery_log, creator_address) = init_recovery_log(); + fn reject_if_signer_not_existing_member() { + let create_request = CreateXid::default(); + let state = get_state(vec![AssociationEvent::CreateXid(create_request)]).unwrap(); - let add_association = AddAssociationEntry { - existing_member_signature: MockSignature::new_boxed( - true, - creator_address.clone(), - SignatureKind::Erc191, - ), - ..AddAssociationEntry::default() - }; + // The default here will create an AddAssociation from a random wallet + let update = AssociationEvent::AddAssociation(AddAssociation { + ..Default::default() + }); - let add_child = AddAssociationEntry { - existing_member_signature: MockSignature::new_boxed( + 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, - add_association.new_member_address(), + initial_wallet.clone(), SignatureKind::Erc191, ), - ..AddAssociationEntry::default() - }; + revoked_member: initial_wallet.clone(), + ..Default::default() + }); - let add_grandchild = AddAssociationEntry { - existing_member_signature: MockSignature::new_boxed( + 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, - add_child.new_member_address(), + recovery_address, SignatureKind::Erc191, ), - ..AddAssociationEntry::default() + revoked_member: installation_to_revoke.id.clone(), + ..Default::default() }; - recovery_log.push(RecoveryLogEntry::RevokeAssociation( - RevokeAssociationEntry { - recovery_address_signature: MockSignature::new_boxed( - true, - // Creator address is the recovery address, so this is valid - creator_address.clone(), - SignatureKind::Erc191, - ), - revoked_association_hash: add_association.hash(), - allowed_child_hashes: vec![add_child.hash()], - // Not setting any allowed children here, since this doesn't have any - ..Default::default() - }, - )); - - let first_member_address = add_association.new_member_address(); - let first_child_address = add_child.new_member_address(); - let grandchild_address = add_grandchild.new_member_address(); - - let state = get_state( - recovery_log, - vec![add_association, add_child, add_grandchild], - ); - assert_eq!(state.get(&first_member_address).unwrap().is_revoked, true); - assert_eq!(state.get(&first_child_address).unwrap().is_revoked, false); - assert_eq!(state.get(&grandchild_address).unwrap().is_revoked, false); + let new_state = update.update_state(Some(state)).unwrap(); + assert_eq!(new_state.entities().len(), 1); } #[test] - fn fail_if_ancestor_missing() { - let (recovery_log, creator_address) = init_recovery_log(); + 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 add_association = AddAssociationEntry { - existing_member_signature: MockSignature::new_boxed( + 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, - creator_address.clone(), + recovery_address.clone(), SignatureKind::Erc191, ), - ..AddAssociationEntry::default() + revoked_member: installation_to_revoke.id.clone(), + ..Default::default() }; - let add_child = AddAssociationEntry { + 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, - add_association.new_member_address(), + recovery_address, SignatureKind::Erc191, ), - ..AddAssociationEntry::default() - }; - - let add_grandchild = AddAssociationEntry { - existing_member_signature: MockSignature::new_boxed( + new_member_signature: MockSignature::new_boxed( true, - add_child.new_member_address(), - SignatureKind::Erc191, + original_installation_id, + SignatureKind::InstallationKey, ), - ..AddAssociationEntry::default() + ..Default::default() }; - let grandchild_address = add_grandchild.new_member_address(); - - let state = get_state( - recovery_log, - // Deliberately omitting the add_child, which is necessary here - vec![add_association, add_grandchild], - ); - - assert_eq!(state.get(&grandchild_address).is_none(), true); + 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) } } diff --git a/xmtp_id/src/associations/signature.rs b/xmtp_id/src/associations/signature.rs new file mode 100644 index 000000000..5de639156 --- /dev/null +++ b/xmtp_id/src/associations/signature.rs @@ -0,0 +1,21 @@ +use thiserror::Error; + +#[derive(Debug, Error, PartialEq)] +pub enum SignatureError { + #[error("Signature validation failed")] + Invalid, +} + +#[derive(Clone, Debug, PartialEq)] +pub enum SignatureKind { + Erc191, + Erc1271, + InstallationKey, + LegacyKey, +} + +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 index 7c5ad4914..efbc5c94c 100644 --- a/xmtp_id/src/associations/state.rs +++ b/xmtp_id/src/associations/state.rs @@ -2,9 +2,9 @@ use std::collections::{HashMap, HashSet}; use thiserror::Error; -use super::{entity::Entity, EntityRole}; +use super::{entity::Entity, hashes::generate_xid, EntityRole}; -#[derive(Debug, Error)] +#[derive(Debug, Error, PartialEq)] pub enum StateError { #[error("Not found")] NotFound, @@ -14,29 +14,23 @@ pub enum StateError { #[derive(Clone, Debug)] pub struct AssociationState { + pub xid: String, pub current_entities: HashMap, - // Stores the entity as it was at the time it was added - pub entities_by_event: HashMap, - pub revoked_association_hashes: HashSet, - pub allowlisted_association_hashes: HashSet, - pub recovery_address: Option, + pub recovery_address: String, + pub seen_events: HashSet, } impl AssociationState { - pub fn add(&self, entity: Entity, event_hash: String) -> Result { - self.replay_check(&event_hash)?; + pub fn add(&self, entity: Entity) -> Self { let mut new_state = self.clone(); - let _ = new_state - .entities_by_event - .insert(event_hash, entity.clone()); let _ = new_state.current_entities.insert(entity.id.clone(), entity); - Ok(new_state) + new_state } pub fn set_recovery_address(&self, recovery_address: String) -> Self { let mut new_state = self.clone(); - new_state.recovery_address = Some(recovery_address); + new_state.recovery_address = recovery_address; new_state } @@ -45,42 +39,36 @@ impl AssociationState { self.current_entities.get(id).map(|e| e.clone()) } - pub fn has_seen(&self, event_hash: &String) -> bool { - self.entities_by_event.contains_key(event_hash) - } + pub fn mark_event_seen(&self, event_hash: String) -> Self { + let mut new_state = self.clone(); + new_state.seen_events.insert(event_hash); - fn replay_check(&self, event_hash: &String) -> Result<(), StateError> { - if self.has_seen(event_hash) { - return Err(StateError::ReplayDetected); - } + new_state + } - Ok(()) + pub fn has_seen(&self, event_hash: &String) -> bool { + self.seen_events.contains(event_hash) } - pub fn apply_revocation( - &self, - revoked_association_hash: String, - allowlisted_association_hashes: Vec, - ) -> Self { + pub fn remove(&self, entity_id: String) -> Self { let mut new_state = self.clone(); - let _ = new_state - .revoked_association_hashes - .insert(revoked_association_hash); - new_state - .allowlisted_association_hashes - .extend(allowlisted_association_hashes); + let _ = new_state.current_entities.remove(&entity_id); new_state } - pub fn was_association_revoked(&self, association_hash: &String) -> bool { - self.revoked_association_hashes.contains(association_hash) - } - 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() @@ -89,23 +77,22 @@ impl AssociationState { .collect() } - pub fn new() -> Self { + 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: HashMap::new(), - entities_by_event: HashMap::new(), - revoked_association_hashes: HashSet::new(), - allowlisted_association_hashes: HashSet::new(), - recovery_address: None, + current_entities: { + let mut entities = HashMap::new(); + entities.insert(account_address.clone(), new_entity); + entities + }, + seen_events: HashSet::new(), + recovery_address: account_address, + xid, } } } -impl Default for AssociationState { - fn default() -> Self { - Self::new() - } -} - #[cfg(test)] mod tests { use crate::associations::test_utils::rand_string; @@ -114,11 +101,9 @@ mod tests { #[test] fn can_add_remove() { - let starting_state = AssociationState::new(); + let starting_state = AssociationState::new(rand_string(), 0); let new_entity = Entity::default(); - let with_add = starting_state - .add(new_entity.clone(), rand_string()) - .unwrap(); + 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()); } From 3365a71047afbb32597b9b04b2cb05d8c4fe3b10 Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Tue, 2 Apr 2024 16:49:10 -0700 Subject: [PATCH 4/5] Clean up comments --- xmtp_id/src/associations/association_log.rs | 4 +++- xmtp_id/src/associations/signature.rs | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/xmtp_id/src/associations/association_log.rs b/xmtp_id/src/associations/association_log.rs index 1852c222b..7abdb3f44 100644 --- a/xmtp_id/src/associations/association_log.rs +++ b/xmtp_id/src/associations/association_log.rs @@ -93,6 +93,8 @@ impl LogEntry for AddAssociation { ) -> 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); @@ -110,7 +112,7 @@ impl LogEntry for AddAssociation { } } - // Get the current version of the entity that added this new entry. If it has been revoked and added back, it will now be unrevoked + // Find the existing entity that authorized this add let existing_entity = existing_state .get(&existing_member_address) .ok_or(AssociationError::MissingExistingMember)?; diff --git a/xmtp_id/src/associations/signature.rs b/xmtp_id/src/associations/signature.rs index 5de639156..3d6619da4 100644 --- a/xmtp_id/src/associations/signature.rs +++ b/xmtp_id/src/associations/signature.rs @@ -8,6 +8,7 @@ pub enum SignatureError { #[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, From 092b1474a629266704f76aa53e1ed9594dc0b40e Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Tue, 2 Apr 2024 19:42:12 -0700 Subject: [PATCH 5/5] Update from latest protos --- xmtp_id/src/associations/association_log.rs | 51 +++++++++++++++------ xmtp_id/src/associations/entity.rs | 1 - xmtp_id/src/associations/mod.rs | 48 +++++++++++++++++-- xmtp_id/src/associations/signature.rs | 2 +- 4 files changed, 80 insertions(+), 22 deletions(-) diff --git a/xmtp_id/src/associations/association_log.rs b/xmtp_id/src/associations/association_log.rs index 7abdb3f44..098917af3 100644 --- a/xmtp_id/src/associations/association_log.rs +++ b/xmtp_id/src/associations/association_log.rs @@ -102,21 +102,21 @@ impl LogEntry for AddAssociation { 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())); } - if self.new_member_role == EntityRole::LegacyKey { + // 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); } } - // Find the existing entity that authorized this add - let existing_entity = existing_state - .get(&existing_member_address) - .ok_or(AssociationError::MissingExistingMember)?; - // Make sure that the signature type lines up with the role if !allowed_signature_for_role( &self.new_member_role, @@ -128,10 +128,30 @@ impl LogEntry for AddAssociation { )); } + 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), + Some(existing_entity_id), ); println!( @@ -166,6 +186,13 @@ impl LogEntry for RevokeAssociation { 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 @@ -239,19 +266,13 @@ pub fn allowed_signature_for_role(role: &EntityRole, signature_kind: &SignatureK SignatureKind::Erc191 => true, SignatureKind::Erc1271 => true, SignatureKind::InstallationKey => false, - SignatureKind::LegacyKey => false, - }, - EntityRole::LegacyKey => match signature_kind { - SignatureKind::Erc191 => false, - SignatureKind::Erc1271 => false, - SignatureKind::InstallationKey => false, - SignatureKind::LegacyKey => true, + SignatureKind::LegacyDelegated => true, }, EntityRole::Installation => match signature_kind { SignatureKind::Erc191 => false, SignatureKind::Erc1271 => false, SignatureKind::InstallationKey => true, - SignatureKind::LegacyKey => false, + SignatureKind::LegacyDelegated => false, }, } } diff --git a/xmtp_id/src/associations/entity.rs b/xmtp_id/src/associations/entity.rs index 8e687f253..9862caf01 100644 --- a/xmtp_id/src/associations/entity.rs +++ b/xmtp_id/src/associations/entity.rs @@ -2,7 +2,6 @@ pub enum EntityRole { Installation, Address, - LegacyKey, } #[derive(Clone, Debug)] diff --git a/xmtp_id/src/associations/mod.rs b/xmtp_id/src/associations/mod.rs index be0a3e8fd..616aa3ae2 100644 --- a/xmtp_id/src/associations/mod.rs +++ b/xmtp_id/src/associations/mod.rs @@ -190,22 +190,22 @@ mod tests { existing_member_signature: MockSignature::new_boxed( true, account_address.clone(), - SignatureKind::Erc191, + SignatureKind::LegacyDelegated, ), new_member_signature: MockSignature::new_boxed( true, rand_string(), - SignatureKind::LegacyKey, + SignatureKind::InstallationKey, ), - new_member_role: EntityRole::LegacyKey, + new_member_role: EntityRole::Installation, }, }; let initial_state = get_state(vec![AssociationEvent::CreateXid(create_request)]).unwrap(); assert_eq!(initial_state.entities().len(), 2); - let legacy_keys = initial_state.entities_by_role(EntityRole::LegacyKey); - assert_eq!(legacy_keys.len(), 1); + let installation_keys = initial_state.entities_by_role(EntityRole::Installation); + assert_eq!(installation_keys.len(), 1); } #[test] @@ -338,4 +338,42 @@ mod tests { 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 index 3d6619da4..3600ccfc6 100644 --- a/xmtp_id/src/associations/signature.rs +++ b/xmtp_id/src/associations/signature.rs @@ -12,7 +12,7 @@ pub enum SignatureKind { Erc191, Erc1271, InstallationKey, - LegacyKey, + LegacyDelegated, } pub trait Signature {