From 6cc75a5a70824c6500181b406faac62c531ce00b Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Thu, 4 Apr 2024 15:01:19 -0700 Subject: [PATCH] added association log and signature --- xmtp_id/src/associations/association_log.rs | 313 ++++++++++++++++++ xmtp_id/src/associations/mod.rs | 332 ++++++++++++++++++++ xmtp_id/src/associations/signature.rs | 24 ++ 3 files changed, 669 insertions(+) create mode 100644 xmtp_id/src/associations/association_log.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..b1f47d1c1 --- /dev/null +++ b/xmtp_id/src/associations/association_log.rs @@ -0,0 +1,313 @@ +use super::hashes::generate_xid; +use super::member::{Member, MemberIdentifier, MemberKind}; +use super::signature::{Signature, SignatureError, SignatureKind}; +use super::state::AssociationState; + +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("Missing existing member")] + MissingExistingMember, + #[error("Legacy key is only allowed to be associated using a legacy signature with nonce 0")] + LegacySignatureReuse, + #[error("The new member identifier does not match the signer")] + NewMemberIdSignatureMismatch, + #[error("Signature not allowed for role {0:?} {1:?}")] + SignatureNotAllowed(MemberKind, SignatureKind), + #[error("Replay detected")] + Replay, +} + +pub trait IdentityAction { + fn update_state( + &self, + existing_state: Option, + ) -> Result; + fn signatures(&self) -> Vec>; + fn replay_check(&self, state: &AssociationState) -> Result<(), AssociationError> { + let signatures = self.signatures(); + for signature in signatures { + if state.has_seen(&signature) { + return Err(AssociationError::Replay); + } + } + + Ok(()) + } +} + +pub struct CreateInbox { + pub nonce: u64, + pub account_address: String, + pub initial_address_signature: Box, +} + +impl IdentityAction for CreateInbox { + 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 recovered_signer = self.initial_address_signature.recover_signer()?; + if recovered_signer.ne(&MemberIdentifier::Address(account_address.clone())) { + return Err(AssociationError::MissingExistingMember); + } + + if self.initial_address_signature.signature_kind() == SignatureKind::LegacyDelegated + && self.nonce != 0 + { + return Err(AssociationError::LegacySignatureReuse); + } + + Ok(AssociationState::new(account_address, self.nonce)) + } + + fn signatures(&self) -> Vec> { + vec![self.initial_address_signature.bytes()] + } +} + +pub struct AddAssociation { + pub client_timestamp_ns: u64, + pub new_member_signature: Box, + pub new_member_identifier: MemberIdentifier, + pub existing_member_signature: Box, +} + +impl IdentityAction for AddAssociation { + fn update_state( + &self, + maybe_existing_state: Option, + ) -> Result { + let existing_state = maybe_existing_state.ok_or(AssociationError::NotCreated)?; + self.replay_check(&existing_state)?; + + let new_member_address = self.new_member_signature.recover_signer()?; + if new_member_address.ne(&self.new_member_identifier) { + return Err(AssociationError::NewMemberIdSignatureMismatch); + } + + let existing_member_identifier = self.existing_member_signature.recover_signer()?; + let recovery_address = existing_state.recovery_address(); + + if new_member_address.ne(&self.new_member_identifier) { + return Err(AssociationError::Generic( + "new member identifier does not match signature".to_string(), + )); + } + + if new_member_address.ne(&self.new_member_identifier) { + return Err(AssociationError::Generic( + "new member identifier does not match signature".to_string(), + )); + } + + // You cannot add yourself + if new_member_address == existing_member_identifier { + 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() + .ne(&generate_xid(&existing_member_identifier.to_string(), &0)) + { + return Err(AssociationError::LegacySignatureReuse); + } + } + + // Make sure that the signature type lines up with the role + if !allowed_signature_for_kind( + &self.new_member_identifier.kind(), + &self.new_member_signature.signature_kind(), + ) { + return Err(AssociationError::SignatureNotAllowed( + self.new_member_identifier.kind(), + self.new_member_signature.signature_kind(), + )); + } + + let existing_member = existing_state.get(&existing_member_identifier); + + let existing_entity_id = match existing_member { + // If there is an existing member of the XID, use that member's ID + Some(member) => member.identifier, + None => { + let recovery_identifier = MemberIdentifier::Address(recovery_address.clone()); + // Check if it is a signature from the recovery address, which is allowed to add members + if existing_member_identifier.ne(&recovery_identifier) { + 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_identifier + } + }; + + let new_member = Member::new(new_member_address, Some(existing_entity_id)); + + println!("Adding new entity to state {:?}", &new_member); + + Ok(existing_state.add(new_member)) + } + + fn signatures(&self) -> Vec> { + vec![ + self.existing_member_signature.bytes(), + self.new_member_signature.bytes(), + ] + } +} + +pub struct RevokeAssociation { + pub client_timestamp_ns: u64, + pub recovery_address_signature: Box, + pub revoked_member: MemberIdentifier, +} + +impl IdentityAction for RevokeAssociation { + fn update_state( + &self, + maybe_existing_state: Option, + ) -> Result { + let existing_state = maybe_existing_state.ok_or(AssociationError::NotCreated)?; + self.replay_check(&existing_state)?; + + if self.recovery_address_signature.signature_kind() == SignatureKind::LegacyDelegated { + return Err(AssociationError::SignatureNotAllowed( + MemberKind::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(); + + // Ensure this message is signed by the recovery address + if recovery_signer.ne(&MemberIdentifier::Address(state_recovery_address.clone())) { + return Err(AssociationError::MissingExistingMember); + } + + let installations_to_remove: Vec = existing_state + .members_by_parent(&self.revoked_member) + .into_iter() + // Only remove children if they are installations + .filter(|child| child.kind() == MemberKind::Installation) + .collect(); + + // Actually apply the revocation to the parent + let new_state = existing_state.remove(&self.revoked_member); + + Ok(installations_to_remove + .iter() + .fold(new_state, |state, installation| { + state.remove(&installation.identifier) + })) + } + + fn signatures(&self) -> Vec> { + vec![self.recovery_address_signature.bytes()] + } +} + +pub enum Action { + CreateInbox(CreateInbox), + AddAssociation(AddAssociation), + RevokeAssociation(RevokeAssociation), +} + +impl IdentityAction for Action { + fn update_state( + &self, + existing_state: Option, + ) -> Result { + match self { + Action::CreateInbox(event) => event.update_state(existing_state), + Action::AddAssociation(event) => event.update_state(existing_state), + Action::RevokeAssociation(event) => event.update_state(existing_state), + } + } + + fn signatures(&self) -> Vec> { + match self { + Action::CreateInbox(event) => event.signatures(), + Action::AddAssociation(event) => event.signatures(), + Action::RevokeAssociation(event) => event.signatures(), + } + } +} + +pub struct IdentityUpdate { + pub actions: Vec, +} + +impl IdentityUpdate { + pub fn new(actions: Vec) -> Self { + Self { actions } + } +} + +impl IdentityAction for IdentityUpdate { + fn update_state( + &self, + existing_state: Option, + ) -> Result { + let mut state = existing_state.clone(); + for action in &self.actions { + state = Some(action.update_state(state)?); + } + + let new_state = state.ok_or(AssociationError::NotCreated)?; + + // After all the updates in the LogEntry have been processed, add the list of signatures to the state + // so that the signatures can not be re-used in subsequent updates + Ok(new_state.add_seen_signatures(self.signatures())) + } + + fn signatures(&self) -> Vec> { + self.actions + .iter() + .flat_map(|action| action.signatures()) + .collect() + } +} + +// Ensure that the type of signature matches the new entity's role. +pub fn allowed_signature_for_kind(role: &MemberKind, signature_kind: &SignatureKind) -> bool { + match role { + MemberKind::Address => match signature_kind { + SignatureKind::Erc191 => true, + SignatureKind::Erc1271 => true, + SignatureKind::InstallationKey => false, + SignatureKind::LegacyDelegated => true, + }, + MemberKind::Installation => match signature_kind { + SignatureKind::Erc191 => false, + SignatureKind::Erc1271 => false, + SignatureKind::InstallationKey => true, + SignatureKind::LegacyDelegated => false, + }, + } +} diff --git a/xmtp_id/src/associations/mod.rs b/xmtp_id/src/associations/mod.rs index 839409895..166567b46 100644 --- a/xmtp_id/src/associations/mod.rs +++ b/xmtp_id/src/associations/mod.rs @@ -1,8 +1,340 @@ +mod association_log; mod hashes; mod member; +mod signature; mod state; #[cfg(test)] mod test_utils; +pub use self::association_log::*; pub use self::member::{Member, MemberIdentifier, MemberKind}; +pub use self::signature::{Signature, SignatureError, SignatureKind}; pub use self::state::AssociationState; + +pub fn apply_update( + initial_state: AssociationState, + update: IdentityUpdate, +) -> Result { + update.update_state(Some(initial_state)) +} + +pub fn get_state(updates: Vec) -> Result { + let new_state = updates.iter().try_fold( + None, + |state, update| -> Result, AssociationError> { + let updated_state = update.update_state(state)?; + Ok(Some(updated_state)) + }, + )?; + + new_state.ok_or(AssociationError::NotCreated) +} + +#[cfg(test)] +mod tests { + use self::test_utils::{rand_string, rand_u64, rand_vec}; + + use super::*; + + #[derive(Clone)] + struct MockSignature { + is_valid: bool, + signer_identity: MemberIdentifier, + signature_kind: SignatureKind, + signature_nonce: u64, + } + + impl MockSignature { + pub fn new_boxed( + is_valid: bool, + signer_identity: MemberIdentifier, + signature_kind: SignatureKind, + // Signature nonce is used to control what the signature bytes are + // Defaults to random + signature_nonce: Option, + ) -> Box { + let nonce = signature_nonce.unwrap_or(rand_u64()); + Box::new(Self { + is_valid, + signer_identity, + signature_kind, + signature_nonce: nonce, + }) + } + } + + 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 bytes(&self) -> Vec { + let sig = format!("{}{}", self.signer_identity, self.signature_nonce); + sig.as_bytes().to_vec() + } + } + + impl Default for AddAssociation { + fn default() -> Self { + let existing_member = rand_string(); + let new_member = rand_vec(); + return Self { + client_timestamp_ns: rand_u64(), + existing_member_signature: MockSignature::new_boxed( + true, + existing_member.into(), + SignatureKind::Erc191, + None, + ), + new_member_signature: MockSignature::new_boxed( + true, + new_member.clone().into(), + SignatureKind::InstallationKey, + None, + ), + new_member_identifier: new_member.into(), + }; + } + } + + // Default will create an inbox with a ERC-191 signature + impl Default for CreateInbox { + fn default() -> Self { + let signer = rand_string(); + return Self { + nonce: rand_u64(), + account_address: signer.clone(), + initial_address_signature: MockSignature::new_boxed( + true, + signer.into(), + SignatureKind::Erc191, + None, + ), + }; + } + } + + impl Default for RevokeAssociation { + fn default() -> Self { + let signer = rand_string(); + return Self { + client_timestamp_ns: rand_u64(), + recovery_address_signature: MockSignature::new_boxed( + true, + signer.into(), + SignatureKind::Erc191, + None, + ), + revoked_member: rand_string().into(), + }; + } + } + + fn new_test_inbox() -> AssociationState { + let create_request = CreateInbox::default(); + let identity_update = IdentityUpdate::new(vec![Action::CreateInbox(create_request)]); + + get_state(vec![identity_update]).unwrap() + } + + #[test] + fn test_create_inbox() { + let create_request = CreateInbox::default(); + let account_address = create_request.account_address.clone(); + let identity_update = IdentityUpdate::new(vec![Action::CreateInbox(create_request)]); + let state = get_state(vec![identity_update]).unwrap(); + assert_eq!(state.members().len(), 1); + + let existing_entity = state.get(&account_address.clone().into()).unwrap(); + assert!(existing_entity.identifier.eq(&account_address.into())); + } + + #[test] + fn create_and_add_separately() { + let initial_state = new_test_inbox(); + let new_installation_identifier: MemberIdentifier = rand_vec().into(); + let first_member: MemberIdentifier = initial_state.recovery_address().clone().into(); + + let update = Action::AddAssociation(AddAssociation { + new_member_identifier: new_installation_identifier.clone(), + new_member_signature: MockSignature::new_boxed( + true, + new_installation_identifier.clone(), + SignatureKind::InstallationKey, + None, + ), + existing_member_signature: MockSignature::new_boxed( + true, + first_member.clone(), + SignatureKind::InstallationKey, + None, + ), + ..Default::default() + }); + + let new_state = apply_update(initial_state, IdentityUpdate::new(vec![update])).unwrap(); + assert_eq!(new_state.members().len(), 2); + + let new_member = new_state.get(&new_installation_identifier).unwrap(); + assert_eq!(new_member.added_by_entity, Some(first_member)); + } + + #[test] + fn create_and_add_together() { + let create_action = CreateInbox::default(); + let account_address = create_action.account_address.clone(); + let new_member_identifier: MemberIdentifier = rand_vec().into(); + let add_action = AddAssociation { + existing_member_signature: MockSignature::new_boxed( + true, + account_address.clone().into(), + SignatureKind::Erc191, + None, + ), + // Add an installation ID + new_member_signature: MockSignature::new_boxed( + true, + new_member_identifier.clone(), + SignatureKind::InstallationKey, + None, + ), + new_member_identifier: new_member_identifier.clone(), + ..Default::default() + }; + let identity_update = IdentityUpdate::new(vec![ + Action::CreateInbox(create_action), + Action::AddAssociation(add_action), + ]); + let state = get_state(vec![identity_update]).unwrap(); + assert_eq!(state.members().len(), 2); + assert_eq!( + state.get(&new_member_identifier).unwrap().added_by_entity, + Some(account_address.into()) + ); + } + + #[test] + fn create_from_legacy_key() { + let member_identifier: MemberIdentifier = rand_string().into(); + let create_action = CreateInbox { + nonce: 0, + account_address: member_identifier.to_string(), + initial_address_signature: MockSignature::new_boxed( + true, + member_identifier.clone(), + SignatureKind::LegacyDelegated, + Some(0), + ), + }; + let state = get_state(vec![IdentityUpdate::new(vec![Action::CreateInbox( + create_action, + )])]) + .unwrap(); + assert_eq!(state.members().len(), 1); + + // The legacy key can only be used once. After this, subsequent updates should fail + let update = Action::AddAssociation(AddAssociation { + existing_member_signature: MockSignature::new_boxed( + true, + member_identifier, + SignatureKind::LegacyDelegated, + // All requests from the same legacy key will have the same signature nonce + Some(0), + ), + ..Default::default() + }); + let update_result = apply_update(state, IdentityUpdate::new(vec![update])); + assert_eq!(update_result.is_err(), true); + assert_eq!(update_result.err().unwrap(), AssociationError::Replay); + } + + #[test] + fn reject_invalid_signature_on_create() { + let bad_signature = + MockSignature::new_boxed(false, rand_string().into(), SignatureKind::Erc191, None); + let action = CreateInbox { + initial_address_signature: bad_signature.clone(), + ..Default::default() + }; + + let state_result = get_state(vec![IdentityUpdate::new(vec![Action::CreateInbox(action)])]); + assert_eq!(state_result.is_err(), true); + assert_eq!( + state_result.err().unwrap(), + AssociationError::Signature(SignatureError::Invalid) + ); + } + + #[test] + fn reject_invalid_signature_on_update() { + let initial_state = new_test_inbox(); + let bad_signature = + MockSignature::new_boxed(false, rand_string().into(), SignatureKind::Erc191, None); + + let update_with_bad_existing_member = Action::AddAssociation(AddAssociation { + existing_member_signature: bad_signature.clone(), + ..Default::default() + }); + + let update_result = apply_update( + initial_state.clone(), + IdentityUpdate::new(vec![update_with_bad_existing_member]), + ); + assert_eq!(update_result.is_err(), true); + assert_eq!( + update_result.err().unwrap(), + AssociationError::Signature(SignatureError::Invalid) + ); + + let update_with_bad_new_member = Action::AddAssociation(AddAssociation { + new_member_signature: bad_signature.clone(), + existing_member_signature: MockSignature::new_boxed( + true, + initial_state.recovery_address().clone().into(), + SignatureKind::Erc191, + None, + ), + ..Default::default() + }); + + let update_result_2 = apply_update( + initial_state, + IdentityUpdate::new(vec![update_with_bad_new_member]), + ); + assert_eq!(update_result_2.is_err(), true); + assert_eq!( + update_result_2.err().unwrap(), + AssociationError::Signature(SignatureError::Invalid) + ); + } + + #[test] + fn reject_if_signer_not_existing_member() { + let create_request = Action::CreateInbox(CreateInbox::default()); + // The default here will create an AddAssociation from a random wallet + let update = Action::AddAssociation(AddAssociation { + // Existing member signature is coming from a random wallet + existing_member_signature: MockSignature::new_boxed( + true, + rand_string().into(), + SignatureKind::Erc191, + None, + ), + ..Default::default() + }); + + let state_result = get_state(vec![IdentityUpdate::new(vec![create_request, update])]); + assert_eq!(state_result.is_err(), true); + assert_eq!( + state_result.err().unwrap(), + AssociationError::MissingExistingMember + ); + } +} diff --git a/xmtp_id/src/associations/signature.rs b/xmtp_id/src/associations/signature.rs new file mode 100644 index 000000000..34bbe14cb --- /dev/null +++ b/xmtp_id/src/associations/signature.rs @@ -0,0 +1,24 @@ +use thiserror::Error; + +use super::MemberIdentifier; + +#[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 bytes(&self) -> Vec; +}