From 9be4a86746a06464eaef975ec0bfcef85fe5593e Mon Sep 17 00:00:00 2001 From: Nicholas Molnar <65710+neekolas@users.noreply.github.com> Date: Fri, 5 Apr 2024 18:29:43 -0700 Subject: [PATCH] Add Identity Update Builder --- xmtp_id/src/associations/builder.rs | 254 +++++++++++++++++++++++++ xmtp_id/src/associations/mod.rs | 53 +----- xmtp_id/src/associations/signature.rs | 21 +- xmtp_id/src/associations/signer.rs | 35 ++++ xmtp_id/src/associations/test_utils.rs | 83 ++++++++ 5 files changed, 397 insertions(+), 49 deletions(-) create mode 100644 xmtp_id/src/associations/builder.rs create mode 100644 xmtp_id/src/associations/signer.rs diff --git a/xmtp_id/src/associations/builder.rs b/xmtp_id/src/associations/builder.rs new file mode 100644 index 000000000..9252926b5 --- /dev/null +++ b/xmtp_id/src/associations/builder.rs @@ -0,0 +1,254 @@ +use std::collections::HashMap; + +use thiserror::Error; +use xmtp_mls::utils::time::now_ns; + +use super::{ + association_log::{AddAssociation, ChangeRecoveryAddress, CreateInbox, RevokeAssociation}, + signer::{Signer, SignerError}, + unsigned_actions::{ + SignatureTextCreator, UnsignedAction, UnsignedAddAssociation, UnsignedCreateInbox, + UnsignedIdentityUpdate, + }, + Action, AssociationState, IdentityUpdate, MemberIdentifier, Signature, +}; + +#[derive(Error, Debug)] +pub enum IdentityBuilderError { + #[error("Missing signer")] + MissingSigner, + #[error("Signer error {0}")] + Signer(#[from] SignerError), +} + +#[derive(Clone, PartialEq)] +enum SignatureField { + InitialAddress, + ExistingMember, + NewMember, + RecoveryAddress, +} + +struct PendingSignature { + signer_identity: MemberIdentifier, + field_name: SignatureField, +} + +struct PendingIdentityAction { + unsigned_action: UnsignedAction, + pending_signatures: Vec, +} + +pub struct IdentityUpdateBuilder { + inbox_id: String, + client_timestamp_ns: u64, + actions: Vec, + signers: HashMap>, +} + +impl IdentityUpdateBuilder { + pub fn new(inbox_id: String) -> Self { + Self { + inbox_id, + client_timestamp_ns: now_ns() as u64, + actions: vec![], + signers: HashMap::new(), + } + } + + pub fn create_inbox(mut self, signer: Box, nonce: u64) -> Self { + let signer_identity = signer.signer_identity(); + let pending_action = PendingIdentityAction { + unsigned_action: UnsignedAction::CreateInbox(UnsignedCreateInbox { + account_address: signer_identity.to_string(), + nonce, + }), + pending_signatures: vec![PendingSignature { + signer_identity: signer_identity.clone(), + field_name: SignatureField::InitialAddress, + }], + }; + self.actions.push(pending_action); + + self.signers.insert(signer_identity, signer); + + self + } + + pub fn add_association( + mut self, + new_member_signer: Box, + existing_member_signer: Box, + ) -> Self { + let new_member_identifier = new_member_signer.signer_identity(); + let existing_member_identifier = existing_member_signer.signer_identity(); + self.actions.push(PendingIdentityAction { + unsigned_action: UnsignedAction::AddAssociation(UnsignedAddAssociation { + new_member_identifier: new_member_identifier.clone(), + inbox_id: self.inbox_id.clone(), + }), + pending_signatures: vec![ + PendingSignature { + signer_identity: existing_member_identifier.clone(), + field_name: SignatureField::ExistingMember, + }, + PendingSignature { + signer_identity: new_member_identifier.clone(), + field_name: SignatureField::NewMember, + }, + ], + }); + self.signers + .insert(new_member_identifier, new_member_signer); + self.signers + .insert(existing_member_identifier, existing_member_signer); + self + } + + pub fn build(self) -> Result { + let unsigned_actions: Vec = self + .actions + .iter() + .map(|pending_action| pending_action.unsigned_action.clone()) + .collect(); + + let signature_text = get_signature_text(unsigned_actions, self.client_timestamp_ns); + + // Go through all the unique signers for this update and sign the signature text for each of them + let signatures: HashMap> = self + .signers + .iter() + .try_fold(HashMap::new(), |mut acc, (signer_identity, signer)| -> Result>, IdentityBuilderError> { + acc.insert( + signer_identity.clone(), + signer.sign(signature_text.as_str())?, + ); + + Ok(acc) + })?; + + let signed_actions = + self.actions + .into_iter() + .map(|pending_action| -> Result { + match pending_action.unsigned_action { + UnsignedAction::CreateInbox(unsigned_action) => { + let signer_identity = find_signer_identity( + &pending_action.pending_signatures, + SignatureField::InitialAddress, + )?; + let initial_address_signature = signatures + .get(&signer_identity) + .cloned() + .ok_or(IdentityBuilderError::MissingSigner)?; + + Ok(Action::CreateInbox(CreateInbox { + nonce: unsigned_action.nonce, + account_address: unsigned_action.account_address, + initial_address_signature, + })) + } + UnsignedAction::AddAssociation(unsigned_action) => { + let existing_member_signer_identity = find_signer_identity( + &pending_action.pending_signatures, + SignatureField::ExistingMember, + )?; + let new_member_signer_identity = find_signer_identity( + &pending_action.pending_signatures, + SignatureField::NewMember, + )?; + + let existing_member_signature = signatures + .get(&existing_member_signer_identity) + .cloned() + .ok_or(IdentityBuilderError::MissingSigner)?; + let new_member_signature = signatures + .get(&new_member_signer_identity) + .cloned() + .ok_or(IdentityBuilderError::MissingSigner)?; + + Ok(Action::AddAssociation(AddAssociation { + new_member_identifier: unsigned_action.new_member_identifier, + existing_member_signature, + new_member_signature, + })) + } + _ => todo!(), + } + }) + .collect::, IdentityBuilderError>>()?; + + Ok(IdentityUpdate::new( + signed_actions, + self.client_timestamp_ns, + )) + } +} + +fn find_signer_identity( + pending_signatures: &Vec, + field: SignatureField, +) -> Result { + Ok(pending_signatures + .iter() + .find(|pending_signature| pending_signature.field_name == field) + .ok_or(IdentityBuilderError::MissingSigner)? + .signer_identity + .clone()) +} + +fn get_signature_text(actions: Vec, client_timestamp_ns: u64) -> String { + let identity_update = UnsignedIdentityUpdate { + client_timestamp_ns, + actions, + }; + + identity_update.signature_text() +} + +#[cfg(test)] +mod tests { + use crate::associations::{ + get_state, + hashes::generate_inbox_id, + test_utils::{rand_vec, MockSigner}, + SignatureKind, + }; + + use super::*; + + #[test] + fn create_inbox() { + let account_address = "account_address".to_string(); + let nonce = 0; + let inbox_id = generate_inbox_id(&account_address, &nonce); + let identity_update = IdentityUpdateBuilder::new(inbox_id) + .create_inbox( + MockSigner::new_boxed(account_address.into(), SignatureKind::Erc191), + nonce, + ) + .build() + .unwrap(); + + get_state(vec![identity_update]).expect("should be valid"); + } + + #[test] + fn create_and_add_identity() { + let account_address = "account_address".to_string(); + let nonce = 0; + let inbox_id = generate_inbox_id(&account_address, &nonce); + let existing_member_signer = + MockSigner::new_boxed(account_address.into(), SignatureKind::Erc191); + let new_member_signer = + MockSigner::new_boxed(rand_vec().into(), SignatureKind::InstallationKey); + + let identity_update = IdentityUpdateBuilder::new(inbox_id) + .create_inbox(existing_member_signer.clone(), nonce) + .add_association(new_member_signer, existing_member_signer) + .build() + .unwrap(); + + get_state(vec![identity_update]).expect("should be valid"); + } +} diff --git a/xmtp_id/src/associations/mod.rs b/xmtp_id/src/associations/mod.rs index d868f6d35..b1b0af985 100644 --- a/xmtp_id/src/associations/mod.rs +++ b/xmtp_id/src/associations/mod.rs @@ -1,7 +1,9 @@ mod association_log; +pub mod builder; mod hashes; mod member; mod signature; +mod signer; mod state; #[cfg(test)] mod test_utils; @@ -35,61 +37,16 @@ pub fn get_state(updates: Vec) -> Result) -> Self { Self::new(actions, rand_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(); @@ -256,7 +213,7 @@ mod tests { true, member_identifier.clone(), SignatureKind::LegacyDelegated, - Some(0), + Some("0".to_string()), ), }; let state = get_state(vec![IdentityUpdate::new_test(vec![Action::CreateInbox( @@ -272,7 +229,7 @@ mod tests { member_identifier, SignatureKind::LegacyDelegated, // All requests from the same legacy key will have the same signature nonce - Some(0), + Some("0".to_string()), ), ..Default::default() }); diff --git a/xmtp_id/src/associations/signature.rs b/xmtp_id/src/associations/signature.rs index b435dee5c..d71bec608 100644 --- a/xmtp_id/src/associations/signature.rs +++ b/xmtp_id/src/associations/signature.rs @@ -28,8 +28,27 @@ impl std::fmt::Display for SignatureKind { } } -pub trait Signature { +pub trait Signature: SignatureClone { fn recover_signer(&self) -> Result; fn signature_kind(&self) -> SignatureKind; fn bytes(&self) -> Vec; } + +pub trait SignatureClone { + fn clone_box(&self) -> Box; +} + +impl SignatureClone for T +where + T: 'static + Signature + Clone, +{ + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +impl Clone for Box { + fn clone(&self) -> Box { + self.clone_box() + } +} diff --git a/xmtp_id/src/associations/signer.rs b/xmtp_id/src/associations/signer.rs new file mode 100644 index 000000000..b606bc5fc --- /dev/null +++ b/xmtp_id/src/associations/signer.rs @@ -0,0 +1,35 @@ +use thiserror::Error; + +use super::{MemberIdentifier, Signature, SignatureKind}; + +#[derive(Error, Debug)] +pub enum SignerError { + #[error("Signature error {0}")] + Generic(String), +} + +#[async_trait::async_trait] +pub trait Signer: SignerClone { + fn signer_identity(&self) -> MemberIdentifier; + fn signature_kind(&self) -> SignatureKind; + fn sign(&self, text: &str) -> Result, SignerError>; +} + +pub trait SignerClone { + fn clone_box(&self) -> Box; +} + +impl SignerClone for T +where + T: 'static + Signer + Clone, +{ + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +impl Clone for Box { + fn clone(&self) -> Box { + self.clone_box() + } +} diff --git a/xmtp_id/src/associations/test_utils.rs b/xmtp_id/src/associations/test_utils.rs index ce8a08a4c..4999566bc 100644 --- a/xmtp_id/src/associations/test_utils.rs +++ b/xmtp_id/src/associations/test_utils.rs @@ -1,5 +1,10 @@ use rand::{distributions::Alphanumeric, Rng}; +use super::{ + signer::{Signer, SignerClone, SignerError}, + MemberIdentifier, Signature, SignatureError, SignatureKind, +}; + pub fn rand_string() -> String { let v: String = rand::thread_rng() .sample_iter(&Alphanumeric) @@ -19,3 +24,81 @@ pub fn rand_vec() -> Vec { rand::thread_rng().fill(&mut buf[..]); buf.to_vec() } + +#[derive(Clone)] +pub struct MockSignature { + is_valid: bool, + signer_identity: MemberIdentifier, + signature_kind: SignatureKind, + signature_nonce: String, +} + +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_string()); + 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() + } +} + +#[derive(Clone)] +pub struct MockSigner { + identity: MemberIdentifier, + signature_kind: SignatureKind, +} + +impl Signer for MockSigner { + fn signature_kind(&self) -> SignatureKind { + self.signature_kind.clone() + } + fn signer_identity(&self) -> MemberIdentifier { + self.identity.clone() + } + + fn sign(&self, text: &str) -> Result, SignerError> { + Ok(MockSignature::new_boxed( + true, + self.signer_identity(), + self.signature_kind(), + Some(text.to_string()), + )) + } +} + +impl MockSigner { + pub fn new_boxed(identity: MemberIdentifier, signature_kind: SignatureKind) -> Box { + Box::new(Self { + identity, + signature_kind, + }) + } +}