xmtp_db/encrypted_store/
identity_cache.rs

1use super::schema::identity_cache;
2use super::{ConnectionExt, Sqlite};
3use crate::{DbConnection, StorageError};
4use crate::{Store, impl_fetch, impl_store};
5use diesel::backend::Backend;
6use diesel::deserialize::{self, FromSql, FromSqlRow};
7use diesel::expression::AsExpression;
8use diesel::serialize::{IsNull, Output, ToSql};
9use diesel::sql_types::Integer;
10use diesel::{Insertable, Queryable};
11use diesel::{prelude::*, serialize};
12use serde::{Deserialize, Serialize};
13use std::any::type_name;
14use std::collections::HashMap;
15use xmtp_proto::ConversionError;
16use xmtp_proto::xmtp::identity::associations::IdentifierKind;
17
18#[derive(Insertable, Queryable, Debug, Clone, Deserialize, Serialize)]
19#[diesel(table_name = identity_cache)]
20#[diesel()]
21pub struct IdentityCache {
22    inbox_id: String,
23    identity: String,
24    identity_kind: StoredIdentityKind,
25}
26
27#[repr(i32)]
28#[derive(Debug, Copy, Clone, Serialize, Deserialize, Eq, PartialEq, AsExpression, FromSqlRow)]
29#[diesel(sql_type = Integer)]
30/// Type of identity stored
31pub enum StoredIdentityKind {
32    Ethereum = 1,
33    Passkey = 2,
34}
35
36impl TryFrom<IdentifierKind> for StoredIdentityKind {
37    type Error = xmtp_proto::ConversionError;
38    fn try_from(kind: IdentifierKind) -> Result<Self, Self::Error> {
39        match kind {
40            IdentifierKind::Ethereum => Ok(StoredIdentityKind::Ethereum),
41            IdentifierKind::Passkey => Ok(StoredIdentityKind::Passkey),
42            IdentifierKind::Unspecified => {
43                Err(ConversionError::Unspecified("IdentifierKind::Unspecified"))
44            }
45        }
46    }
47}
48
49impl TryFrom<i32> for StoredIdentityKind {
50    type Error = ConversionError;
51
52    fn try_from(value: i32) -> Result<Self, Self::Error> {
53        match value {
54            1 => Ok(StoredIdentityKind::Ethereum),
55            2 => Ok(StoredIdentityKind::Passkey),
56            v => Err(ConversionError::InvalidValue {
57                item: type_name::<StoredIdentityKind>(),
58                expected: "a integer value of `1` or `2`",
59                got: v.to_string(),
60            }),
61        }
62    }
63}
64
65impl From<&StoredIdentityKind> for i32 {
66    fn from(value: &StoredIdentityKind) -> Self {
67        use StoredIdentityKind::*;
68        match value {
69            Ethereum => 1,
70            Passkey => 2,
71        }
72    }
73}
74
75impl From<StoredIdentityKind> for IdentifierKind {
76    fn from(value: StoredIdentityKind) -> Self {
77        use StoredIdentityKind::*;
78        match value {
79            Ethereum => IdentifierKind::Ethereum,
80            Passkey => IdentifierKind::Passkey,
81        }
82    }
83}
84
85impl_store!(IdentityCache, identity_cache);
86impl_fetch!(IdentityCache, identity_cache);
87
88pub trait QueryIdentityCache {
89    /// Returns a HashMap of WalletAddress -> InboxId
90    fn fetch_cached_inbox_ids(
91        &self,
92        identifiers: &[(Address, StoredIdentityKind)],
93    ) -> Result<HashMap<String, String>, StorageError>;
94
95    fn cache_inbox_id<S: ToString>(
96        &self,
97        kind: StoredIdentityKind,
98        identity: String,
99        inbox_id: S,
100    ) -> Result<(), StorageError>;
101}
102
103impl<G> QueryIdentityCache for &G
104where
105    G: QueryIdentityCache,
106{
107    fn fetch_cached_inbox_ids(
108        &self,
109        identifiers: &[(Address, StoredIdentityKind)],
110    ) -> Result<HashMap<String, String>, StorageError> {
111        (**self).fetch_cached_inbox_ids(identifiers)
112    }
113
114    fn cache_inbox_id<S: ToString>(
115        &self,
116        kind: StoredIdentityKind,
117        identity: String,
118        inbox_id: S,
119    ) -> Result<(), StorageError> {
120        (**self).cache_inbox_id(kind, identity, inbox_id)
121    }
122}
123
124type Address = String;
125
126impl<C: ConnectionExt> QueryIdentityCache for DbConnection<C> {
127    /// Returns a HashMap of WalletAddress -> InboxId
128    fn fetch_cached_inbox_ids(
129        &self,
130        identifiers: &[(Address, StoredIdentityKind)],
131    ) -> Result<HashMap<String, String>, StorageError> {
132        use crate::encrypted_store::schema::identity_cache::*;
133
134        let mut conditions = identity_cache::table.into_boxed();
135
136        for (addr, ident) in identifiers {
137            let kind: i32 = ident.into();
138            let cond = identity.eq(addr).and(identity_kind.eq(kind));
139            conditions = conditions.or_filter(cond);
140        }
141
142        let result = self
143            .raw_query(|conn| conditions.load::<IdentityCache>(conn))?
144            .into_iter()
145            .map(|entry| (entry.identity, entry.inbox_id))
146            .collect();
147        Ok(result)
148    }
149
150    fn cache_inbox_id<S: ToString>(
151        &self,
152        kind: StoredIdentityKind,
153        identity: String,
154        inbox_id: S,
155    ) -> Result<(), StorageError> {
156        IdentityCache {
157            inbox_id: inbox_id.to_string(),
158            identity,
159            identity_kind: kind,
160        }
161        .store(self)
162    }
163}
164
165impl ToSql<Integer, Sqlite> for StoredIdentityKind
166where
167    i32: ToSql<Integer, Sqlite>,
168{
169    fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Sqlite>) -> serialize::Result {
170        out.set_value(*self as i32);
171        Ok(IsNull::No)
172    }
173}
174
175impl FromSql<Integer, Sqlite> for StoredIdentityKind
176where
177    i32: FromSql<Integer, Sqlite>,
178{
179    fn from_sql(bytes: <Sqlite as Backend>::RawValue<'_>) -> deserialize::Result<Self> {
180        match i32::from_sql(bytes)? {
181            1 => Ok(Self::Ethereum),
182            2 => Ok(Self::Passkey),
183            x => Err(format!("Unrecognized variant {}", x).into()),
184        }
185    }
186}
187
188#[cfg(test)]
189pub(crate) mod tests {
190    use super::IdentityCache;
191    use crate::{
192        Store, identity_cache::StoredIdentityKind, prelude::*, test_utils::with_connection,
193    };
194
195    #[derive(Clone)]
196    struct MockIdentity {
197        identity: String,
198        inbox_id: String,
199    }
200
201    impl MockIdentity {
202        fn create() -> Self {
203            Self {
204                identity: xmtp_common::rand_hexstring(),
205                inbox_id: xmtp_common::rand_string::<32>(),
206            }
207        }
208    }
209
210    // Test storing duplicated wallets (same inbox_id and wallet_address)
211    #[xmtp_common::test]
212    fn test_store_duplicated_wallets() {
213        with_connection(|conn| {
214            let entry1 = IdentityCache {
215                inbox_id: "test_dup".to_string(),
216                identity: "wallet_dup".to_string(),
217                identity_kind: StoredIdentityKind::Ethereum,
218            };
219            let entry2 = IdentityCache {
220                inbox_id: "test_dup".to_string(),
221                identity: "wallet_dup".to_string(),
222                identity_kind: StoredIdentityKind::Ethereum,
223            };
224            entry1.store(conn).expect("Failed to store wallet");
225            let result = entry2.store(conn);
226            assert!(
227                result.is_err(),
228                "Duplicated wallet stored without error, expected failure"
229            );
230        })
231    }
232
233    // Test storing and fetching multiple wallet addresses with multiple keys
234    #[xmtp_common::test]
235    fn test_fetch_and_store_identity_cache() {
236        with_connection(|conn| {
237            let ident1 = MockIdentity::create();
238            let ident2 = MockIdentity::create();
239
240            conn.cache_inbox_id(
241                StoredIdentityKind::Ethereum,
242                ident1.identity.clone(),
243                &ident1.inbox_id,
244            )
245            .unwrap();
246
247            let idents = &[
248                (ident1.identity.clone(), StoredIdentityKind::Ethereum),
249                (ident2.identity.clone(), StoredIdentityKind::Ethereum),
250            ];
251            let stored_wallets = conn.fetch_cached_inbox_ids(idents).unwrap();
252
253            // Verify that 1 entries are fetched
254            assert_eq!(stored_wallets.len(), 1);
255
256            // Verify it's the correct inbox_id
257            let cached_inbox_id = stored_wallets.get(&idents[0].0).unwrap();
258            assert_eq!(*cached_inbox_id, ident1.inbox_id);
259
260            // Fetch wallets with a non-existent list of inbox_ids
261            let ident = MockIdentity::create();
262            let non_existent_wallets = conn
263                .fetch_cached_inbox_ids(&[(ident.identity, StoredIdentityKind::Ethereum)])
264                .unwrap_or_default();
265            assert!(
266                non_existent_wallets.is_empty(),
267                "Expected no wallets, found some"
268            );
269        })
270    }
271}
xmtp_db/encrypted_store/identity_cache.rs

xmtp_db/encrypted_store/
identity_cache.rs
