Skip to main content

xmtp_id/associations/
state.rs

1//! [`AssociationState`] describes a single point in time for an Inbox where it contains a set of
2//! associated [`MemberIdentifier`]'s, which may be one of [`MemberKind::Address`]
3//! or[`MemberKind::Installation`]. A diff between two states can be calculated to determine
4//! a change of membership between two periods of time. [XIP-46](https://github.com/xmtp/XIPs/pull/53)
5
6use super::{
7    AssociationError, Handle, MemberIdentifier, MemberKind, ident,
8    member::{Identifier, Member},
9};
10use crate::InboxIdRef;
11use std::{
12    collections::{HashMap, HashSet},
13    fmt::{Debug, Write},
14};
15
16#[derive(Debug, Clone)]
17pub struct AssociationStateDiff {
18    pub new_members: Vec<MemberIdentifier>,
19    pub removed_members: Vec<MemberIdentifier>,
20}
21
22#[derive(Debug)]
23pub struct Installation {
24    pub id: Vec<u8>,
25    pub client_timestamp_ns: Option<u64>,
26}
27
28impl AssociationStateDiff {
29    pub fn new_installations(&self) -> Vec<Vec<u8>> {
30        self.new_members
31            .iter()
32            .filter_map(|member| match member {
33                MemberIdentifier::Installation(ident::Installation(key)) => Some(key.clone()),
34                _ => None,
35            })
36            .collect()
37    }
38
39    pub fn removed_installations(&self) -> Vec<Vec<u8>> {
40        self.removed_members
41            .iter()
42            .filter_map(|member| match member {
43                MemberIdentifier::Installation(ident::Installation(key)) => Some(key.clone()),
44                _ => None,
45            })
46            .collect()
47    }
48}
49
50#[derive(Clone)]
51pub struct AssociationState {
52    pub(crate) inbox_id: String,
53    pub(crate) members: HashMap<MemberIdentifier, Member>,
54    pub(crate) recovery_identifier: Identifier,
55    pub(crate) seen_signatures: HashSet<Vec<u8>>,
56    /// Claimed handles in claim order; the last is the display handle.
57    pub(crate) handles: Vec<Handle>,
58    /// The number of `ClaimHandle` actions applied to produce this state.
59    pub(crate) handle_claims: u32,
60}
61
62impl std::fmt::Debug for AssociationState {
63    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        let mut members = String::new();
65        for member in self.members.keys() {
66            write!(members, "{:?}", member)?;
67            write!(members, ",")?;
68        }
69
70        let mut signatures = String::new();
71        for signature in self.seen_signatures.iter() {
72            write!(
73                signatures,
74                "{}",
75                xmtp_common::fmt::truncate_hex(hex::encode(signature))
76            )?;
77            write!(signatures, ",")?;
78        }
79
80        // A handle names an inbox, so a log line carries only the count.
81        write!(
82            f,
83            "AssociationState {{ inbox_id: {}, members: {}, recovery: {}, seen_signatures: {}, handles: {}, handle_claims: {} }}",
84            self.inbox_id,
85            members,
86            self.recovery_identifier,
87            signatures,
88            self.handles.len(),
89            self.handle_claims
90        )
91    }
92}
93
94impl TryFrom<MemberIdentifier> for Identifier {
95    type Error = AssociationError;
96    fn try_from(ident: MemberIdentifier) -> Result<Self, Self::Error> {
97        let ident = match ident {
98            MemberIdentifier::Ethereum(eth) => Self::Ethereum(eth),
99            MemberIdentifier::Passkey(passkey) => Self::Passkey(passkey),
100            MemberIdentifier::Installation(_) => {
101                return Err(AssociationError::NotIdentifier(
102                    "Installation Keys".to_string(),
103                ));
104            }
105        };
106        Ok(ident)
107    }
108}
109
110impl AssociationState {
111    pub fn add(&self, member: Member) -> Self {
112        let mut new_state = self.clone();
113        let _ = new_state.members.insert(member.identifier.clone(), member);
114
115        new_state
116    }
117
118    pub fn remove(&self, identifier: &MemberIdentifier) -> Self {
119        let mut new_state = self.clone();
120        let _ = new_state.members.remove(identifier);
121
122        new_state
123    }
124
125    pub fn set_recovery_identifier(&self, recovery_identifier: Identifier) -> Self {
126        let mut new_state = self.clone();
127        new_state.recovery_identifier = recovery_identifier;
128
129        new_state
130    }
131
132    pub fn get(&self, identifier: &MemberIdentifier) -> Option<&Member> {
133        self.members.get(identifier)
134    }
135
136    pub fn add_seen_signatures(&self, signatures: Vec<Vec<u8>>) -> Self {
137        let mut new_state = self.clone();
138        new_state.seen_signatures.extend(signatures);
139
140        new_state
141    }
142
143    pub fn has_seen(&self, signature: &Vec<u8>) -> bool {
144        self.seen_signatures.contains(signature)
145    }
146
147    /// The state after a claim of `handle`: `handle` moves to the end of the
148    /// handles, or is added there, and the claim count rises by one. Members,
149    /// the recovery identifier, and chain ids are unchanged.
150    pub fn claim_handle(&self, handle: Handle) -> Self {
151        let mut new_state = self.clone();
152        new_state.handles.retain(|held| *held != handle);
153        new_state.handles.push(handle);
154        new_state.handle_claims += 1;
155
156        new_state
157    }
158
159    /// Every handle the inbox has claimed, in claim order, with a handle
160    /// claimed again at the position of its last claim.
161    pub fn handles(&self) -> &[Handle] {
162        &self.handles
163    }
164
165    /// The number of `ClaimHandle` actions applied to produce this state.
166    /// A claim of a handle the inbox already holds counts.
167    pub fn handle_claims(&self) -> u32 {
168        self.handle_claims
169    }
170
171    /// The inbox's current handle: the last one it claimed, or `None` when it
172    /// has claimed none.
173    pub fn display_handle(&self) -> Option<&Handle> {
174        self.handles.last()
175    }
176
177    /// The handles other than the display handle, the most recently claimed
178    /// first.
179    pub fn other_handles(&self) -> impl Iterator<Item = &Handle> {
180        self.handles.iter().rev().skip(1)
181    }
182
183    pub fn members(&self) -> Vec<Member> {
184        let mut sorted_members: Vec<_> = self.members.values().cloned().collect();
185        sorted_members.sort_by_key(|m| m.client_timestamp_ns.unwrap_or(u64::MAX));
186        sorted_members
187    }
188
189    pub fn inbox_id(&self) -> InboxIdRef<'_> {
190        &self.inbox_id
191    }
192
193    pub fn recovery_identifier(&self) -> &Identifier {
194        &self.recovery_identifier
195    }
196
197    pub fn members_by_parent(&self, parent_id: &MemberIdentifier) -> Vec<Member> {
198        self.members
199            .values()
200            .filter(|e| e.added_by_entity.eq(&Some(parent_id.clone())))
201            .cloned()
202            .collect()
203    }
204
205    pub fn members_by_kind(&self, kind: MemberKind) -> Vec<Member> {
206        self.members
207            .values()
208            .filter(|e| e.kind() == kind)
209            .cloned()
210            .collect()
211    }
212
213    pub fn identifiers(&self) -> Vec<Identifier> {
214        let mut address_members: Vec<_> = self.members.values().cloned().collect();
215
216        address_members.sort_by_key(|m| m.client_timestamp_ns.unwrap_or(u64::MAX));
217
218        address_members
219            .into_iter()
220            .filter_map(|member| match member.identifier {
221                MemberIdentifier::Ethereum(eth) => Some(Identifier::Ethereum(eth)),
222                MemberIdentifier::Passkey(pk) => Some(Identifier::Passkey(pk)),
223                _ => None,
224            })
225            .collect()
226    }
227
228    pub fn installation_ids(&self) -> Vec<Vec<u8>> {
229        self.members_by_kind(MemberKind::Installation)
230            .into_iter()
231            .filter_map(|member| match member.identifier {
232                MemberIdentifier::Installation(ident::Installation(key)) => Some(key),
233                _ => None,
234            })
235            .collect()
236    }
237
238    pub fn installations(&self) -> Vec<Installation> {
239        self.members()
240            .into_iter()
241            .filter_map(|member| match member.identifier {
242                MemberIdentifier::Installation(ident::Installation(id)) => Some(Installation {
243                    id,
244                    client_timestamp_ns: member.client_timestamp_ns,
245                }),
246                _ => None,
247            })
248            .collect()
249    }
250
251    pub fn diff(&self, new_state: &Self) -> AssociationStateDiff {
252        let new_members: Vec<MemberIdentifier> = new_state
253            .members
254            .keys()
255            .filter(|new_member_identifier| !self.members.contains_key(new_member_identifier))
256            .cloned()
257            .collect();
258
259        let removed_members: Vec<MemberIdentifier> = self
260            .members
261            .keys()
262            .filter(|existing_member_identifier| {
263                !new_state.members.contains_key(existing_member_identifier)
264            })
265            .cloned()
266            .collect();
267
268        AssociationStateDiff {
269            new_members,
270            removed_members,
271        }
272    }
273
274    /// Converts the [`AssociationState`] to a diff that represents all members
275    /// of the inbox at the current state.
276    pub fn as_diff(&self) -> AssociationStateDiff {
277        AssociationStateDiff {
278            new_members: self.members.keys().cloned().collect(),
279            removed_members: vec![],
280        }
281    }
282
283    pub fn new(
284        account_identifier: Identifier,
285        nonce: u64,
286        chain_id: Option<u64>,
287    ) -> Result<Self, AssociationError> {
288        let member_identifier: MemberIdentifier = account_identifier.clone().into();
289
290        let inbox_id = account_identifier.inbox_id(nonce)?;
291        let new_member = Member::new(member_identifier.clone(), None, None, chain_id);
292        Ok(Self {
293            members: HashMap::from_iter([(member_identifier, new_member)]),
294            seen_signatures: HashSet::new(),
295            recovery_identifier: account_identifier,
296            inbox_id,
297            handles: Vec::new(),
298            handle_claims: 0,
299        })
300    }
301}
302
303#[cfg(test)]
304pub(crate) mod tests {
305    use super::*;
306
307    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
308    #[cfg_attr(not(target_arch = "wasm32"), test)]
309    fn can_add_remove() {
310        let starting_state = AssociationState::new(Identifier::rand_ethereum(), 0, None).unwrap();
311        let new_entity = Member::default();
312        let with_add = starting_state.add(new_entity.clone());
313        assert!(with_add.get(&new_entity.identifier).is_some());
314        assert!(starting_state.get(&new_entity.identifier).is_none());
315    }
316
317    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
318    #[cfg_attr(not(target_arch = "wasm32"), test)]
319    fn can_diff() {
320        let starting_state = AssociationState::new(Identifier::rand_ethereum(), 0, None).unwrap();
321        let entity_1 = Member::default();
322        let entity_2 = Member::default();
323        let entity_3 = Member::default();
324
325        let state_1 = starting_state.add(entity_1.clone()).add(entity_2.clone());
326        let state_2 = state_1.remove(&entity_1.identifier).add(entity_3.clone());
327
328        let diff = state_1.diff(&state_2);
329
330        assert_eq!(diff.new_members, vec![entity_3.identifier]);
331        assert_eq!(diff.removed_members, vec![entity_1.identifier]);
332    }
333}