1use 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 pub(crate) handles: Vec<Handle>,
58 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 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 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 pub fn handles(&self) -> &[Handle] {
162 &self.handles
163 }
164
165 pub fn handle_claims(&self) -> u32 {
168 self.handle_claims
169 }
170
171 pub fn display_handle(&self) -> Option<&Handle> {
174 self.handles.last()
175 }
176
177 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 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}