Add Identity Update Builder

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<PendingSignature>,
+}
+
+pub struct IdentityUpdateBuilder {
+    inbox_id: String,
+    client_timestamp_ns: u64,
+    actions: Vec<PendingIdentityAction>,
+    signers: HashMap<MemberIdentifier, Box<dyn Signer>>,
+}
+
+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<dyn Signer>, 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<dyn Signer>,
+        existing_member_signer: Box<dyn Signer>,
+    ) -> 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<IdentityUpdate, IdentityBuilderError> {
+        let unsigned_actions: Vec<UnsignedAction> = 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<MemberIdentifier, Box<dyn Signature>> = self
+            .signers
+            .iter()
+            .try_fold(HashMap::new(), |mut acc, (signer_identity, signer)| -> Result<HashMap<MemberIdentifier, Box<dyn Signature>>, 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<Action, IdentityBuilderError> {
+                    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::<Result<Vec<Action>, IdentityBuilderError>>()?;
+
+        Ok(IdentityUpdate::new(
+            signed_actions,
+            self.client_timestamp_ns,
+        ))
+    }
+}
+
+fn find_signer_identity(
+    pending_signatures: &Vec<PendingSignature>,
+    field: SignatureField,
+) -> Result<MemberIdentifier, IdentityBuilderError> {
+    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<UnsignedAction>, 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");
+    }
+}

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<IdentityUpdate>) -> Result<AssociationState, Assoc
 
 #[cfg(test)]
 mod tests {
-    use self::test_utils::{rand_string, rand_u64, rand_vec};
+    use self::test_utils::{rand_string, rand_u64, rand_vec, MockSignature};
 
     use super::*;
 
-    #[derive(Clone)]
-    struct MockSignature {
-        is_valid: bool,
-        signer_identity: MemberIdentifier,
-        signature_kind: SignatureKind,
-        signature_nonce: u64,
-    }
-
     impl IdentityUpdate {
         pub fn new_test(actions: Vec<Action>) -> 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<u64>,
-        ) -> Box<Self> {
-            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<MemberIdentifier, SignatureError> {
-            match self.is_valid {
-                true => Ok(self.signer_identity.clone()),
-                false => Err(SignatureError::Invalid),
-            }
-        }
-
-        fn bytes(&self) -> Vec<u8> {
-            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()
         });

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<MemberIdentifier, SignatureError>;
     fn signature_kind(&self) -> SignatureKind;
     fn bytes(&self) -> Vec<u8>;
 }
+
+pub trait SignatureClone {
+    fn clone_box(&self) -> Box<dyn Signature>;
+}
+
+impl<T> SignatureClone for T
+where
+    T: 'static + Signature + Clone,
+{
+    fn clone_box(&self) -> Box<dyn Signature> {
+        Box::new(self.clone())
+    }
+}
+
+impl Clone for Box<dyn Signature> {
+    fn clone(&self) -> Box<dyn Signature> {
+        self.clone_box()
+    }
+}