xmtp_mls_validation/
group_membership.rs

1use prost::{DecodeError, Message};
2use std::collections::HashMap;
3use xmtp_proto::xmtp::mls::message_contents::GroupMembership as GroupMembershipProto;
4
5#[derive(Debug, Clone, PartialEq)]
6pub struct GroupMembership {
7    pub members: HashMap<String, u64>,
8    pub failed_installations: Vec<Vec<u8>>,
9}
10
11impl GroupMembership {
12    pub fn new() -> Self {
13        GroupMembership {
14            members: HashMap::new(),
15            failed_installations: Vec::new(),
16        }
17    }
18
19    pub fn add(&mut self, inbox_id: String, last_sequence_id: u64) {
20        self.members.insert(inbox_id, last_sequence_id);
21    }
22
23    pub fn remove<InboxId: AsRef<str>>(&mut self, inbox_id: InboxId) {
24        self.members.remove(inbox_id.as_ref());
25    }
26
27    pub fn get<InboxId: AsRef<str>>(&self, inbox_id: InboxId) -> Option<&u64> {
28        self.members.get(inbox_id.as_ref())
29    }
30
31    pub fn inbox_ids(&self) -> Vec<&str> {
32        self.members.keys().map(AsRef::as_ref).collect()
33    }
34
35    // Convert the mapping to a vector of `inbox_id`/`sequence_id` tuples
36    pub fn to_filters(&self) -> Vec<(&str, i64)> {
37        self.members
38            .iter()
39            .map(|(inbox_id, sequence_id)| (inbox_id.as_str(), *sequence_id as i64))
40            .collect()
41    }
42
43    pub fn diff<'inbox_id>(
44        &'inbox_id self,
45        new_group_membership: &'inbox_id Self,
46    ) -> MembershipDiff<'inbox_id> {
47        let mut removed_inboxes: Vec<&String> = vec![];
48        let mut updated_inboxes: Vec<&String> = vec![];
49
50        for (inbox_id, last_sequence_id) in self.members.iter() {
51            match new_group_membership.get(inbox_id) {
52                Some(new_last_sequence_id) => {
53                    if new_last_sequence_id.ne(last_sequence_id) {
54                        updated_inboxes.push(inbox_id);
55                    }
56                }
57                None => {
58                    removed_inboxes.push(inbox_id);
59                }
60            }
61        }
62
63        let added_inboxes = new_group_membership
64            .members
65            .keys()
66            .filter(|&inbox_id| !self.members.contains_key(inbox_id))
67            .collect::<Vec<&String>>();
68
69        MembershipDiff {
70            added_inboxes,
71            removed_inboxes,
72            updated_inboxes,
73        }
74    }
75}
76
77impl Default for GroupMembership {
78    fn default() -> Self {
79        GroupMembership::new()
80    }
81}
82
83impl TryFrom<Vec<u8>> for GroupMembership {
84    type Error = DecodeError;
85
86    fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
87        let membership_proto = GroupMembershipProto::decode(value.as_slice())?;
88
89        Ok(GroupMembership {
90            members: membership_proto.members,
91            failed_installations: membership_proto.failed_installations,
92        })
93    }
94}
95
96impl From<&GroupMembership> for Vec<u8> {
97    fn from(value: &GroupMembership) -> Self {
98        let membership_proto = GroupMembershipProto {
99            members: value.members.clone(),
100            failed_installations: value.failed_installations.clone(),
101        };
102
103        membership_proto.encode_to_vec()
104    }
105}
106
107#[derive(Debug, Clone)]
108pub struct MembershipDiff<'inbox_id> {
109    pub added_inboxes: Vec<&'inbox_id String>,
110    pub removed_inboxes: Vec<&'inbox_id String>,
111    pub updated_inboxes: Vec<&'inbox_id String>,
112}
113
114#[cfg(test)]
115pub(crate) mod tests {
116    use super::GroupMembership;
117
118    #[xmtp_common::test]
119    fn test_equality_works() {
120        let inbox_id_1 = "inbox_1".to_string();
121        let sequence_id_1: u64 = 1;
122        let mut member_map_1 = GroupMembership::new();
123        let mut member_map_2 = GroupMembership::new();
124
125        member_map_1.add(inbox_id_1.clone(), sequence_id_1);
126
127        assert!(member_map_1.ne(&member_map_2));
128
129        member_map_2.add(inbox_id_1.clone(), sequence_id_1);
130        assert!(member_map_1.eq(&member_map_2));
131
132        // Now change the sequence ID and make sure it is not equal again
133        member_map_2.add(inbox_id_1.clone(), 2);
134        assert!(member_map_1.ne(&member_map_2));
135    }
136
137    #[xmtp_common::test]
138    fn test_diff() {
139        let mut initial_members = GroupMembership::new();
140        initial_members.add("inbox_1".into(), 1);
141        initial_members.add("inbox_2".into(), 1);
142
143        let mut updated_list = initial_members.clone();
144        updated_list.remove("inbox_1");
145        updated_list.add("inbox_2".into(), 2);
146        updated_list.add("inbox_3".into(), 1);
147
148        let diff = initial_members.diff(&updated_list);
149        assert_eq!(diff.added_inboxes, vec!["inbox_3"]);
150        assert_eq!(diff.updated_inboxes, vec!["inbox_2"]);
151        assert_eq!(diff.removed_inboxes, vec!["inbox_1"]);
152    }
153}
xmtp_mls_validation/group_membership.rs

xmtp_mls_validation/
group_membership.rs
