Association History Processing (Part 1)

xmtp_id/Cargo.toml

@@ -21,4 +21,7 @@ chrono.workspace = true
 serde.workspace = true
 async-trait.workspace = true
 futures.workspace = true
+sha2 = "0.10.8"
+rand.workspace = true
+hex.workspace = true
 

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: &u64) -> String {
+    sha256_string(format!("{}{}", account_address, nonce))
+}

xmtp_id/src/associations/member.rs

@@ -0,0 +1,120 @@
+#[derive(Clone, Debug, PartialEq)]
+pub enum MemberKind {
+    Installation,
+    Address,
+}
+
+impl std::fmt::Display for MemberKind {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            MemberKind::Installation => write!(f, "installation"),
+            MemberKind::Address => write!(f, "address"),
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub enum MemberIdentifier {
+    Address(String),
+    Installation(Vec<u8>),
+}
+
+impl MemberIdentifier {
+    pub fn kind(&self) -> MemberKind {
+        match self {
+            MemberIdentifier::Address(_) => MemberKind::Address,
+            MemberIdentifier::Installation(_) => MemberKind::Installation,
+        }
+    }
+}
+
+impl std::fmt::Display for MemberIdentifier {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        let as_string = match self {
+            MemberIdentifier::Address(address) => address.to_string(),
+            MemberIdentifier::Installation(installation) => hex::encode(installation),
+        };
+
+        write!(f, "{}", as_string)
+    }
+}
+
+impl From<String> for MemberIdentifier {
+    fn from(address: String) -> Self {
+        MemberIdentifier::Address(address)
+    }
+}
+
+impl From<Vec<u8>> for MemberIdentifier {
+    fn from(installation: Vec<u8>) -> Self {
+        MemberIdentifier::Installation(installation)
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct Member {
+    pub identifier: MemberIdentifier,
+    pub added_by_entity: Option<MemberIdentifier>,
+}
+
+impl Member {
+    pub fn new(identifier: MemberIdentifier, added_by_entity: Option<MemberIdentifier>) -> Self {
+        Self {
+            identifier,
+            added_by_entity,
+        }
+    }
+
+    pub fn kind(&self) -> MemberKind {
+        self.identifier.kind()
+    }
+}
+
+impl PartialEq<MemberIdentifier> for Member {
+    fn eq(&self, other: &MemberIdentifier) -> bool {
+        self.identifier.eq(other)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::associations::test_utils;
+
+    use super::*;
+
+    use test_utils::rand_string;
+
+    impl Default for MemberIdentifier {
+        fn default() -> Self {
+            MemberIdentifier::Address(rand_string())
+        }
+    }
+
+    impl Default for Member {
+        fn default() -> Self {
+            Self {
+                identifier: MemberIdentifier::default(),
+                added_by_entity: None,
+            }
+        }
+    }
+
+    #[test]
+    fn test_identifier_comparisons() {
+        let address_1 = MemberIdentifier::Address("0x123".to_string());
+        let address_2 = MemberIdentifier::Address("0x456".to_string());
+        let address_1_copy = MemberIdentifier::Address("0x123".to_string());
+
+        assert!(address_1 != address_2);
+        assert!(address_1.ne(&address_2));
+        assert!(address_1 == address_1_copy);
+
+        let installation_1 = MemberIdentifier::Installation(vec![1, 2, 3]);
+        let installation_2 = MemberIdentifier::Installation(vec![4, 5, 6]);
+        let installation_1_copy = MemberIdentifier::Installation(vec![1, 2, 3]);
+
+        assert!(installation_1 != installation_2);
+        assert!(installation_1.ne(&installation_2));
+        assert!(installation_1 == installation_1_copy);
+    }
+}

xmtp_id/src/associations/mod.rs

@@ -0,0 +1,8 @@
+mod hashes;
+mod member;
+mod state;
+#[cfg(test)]
+mod test_utils;
+
+pub use self::member::{Member, MemberIdentifier, MemberKind};
+pub use self::state::AssociationState;

xmtp_id/src/associations/state.rs

@@ -0,0 +1,109 @@
+use std::collections::{HashMap, HashSet};
+
+use super::{hashes::generate_xid, member::Member, MemberIdentifier, MemberKind};
+
+#[derive(Clone, Debug)]
+pub struct AssociationState {
+    xid: String,
+    members: HashMap<MemberIdentifier, Member>,
+    recovery_address: String,
+    seen_signatures: HashSet<Vec<u8>>,
+}
+
+impl AssociationState {
+    pub fn add(&self, member: Member) -> Self {
+        let mut new_state = self.clone();
+        let _ = new_state.members.insert(member.identifier.clone(), member);
+
+        new_state
+    }
+
+    pub fn remove(&self, identifier: &MemberIdentifier) -> Self {
+        let mut new_state = self.clone();
+        let _ = new_state.members.remove(identifier);
+
+        new_state
+    }
+
+    pub fn set_recovery_address(&self, recovery_address: String) -> Self {
+        let mut new_state = self.clone();
+        new_state.recovery_address = recovery_address;
+
+        new_state
+    }
+
+    pub fn get(&self, identifier: &MemberIdentifier) -> Option<Member> {
+        self.members.get(identifier).cloned()
+    }
+
+    pub fn add_seen_signatures(&self, signatures: Vec<Vec<u8>>) -> Self {
+        let mut new_state = self.clone();
+        new_state.seen_signatures.extend(signatures);
+
+        new_state
+    }
+
+    pub fn has_seen(&self, signature: &Vec<u8>) -> bool {
+        self.seen_signatures.contains(signature)
+    }
+
+    pub fn members(&self) -> Vec<Member> {
+        self.members.values().cloned().collect()
+    }
+
+    pub fn xid(&self) -> &String {
+        &self.xid
+    }
+
+    pub fn recovery_address(&self) -> &String {
+        &self.recovery_address
+    }
+
+    pub fn members_by_parent(&self, parent_id: &MemberIdentifier) -> Vec<Member> {
+        self.members
+            .values()
+            .filter(|e| e.added_by_entity.eq(&Some(parent_id.clone())))
+            .cloned()
+            .collect()
+    }
+
+    pub fn members_by_kind(&self, kind: MemberKind) -> Vec<Member> {
+        self.members
+            .values()
+            .filter(|e| e.kind() == kind)
+            .cloned()
+            .collect()
+    }
+
+    pub fn new(account_address: String, nonce: u64) -> Self {
+        let xid = generate_xid(&account_address, &nonce);
+        let identifier = MemberIdentifier::Address(account_address.clone());
+        let new_member = Member::new(identifier.clone(), None);
+        Self {
+            members: {
+                let mut members = HashMap::new();
+                members.insert(identifier, new_member);
+                members
+            },
+            seen_signatures: HashSet::new(),
+            recovery_address: account_address,
+            xid,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::associations::test_utils::rand_string;
+
+    use super::*;
+
+    #[test]
+    fn can_add_remove() {
+        let starting_state = AssociationState::new(rand_string(), 0);
+        let new_entity = Member::default();
+        let with_add = starting_state.add(new_entity.clone());
+        assert!(with_add.get(&new_entity.identifier).is_some());
+        assert!(starting_state.get(&new_entity.identifier).is_none());
+    }
+}

xmtp_id/src/associations/test_utils.rs

@@ -0,0 +1,21 @@
+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_u64() -> u64 {
+    rand::thread_rng().gen()
+}
+
+pub fn rand_vec() -> Vec<u8> {
+    let mut buf = [0u8; 32];
+    rand::thread_rng().fill(&mut buf[..]);
+    buf.to_vec()
+}

xmtp_id/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod associations;
 pub mod credential_verifier;
 pub mod verified_key_package;