xmtp_db/encrypted_store/
identity.rs

1use crate::encrypted_store::schema::identity;
2use crate::schema::identity::dsl;
3use crate::{ConnectionExt, DbConnection, StorageError, impl_fetch, impl_store};
4use derive_builder::Builder;
5use diesel::prelude::*;
6use serde::{Deserialize, Serialize};
7use xmtp_common::time::now_ns;
8use xmtp_configuration::KEY_PACKAGE_QUEUE_INTERVAL_NS;
9
10/// Identity of this installation
11/// There can only be one.
12#[derive(Insertable, Queryable, Debug, Clone, Builder, Serialize, Deserialize)]
13#[diesel(table_name = identity)]
14#[builder(setter(into), build_fn(error = "crate::StorageError"))]
15pub struct StoredIdentity {
16    pub inbox_id: String,
17    pub installation_keys: Vec<u8>,
18    pub credential_bytes: Vec<u8>,
19    #[builder(setter(skip))]
20    rowid: Option<i32>,
21    pub next_key_package_rotation_ns: Option<i64>,
22    #[builder(default)]
23    pub registration_cursor_sequence_id: Option<i64>,
24}
25
26impl_fetch!(StoredIdentity, identity);
27impl_store!(StoredIdentity, identity);
28
29impl StoredIdentity {
30    pub fn builder() -> StoredIdentityBuilder {
31        StoredIdentityBuilder::default()
32    }
33
34    pub fn new(inbox_id: String, installation_keys: Vec<u8>, credential_bytes: Vec<u8>) -> Self {
35        Self {
36            inbox_id,
37            installation_keys,
38            credential_bytes,
39            rowid: None,
40            next_key_package_rotation_ns: None,
41
42            registration_cursor_sequence_id: None,
43        }
44    }
45}
46pub trait QueryIdentity {
47    fn queue_key_package_rotation(&self) -> Result<(), StorageError>;
48    /// Atomically lower/initialize the rotation column (5s debounce) AND enqueue a
49    /// `PullInDeadline` task targeting `rotation_task_hash` at the resulting column
50    /// value — one transaction, so neither write can land without the other.
51    /// `rotation_seed` is insert-or-ignored first so the pull-in always has a live
52    /// target (commit-target-first), even if startup seeding never ran.
53    /// Callers wake the TaskWorker AFTER this returns (never inside a tx).
54    fn queue_key_rotation_with_nudge(
55        &self,
56        rotation_task_hash: &crate::tasks::TaskDataHash,
57        rotation_seed: crate::tasks::NewTask,
58    ) -> Result<(), StorageError>;
59    fn reset_key_package_rotation_queue(
60        &self,
61        rotation_interval_ns: i64,
62    ) -> Result<(), StorageError>;
63    fn is_identity_needs_rotation(&self) -> Result<bool, StorageError>;
64    /// The identity's absolute rotation deadline (`next_key_package_rotation_ns`).
65    /// `None` if NULL or if no identity row exists yet (indistinguishable to callers;
66    /// treat as "no scheduled deadline").
67    fn next_key_package_rotation_ns(&self) -> Result<Option<i64>, StorageError>;
68}
69
70impl<T> QueryIdentity for &T
71where
72    T: QueryIdentity,
73{
74    fn queue_key_package_rotation(&self) -> Result<(), StorageError> {
75        (**self).queue_key_package_rotation()
76    }
77
78    fn queue_key_rotation_with_nudge(
79        &self,
80        rotation_task_hash: &crate::tasks::TaskDataHash,
81        rotation_seed: crate::tasks::NewTask,
82    ) -> Result<(), StorageError> {
83        (**self).queue_key_rotation_with_nudge(rotation_task_hash, rotation_seed)
84    }
85
86    fn reset_key_package_rotation_queue(
87        &self,
88        rotation_interval_ns: i64,
89    ) -> Result<(), StorageError> {
90        (**self).reset_key_package_rotation_queue(rotation_interval_ns)
91    }
92
93    fn is_identity_needs_rotation(&self) -> Result<bool, StorageError> {
94        (**self).is_identity_needs_rotation()
95    }
96
97    fn next_key_package_rotation_ns(&self) -> Result<Option<i64>, StorageError> {
98        (**self).next_key_package_rotation_ns()
99    }
100}
101
102impl<C: ConnectionExt> QueryIdentity for DbConnection<C> {
103    fn queue_key_package_rotation(&self) -> Result<(), StorageError> {
104        self.raw_query(|conn| {
105            let rotate_at_ns = now_ns() + KEY_PACKAGE_QUEUE_INTERVAL_NS;
106            // NULL (migrated DBs) counts as unscheduled: initialize it here so the
107            // 5s debounce applies and nudge payloads stay stable (coalescing).
108            diesel::update(dsl::identity)
109                .filter(
110                    dsl::next_key_package_rotation_ns
111                        .gt(rotate_at_ns)
112                        .or(dsl::next_key_package_rotation_ns.is_null()),
113                )
114                .set(dsl::next_key_package_rotation_ns.eq(rotate_at_ns))
115                .execute(conn)?;
116
117            Ok(())
118        })?;
119
120        Ok(())
121    }
122
123    fn queue_key_rotation_with_nudge(
124        &self,
125        rotation_task_hash: &crate::tasks::TaskDataHash,
126        rotation_seed: crate::tasks::NewTask,
127    ) -> Result<(), StorageError> {
128        use crate::schema::tasks;
129        use diesel::Connection;
130        use xmtp_proto::xmtp::mls::database::{PullInDeadline, Task as TaskProto, task::Task};
131
132        let hash = rotation_task_hash.to_vec();
133        self.raw_query(|conn| {
134            conn.transaction::<_, diesel::result::Error, _>(|conn| {
135                let rotate_at_ns = now_ns() + KEY_PACKAGE_QUEUE_INTERVAL_NS;
136                diesel::update(dsl::identity)
137                    .filter(
138                        dsl::next_key_package_rotation_ns
139                            .gt(rotate_at_ns)
140                            .or(dsl::next_key_package_rotation_ns.is_null()),
141                    )
142                    .set(dsl::next_key_package_rotation_ns.eq(rotate_at_ns))
143                    .execute(conn)?;
144
145                // Read back inside the tx: the column is stable between rotations,
146                // so repeat calls produce byte-identical pull-ins that coalesce.
147                let deadline: Option<Option<i64>> = dsl::identity
148                    .select(dsl::next_key_package_rotation_ns)
149                    .first::<Option<i64>>(conn)
150                    .optional()?;
151                // Pre-registration (no identity row): match the old zero-rows-
152                // matched no-op instead of erroring; nothing to rotate yet.
153                let Some(deadline) = deadline else {
154                    return Ok(());
155                };
156
157                // Ensure the pull-in's target exists (no-op when already seeded):
158                // a client whose startup seeding never ran must not enqueue a
159                // dropped-on-miss nudge.
160                diesel::insert_or_ignore_into(tasks::table)
161                    .values(rotation_seed)
162                    .execute(conn)?;
163
164                let pull_in = crate::tasks::NewTask::builder()
165                    .originating_message_sequence_id(0)
166                    .expires_at_ns(crate::tasks::NEVER_EXPIRES)
167                    .max_attempts(i32::MAX)
168                    .build(TaskProto {
169                        task: Some(Task::PullInDeadline(PullInDeadline {
170                            target_data_hash: hash,
171                            not_later_than_ns: deadline.unwrap_or(rotate_at_ns),
172                        })),
173                    })
174                    // All required builder fields are set above; unreachable.
175                    .map_err(|_| diesel::result::Error::RollbackTransaction)?;
176                diesel::insert_or_ignore_into(tasks::table)
177                    .values(pull_in)
178                    .execute(conn)?;
179                Ok(())
180            })
181        })?;
182        Ok(())
183    }
184
185    fn reset_key_package_rotation_queue(
186        &self,
187        rotation_interval_ns: i64,
188    ) -> Result<(), StorageError> {
189        use crate::schema::identity::dsl;
190
191        self.raw_query(|conn| {
192            diesel::update(dsl::identity)
193                .filter(
194                    dsl::next_key_package_rotation_ns
195                        .is_null()
196                        .or(dsl::next_key_package_rotation_ns.le(now_ns())),
197                )
198                .set(dsl::next_key_package_rotation_ns.eq(Some(now_ns() + rotation_interval_ns)))
199                .execute(conn)?;
200            Ok(())
201        })?;
202
203        Ok(())
204    }
205
206    fn is_identity_needs_rotation(&self) -> Result<bool, StorageError> {
207        use crate::schema::identity::dsl;
208
209        let next_rotation_opt: Option<Option<i64>> = self.raw_query(|conn| {
210            dsl::identity
211                .select(dsl::next_key_package_rotation_ns)
212                .first::<Option<i64>>(conn)
213                .optional()
214        })?;
215
216        Ok(match next_rotation_opt {
217            // No identity row (pre-registration): nothing to rotate yet.
218            None => false,
219            // NULL column on an existing row: rotation is due now.
220            Some(None) => true,
221            Some(Some(rotate_at)) => now_ns() >= rotate_at,
222        })
223    }
224
225    fn next_key_package_rotation_ns(&self) -> Result<Option<i64>, StorageError> {
226        use crate::schema::identity::dsl;
227        // Use optional() so an empty table (pre-registration) returns Ok(None).
228        let v: Option<Option<i64>> = self.raw_query(|conn| {
229            dsl::identity
230                .select(dsl::next_key_package_rotation_ns)
231                .first::<Option<i64>>(conn)
232                .optional()
233        })?;
234        Ok(v.flatten())
235    }
236}
237
238#[cfg(test)]
239pub(crate) mod tests {
240    use super::StoredIdentity;
241    use crate::{Store, XmtpTestDb};
242    use xmtp_common::rand_vec;
243
244    /// A stand-in rotation seed for exercising `queue_key_rotation_with_nudge`
245    /// (the real seed payload lives in xmtp_mls).
246    fn test_rotation_seed() -> crate::tasks::NewTask {
247        use xmtp_proto::xmtp::mls::database::{KpRotation, Task as TaskProto, task::Task};
248        crate::tasks::NewTask::builder()
249            .originating_message_sequence_id(0)
250            .expires_at_ns(crate::tasks::NEVER_EXPIRES)
251            .max_attempts(i32::MAX)
252            .next_attempt_at_ns(0)
253            .build(TaskProto {
254                task: Some(Task::KpRotation(KpRotation {})),
255            })
256            .unwrap()
257    }
258
259    #[xmtp_common::test]
260    fn queue_with_nudge_is_noop_before_registration() {
261        use crate::prelude::{QueryIdentity, QueryTasks};
262        use crate::test_utils::with_connection;
263        with_connection(|conn| {
264            // Empty identity table (pre-registration): must be a no-op like the
265            // old column-only path, not a NotFound error. The seed must NOT be
266            // inserted either — pre-registration means zero writes.
267            let hash = crate::tasks::TaskDataHash::try_from([0x11u8; 32].as_slice()).unwrap();
268            conn.queue_key_rotation_with_nudge(&hash, test_rotation_seed())
269                .unwrap();
270            assert!(
271                conn.get_tasks().unwrap().is_empty(),
272                "no pull-in without an identity row"
273            );
274        })
275    }
276
277    #[xmtp_common::test]
278    fn queue_with_nudge_selfheals_missing_seed() {
279        use crate::prelude::{QueryIdentity, QueryTasks};
280        use crate::test_utils::with_connection;
281        with_connection(|conn| {
282            StoredIdentity::new("".to_string(), rand_vec::<24>(), rand_vec::<24>())
283                .store(conn)
284                .unwrap();
285            let seed = test_rotation_seed();
286            let hash = crate::tasks::TaskDataHash::try_from(seed.data_hash.as_slice()).unwrap();
287            conn.queue_key_rotation_with_nudge(&hash, seed).unwrap();
288            let tasks = conn.get_tasks().unwrap();
289            assert!(
290                tasks.iter().any(|t| t.data_hash == hash.as_ref()),
291                "nudge must insert the missing rotation seed (pull-in target)"
292            );
293            assert_eq!(tasks.len(), 2, "seed + pull-in");
294        })
295    }
296
297    #[xmtp_common::test]
298    fn queue_initializes_null_rotation_column() {
299        use crate::prelude::QueryIdentity;
300        use crate::test_utils::with_connection;
301        use xmtp_configuration::KEY_PACKAGE_QUEUE_INTERVAL_NS;
302        with_connection(|conn| {
303            StoredIdentity::new("".to_string(), rand_vec::<24>(), rand_vec::<24>())
304                .store(conn)
305                .unwrap();
306
307            // Migrated DBs have NULL here; queueing must initialize it (5s
308            // debounce) rather than skip the row.
309            conn.queue_key_package_rotation().unwrap();
310            let v = conn
311                .next_key_package_rotation_ns()
312                .unwrap()
313                .expect("NULL column must be initialized");
314            let now = xmtp_common::time::now_ns();
315            assert!(v > now && v <= now + KEY_PACKAGE_QUEUE_INTERVAL_NS);
316
317            // Lower-only: a later queue call never raises the deadline.
318            conn.queue_key_package_rotation().unwrap();
319            assert_eq!(conn.next_key_package_rotation_ns().unwrap().unwrap(), v);
320        })
321    }
322
323    #[xmtp_common::test]
324    async fn can_only_store_one_identity() {
325        let store = crate::TestDb::create_ephemeral_store().await;
326        let conn = &store.conn();
327
328        StoredIdentity::new("".to_string(), rand_vec::<24>(), rand_vec::<24>())
329            .store(conn)
330            .unwrap();
331
332        let duplicate_insertion =
333            StoredIdentity::new("".to_string(), rand_vec::<24>(), rand_vec::<24>()).store(conn);
334        assert!(duplicate_insertion.is_err());
335    }
336}
xmtp_db/encrypted_store/identity.rs

xmtp_db/encrypted_store/
identity.rs
