xmtp_db/encrypted_store/
key_package_history.rs

1use super::{
2    ConnectionExt, StorageError, db_connection::DbConnection, schema::key_package_history,
3};
4use crate::{StoreOrIgnore, impl_store_or_ignore};
5use diesel::prelude::*;
6use xmtp_common::time::now_ns;
7use xmtp_configuration::KEYS_EXPIRATION_INTERVAL_NS;
8
9#[derive(Insertable, Debug, Clone)]
10#[diesel(table_name = key_package_history)]
11pub struct NewKeyPackageHistoryEntry {
12    pub key_package_hash_ref: Vec<u8>,
13    pub post_quantum_public_key: Option<Vec<u8>>,
14    pub created_at_ns: i64,
15}
16
17#[derive(Queryable, Selectable, Debug, Clone)]
18#[diesel(table_name = key_package_history)]
19pub struct StoredKeyPackageHistoryEntry {
20    pub id: i32,
21    pub key_package_hash_ref: Vec<u8>,
22    pub created_at_ns: i64,
23    pub delete_at_ns: Option<i64>,
24    pub post_quantum_public_key: Option<Vec<u8>>,
25}
26
27impl_store_or_ignore!(NewKeyPackageHistoryEntry, key_package_history);
28
29pub trait QueryKeyPackageHistory {
30    fn store_key_package_history_entry(
31        &self,
32        key_package_hash_ref: Vec<u8>,
33        post_quantum_public_key: Option<Vec<u8>>,
34    ) -> Result<StoredKeyPackageHistoryEntry, StorageError>;
35
36    fn find_key_package_history_entry_by_hash_ref(
37        &self,
38        hash_ref: Vec<u8>,
39    ) -> Result<StoredKeyPackageHistoryEntry, StorageError>;
40
41    fn find_key_package_history_entries_before_id(
42        &self,
43        id: i32,
44    ) -> Result<Vec<StoredKeyPackageHistoryEntry>, StorageError>;
45
46    fn mark_key_package_before_id_to_be_deleted(&self, id: i32) -> Result<(), StorageError>;
47
48    fn get_expired_key_packages(&self) -> Result<Vec<StoredKeyPackageHistoryEntry>, StorageError>;
49
50    fn delete_key_package_history_up_to_id(&self, id: i32) -> Result<(), StorageError>;
51
52    fn delete_key_package_entry_with_id(&self, id: i32) -> Result<(), StorageError>;
53}
54
55impl<T> QueryKeyPackageHistory for &T
56where
57    T: QueryKeyPackageHistory,
58{
59    fn store_key_package_history_entry(
60        &self,
61        key_package_hash_ref: Vec<u8>,
62        post_quantum_public_key: Option<Vec<u8>>,
63    ) -> Result<StoredKeyPackageHistoryEntry, StorageError> {
64        (**self).store_key_package_history_entry(key_package_hash_ref, post_quantum_public_key)
65    }
66
67    fn find_key_package_history_entry_by_hash_ref(
68        &self,
69        hash_ref: Vec<u8>,
70    ) -> Result<StoredKeyPackageHistoryEntry, StorageError> {
71        (**self).find_key_package_history_entry_by_hash_ref(hash_ref)
72    }
73
74    fn find_key_package_history_entries_before_id(
75        &self,
76        id: i32,
77    ) -> Result<Vec<StoredKeyPackageHistoryEntry>, StorageError> {
78        (**self).find_key_package_history_entries_before_id(id)
79    }
80
81    fn mark_key_package_before_id_to_be_deleted(&self, id: i32) -> Result<(), StorageError> {
82        (**self).mark_key_package_before_id_to_be_deleted(id)
83    }
84
85    fn get_expired_key_packages(&self) -> Result<Vec<StoredKeyPackageHistoryEntry>, StorageError> {
86        (**self).get_expired_key_packages()
87    }
88
89    fn delete_key_package_history_up_to_id(&self, id: i32) -> Result<(), StorageError> {
90        (**self).delete_key_package_history_up_to_id(id)
91    }
92
93    fn delete_key_package_entry_with_id(&self, id: i32) -> Result<(), StorageError> {
94        (**self).delete_key_package_entry_with_id(id)
95    }
96}
97
98impl<C: ConnectionExt> QueryKeyPackageHistory for DbConnection<C> {
99    fn store_key_package_history_entry(
100        &self,
101        key_package_hash_ref: Vec<u8>,
102        post_quantum_public_key: Option<Vec<u8>>,
103    ) -> Result<StoredKeyPackageHistoryEntry, StorageError> {
104        let entry = NewKeyPackageHistoryEntry {
105            key_package_hash_ref: key_package_hash_ref.clone(),
106            post_quantum_public_key: post_quantum_public_key.clone(),
107            created_at_ns: now_ns(),
108        };
109        entry.store_or_ignore(self)?;
110
111        self.find_key_package_history_entry_by_hash_ref(key_package_hash_ref)
112    }
113
114    fn find_key_package_history_entry_by_hash_ref(
115        &self,
116        hash_ref: Vec<u8>,
117    ) -> Result<StoredKeyPackageHistoryEntry, StorageError> {
118        let result = self.raw_query_read(|conn| {
119            key_package_history::dsl::key_package_history
120                .filter(key_package_history::dsl::key_package_hash_ref.eq(hash_ref))
121                .first::<StoredKeyPackageHistoryEntry>(conn)
122        })?;
123
124        Ok(result)
125    }
126
127    fn find_key_package_history_entries_before_id(
128        &self,
129        id: i32,
130    ) -> Result<Vec<StoredKeyPackageHistoryEntry>, StorageError> {
131        let result = self.raw_query_read(|conn| {
132            key_package_history::dsl::key_package_history
133                .filter(key_package_history::dsl::id.lt(id))
134                .load::<StoredKeyPackageHistoryEntry>(conn)
135        })?;
136
137        Ok(result)
138    }
139
140    fn mark_key_package_before_id_to_be_deleted(&self, id: i32) -> Result<(), StorageError> {
141        use crate::schema::key_package_history::dsl;
142        let delete_at_24_hrs_ns = now_ns() + KEYS_EXPIRATION_INTERVAL_NS;
143        self.raw_query_write(|conn| {
144            diesel::update(
145                dsl::key_package_history
146                    .filter(dsl::id.lt(id))
147                    .filter(dsl::delete_at_ns.is_null()), // Only set if not already set
148            )
149            .set(dsl::delete_at_ns.eq(delete_at_24_hrs_ns))
150            .execute(conn)
151        })?;
152
153        Ok(())
154    }
155
156    fn get_expired_key_packages(&self) -> Result<Vec<StoredKeyPackageHistoryEntry>, StorageError> {
157        use crate::schema::key_package_history::dsl;
158        self.raw_query_read(|conn| {
159            dsl::key_package_history
160                .filter(dsl::delete_at_ns.le(now_ns()))
161                .load::<StoredKeyPackageHistoryEntry>(conn)
162        })
163        .map_err(StorageError::from) // convert ConnectionError into StorageError
164    }
165
166    fn delete_key_package_history_up_to_id(&self, id: i32) -> Result<(), StorageError> {
167        self.raw_query_write(|conn| {
168            diesel::delete(
169                key_package_history::dsl::key_package_history
170                    .filter(key_package_history::dsl::id.le(id)),
171            )
172            .execute(conn)
173        })?;
174
175        Ok(())
176    }
177
178    fn delete_key_package_entry_with_id(&self, id: i32) -> Result<(), StorageError> {
179        self.raw_query_write(|conn| {
180            diesel::delete(
181                key_package_history::dsl::key_package_history
182                    .filter(key_package_history::dsl::id.eq(id)),
183            )
184            .execute(conn)
185        })?;
186
187        Ok(())
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use crate::prelude::*;
194    use crate::test_utils::with_connection;
195    use xmtp_common::rand_vec;
196    #[cfg(target_arch = "wasm32")]
197    wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_dedicated_worker);
198
199    #[xmtp_common::test]
200    fn test_store_key_package_history_entry() {
201        with_connection(|conn| {
202            let hash_ref = rand_vec::<24>();
203            let post_quantum_public_key = rand_vec::<32>();
204            let new_entry = conn
205                .store_key_package_history_entry(
206                    hash_ref.clone(),
207                    Some(post_quantum_public_key.clone()),
208                )
209                .unwrap();
210            assert_eq!(new_entry.key_package_hash_ref, hash_ref);
211            assert_eq!(
212                new_entry.post_quantum_public_key,
213                Some(post_quantum_public_key)
214            );
215            assert_eq!(new_entry.id, 1);
216
217            // Now delete it
218            conn.delete_key_package_entry_with_id(1).unwrap();
219            let all_entries = conn
220                .find_key_package_history_entries_before_id(100)
221                .unwrap();
222            assert!(all_entries.is_empty());
223        })
224    }
225
226    #[xmtp_common::test]
227    fn test_store_multiple() {
228        with_connection(|conn| {
229            let post_quantum_public_key = rand_vec::<32>();
230            let hash_ref1 = rand_vec::<24>();
231            let hash_ref2 = rand_vec::<24>();
232            let hash_ref3 = rand_vec::<24>();
233
234            conn.store_key_package_history_entry(
235                hash_ref1.clone(),
236                Some(post_quantum_public_key.clone()),
237            )
238            .unwrap();
239            conn.store_key_package_history_entry(
240                hash_ref2.clone(),
241                Some(post_quantum_public_key.clone()),
242            )
243            .unwrap();
244            let entry_3 = conn
245                .store_key_package_history_entry(hash_ref3.clone(), None)
246                .unwrap();
247
248            let all_entries = conn
249                .find_key_package_history_entries_before_id(100)
250                .unwrap();
251
252            assert_eq!(all_entries.len(), 3);
253
254            let earlier_entries = conn
255                .find_key_package_history_entries_before_id(entry_3.id)
256                .unwrap();
257            assert_eq!(earlier_entries.len(), 2);
258        })
259    }
260}
xmtp_db/encrypted_store/key_package_history.rs

xmtp_db/encrypted_store/
key_package_history.rs
