xmtp_mls_identity/
key_package_maintenance.rs

1//! Key-package maintenance as TaskRunner consumers: payload/seed helpers and the
2//! rotate/sweep work the `KpRotation`/`KpDeletion` dispatch arms call into.
3//! Recurrence + nudging come from the generic layer (TaskOutcome, PullInDeadline).
4
5use crate::identity_updates::IdentityContext;
6use thiserror::Error;
7use tls_codec::Serialize;
8use xmtp_configuration::CREATE_PQ_KEY_PACKAGE_EXTENSION;
9use xmtp_db::StorageError;
10use xmtp_db::TransactionOutcome::Continue;
11use xmtp_db::prelude::*;
12use xmtp_db::sql_key_store::{KEY_PACKAGE_REFERENCES, KEY_PACKAGE_WRAPPER_PRIVATE_KEY};
13use xmtp_db::tasks::{NEVER_EXPIRES, NewTask, TaskDataHash, data_hash_for};
14use xmtp_mls_core::identity::IdentityError;
15use xmtp_mls_core::state_tx::state_write;
16use xmtp_mls_core::worker::NeedsDbReconnect;
17use xmtp_mls_core::worker::tasks::enqueue_pull_in;
18use xmtp_proto::xmtp::mls::database::{
19    KpDeletion, KpRotation, Task as TaskProto, task::Task as TaskKind,
20};
21
22#[derive(Debug, Error)]
23pub enum KeyPackageMaintenanceError {
24    #[error("generic storage error: {0}")]
25    Storage(#[from] StorageError),
26    #[error("generic identity error: {0}")]
27    Identity(#[from] IdentityError),
28    #[error("metadata error: {0}")]
29    Metadata(StorageError),
30    #[error("failed to fetch expired key packages: {0}")]
31    Fetch(StorageError),
32    #[error("failed to delete key package: {0}")]
33    DeleteKeyPackage(IdentityError),
34    #[error("deletion error: {0}")]
35    Deletion(StorageError),
36    #[error("rotation error: {0}")]
37    Rotation(IdentityError),
38}
39
40impl xmtp_common::RetryableError for KeyPackageMaintenanceError {
41    #[deny(clippy::wildcard_enum_match_arm)]
42    fn is_retryable(&self) -> bool {
43        match self {
44            Self::Identity(e) | Self::DeleteKeyPackage(e) | Self::Rotation(e) => e.is_retryable(),
45            Self::Storage(e) | Self::Metadata(e) | Self::Fetch(e) | Self::Deletion(e) => {
46                e.is_retryable()
47            }
48        }
49    }
50
51    /// The delay of a rate-limit rejection that failed this maintenance step.
52    fn rate_limit_delay(&self) -> Option<xmtp_common::time::Duration> {
53        match self {
54            Self::Identity(e) | Self::DeleteKeyPackage(e) | Self::Rotation(e) => {
55                e.rate_limit_delay()
56            }
57            Self::Storage(_) | Self::Metadata(_) | Self::Fetch(_) | Self::Deletion(_) => None,
58        }
59    }
60}
61
62impl NeedsDbReconnect for KeyPackageMaintenanceError {
63    #[deny(clippy::wildcard_enum_match_arm)]
64    fn needs_db_reconnect(&self) -> bool {
65        match self {
66            Self::Storage(s) => s.db_needs_connection(),
67            Self::Identity(s) => s.needs_db_reconnect(),
68            Self::Metadata(s) => s.db_needs_connection(),
69            Self::Fetch(s) => s.db_needs_connection(),
70            Self::DeleteKeyPackage(s) => s.needs_db_reconnect(),
71            Self::Deletion(s) => s.db_needs_connection(),
72            Self::Rotation(s) => s.needs_db_reconnect(),
73        }
74    }
75}
76
77/// The key package rotation task.
78pub fn kp_rotation_proto() -> TaskProto {
79    TaskProto {
80        task: Some(TaskKind::KpRotation(KpRotation {})),
81    }
82}
83
84/// The key package deletion task.
85pub fn kp_deletion_proto() -> TaskProto {
86    TaskProto {
87        task: Some(TaskKind::KpDeletion(KpDeletion {})),
88    }
89}
90
91/// The data hash of [`kp_rotation_proto`].
92pub fn kp_rotation_hash() -> TaskDataHash {
93    data_hash_for(&kp_rotation_proto())
94}
95
96/// The data hash of [`kp_deletion_proto`].
97pub fn kp_deletion_hash() -> TaskDataHash {
98    data_hash_for(&kp_deletion_proto())
99}
100
101/// Never-expire recurring seed: the reaper's
102/// `expires_at_ns < now || attempts >= max_attempts` check can never fire.
103pub fn kp_seed(proto: TaskProto, now: i64) -> Result<NewTask, StorageError> {
104    NewTask::builder()
105        .originating_message_sequence_id(0)
106        .expires_at_ns(NEVER_EXPIRES)
107        .max_attempts(i32::MAX)
108        .next_attempt_at_ns(now)
109        .build(proto)
110}
111
112/// Upload a fresh key package when rotation is due and return whether it occurred.
113/// A confirmed receipt starts retirement of keys published earlier by the backend.
114pub async fn rotate_if_needed<Context: IdentityContext>(
115    context: &Context,
116) -> Result<bool, KeyPackageMaintenanceError> {
117    if !context
118        .db()
119        .is_identity_needs_rotation()
120        .map_err(KeyPackageMaintenanceError::Metadata)?
121    {
122        return Ok(false);
123    }
124    crate::identity_lifecycle::rotate_and_upload_key_package(
125        context.identity(),
126        context.api(),
127        context.mls_storage(),
128        CREATE_PQ_KEY_PACKAGE_EXTENSION,
129    )
130    .await
131    .map_err(KeyPackageMaintenanceError::Rotation)?;
132    Ok(true)
133}
134
135/// Delete one key package's local material (keystore entry + PQ references).
136#[cfg(any(test, feature = "test-utils"))]
137pub fn delete_key_package<Context: IdentityContext>(
138    context: &Context,
139    hash_ref: Vec<u8>,
140    pq_pub_key: Option<Vec<u8>>,
141) -> Result<(), IdentityError> {
142    state_write(context.mls_storage(), |tx| {
143        delete_key_package_material(&tx.storage(), &hash_ref, pq_pub_key.as_deref())?;
144        Ok::<_, IdentityError>(Continue(()))
145    })?;
146    Ok(())
147}
148
149/// Delete all parts of one key package under the current database writer.
150fn delete_key_package_material(
151    key_store: &impl XmtpMlsStorageProvider,
152    hash_ref: &[u8],
153    pq_pub_key: Option<&[u8]>,
154) -> Result<(), IdentityError> {
155    let openmls_hash_ref = xmtp_mls_core::identity::deserialize_key_package_hash_ref(hash_ref)?;
156    let bundle: Option<openmls::prelude::KeyPackageBundle> =
157        key_store.key_package(&openmls_hash_ref)?;
158    if let Some(bundle) = bundle {
159        let init_key = bundle
160            .key_package()
161            .hpke_init_key()
162            .tls_serialize_detached()?;
163        key_store.delete(KEY_PACKAGE_REFERENCES, &init_key)?;
164    }
165    key_store.delete_key_package(&openmls_hash_ref)?;
166
167    if let Some(pq_pub_key) = pq_pub_key {
168        key_store.delete(
169            KEY_PACKAGE_REFERENCES,
170            xmtp_mls_core::identity::pq_key_package_references_key(pq_pub_key)?.as_slice(),
171        )?;
172        key_store.delete(KEY_PACKAGE_WRAPPER_PRIVATE_KEY, hash_ref)?;
173    }
174
175    Ok(())
176}
177
178/// Delete a refused key package's material and history in one transaction.
179// implements: JOIN-020
180pub fn discard_key_package(
181    mls_storage: &impl XmtpMlsStorageProvider,
182    history_id: i32,
183) -> Result<(), IdentityError> {
184    state_write(mls_storage, |tx| {
185        let storage = tx.storage();
186        let history = storage
187            .db()
188            .find_key_package_history_entries_before_id(i32::MAX)?;
189        for kp in history.iter().filter(|kp| kp.id == history_id) {
190            let pq_pub_key = kp.post_quantum_public_key.as_deref();
191            delete_key_package_material(&storage, &kp.key_package_hash_ref, pq_pub_key)?;
192            storage.db().delete_key_package_entry_with_id(kp.id)?;
193        }
194        Ok::<_, IdentityError>(Continue(()))
195    })?;
196    Ok(())
197}
198
199/// Keep expired keys while any durable Welcome remains pending.
200/// Otherwise, delete each expired key's material and history in one transaction.
201pub fn sweep_expired<Context: IdentityContext>(
202    context: &Context,
203) -> Result<(), KeyPackageMaintenanceError> {
204    state_write(context.mls_storage(), |tx| {
205        let storage = tx.storage();
206        let conn = storage.db();
207        // The durable Welcome queue owns the key-retention obligation.
208        if conn.has_pending_welcomes()? {
209            return Ok::<_, IdentityError>(Continue(()));
210        }
211        let expired = conn.get_expired_key_packages()?;
212        for kp in &expired {
213            delete_key_package_material(
214                &storage,
215                &kp.key_package_hash_ref,
216                kp.post_quantum_public_key.as_deref(),
217            )?;
218            conn.delete_key_package_entry_with_id(kp.id)?;
219        }
220        Ok(Continue(()))
221    })?;
222    Ok(())
223}
224
225/// Queue rotation within five seconds and persist its task in one transaction.
226/// Create the recurring task before its deadline update, then emit after commit.
227pub fn queue_key_rotation<Context: IdentityContext>(context: &Context) -> Result<(), StorageError> {
228    xmtp_mls_core::state_tx::state_write_with_events(
229        context.mls_storage(),
230        context.events(),
231        |tx, events| {
232            queue_key_rotation_in(&tx.storage())?;
233            xmtp_events::EventWriter::emit(
234                events,
235                None,
236                Some(xmtp_mls_core::events::InternalEvent::TaskScheduled),
237            );
238            Ok::<_, StorageError>(Continue(()))
239        },
240    )?;
241    Ok(())
242}
243
244/// Queue rotation on the caller's writer so a Welcome install can commit with it.
245/// The caller emits the task fact in its event buffer.
246pub fn queue_key_rotation_in(storage: &impl XmtpMlsStorageProvider) -> Result<(), StorageError> {
247    let now = xmtp_common::time::now_ns();
248    storage
249        .db()
250        .queue_key_rotation_with_nudge(&kp_rotation_hash(), kp_seed(kp_rotation_proto(), now)?)
251}
252
253/// After anything marks superseded KPs for deletion: ensure the KpDeletion
254/// singleton exists (pull-in against a missing target is a no-op), then pull
255/// it in to the earliest pending delete_at. No-op when nothing is marked, so
256/// it is safe (and idempotent) to call on every dispatch.
257pub fn nudge_deletion<Context: IdentityContext>(context: &Context) -> Result<(), StorageError> {
258    let db = context.db();
259    let now = xmtp_common::time::now_ns();
260    if let Some(at) = db.min_key_package_delete_at_ns()? {
261        db.create_or_ignore_task(kp_seed(kp_deletion_proto(), now)?)?;
262        enqueue_pull_in(context, kp_deletion_hash(), at, NEVER_EXPIRES)?;
263    }
264    Ok(())
265}
266
267/// Idempotent startup seeding + reconcile: pull-ins only LOWER task deadlines to
268/// the live DB columns, repairing rows stranded by a crash mid-nudge.
269pub fn seed_and_reconcile_kp_tasks<Context: IdentityContext>(
270    context: &Context,
271) -> Result<(), StorageError> {
272    let db = context.db();
273    let now = xmtp_common::time::now_ns();
274    db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
275    db.create_or_ignore_task(kp_seed(kp_deletion_proto(), now)?)?;
276    // None = pre-registration (no identity row): the seed row already fires at
277    // startup; a pull-in to `now` would be redundant noise.
278    if let Some(rot) = db.next_key_package_rotation_ns()? {
279        enqueue_pull_in(context, kp_rotation_hash(), rot, NEVER_EXPIRES)?;
280    }
281    if let Some(del) = db.min_key_package_delete_at_ns()? {
282        enqueue_pull_in(context, kp_deletion_hash(), del, NEVER_EXPIRES)?;
283    }
284    Ok(())
285}
286
287// These tests use the native connection pool and its error variants.
288#[cfg(all(test, not(target_arch = "wasm32")))]
289mod tests {
290    use super::*;
291    use crate::identity_lifecycle::test_support::IdentityTestPeer;
292    use openmls_traits::storage::StorageProvider;
293    use prost::Message;
294    use xmtp_db::ConnectionExt;
295    use xmtp_mls_core::context::CoreContext;
296    use xmtp_proto::xmtp::mls::database::Task as TaskProtoDecode;
297
298    fn row_by_hash(db: &impl QueryTasks, hash: impl AsRef<[u8]>) -> Option<xmtp_db::tasks::Task> {
299        db.get_tasks()
300            .expect("get_tasks should not fail")
301            .into_iter()
302            .find(|t| t.data_hash == hash.as_ref())
303    }
304
305    #[xmtp_common::test(unwrap_try = true)]
306    async fn pending_welcome_preserves_expired_keys_until_completion() {
307        use xmtp_db::diesel::prelude::*;
308        use xmtp_db::incoming_envelope::{
309            IncomingLimits, NetworkEntityKind, NewIncomingEnvelope, PendingBudget, StreamTopic,
310        };
311        use xmtp_db::schema::key_package_history::dsl;
312        use xmtp_proto::types::Cursor;
313
314        let alix = IdentityTestPeer::new().await;
315        let db = alix.context.db();
316        let history = db
317            .find_key_package_history_entries_before_id(i32::MAX)?
318            .pop()?;
319        let hash = xmtp_mls_core::identity::deserialize_key_package_hash_ref(
320            &history.key_package_hash_ref,
321        )?;
322        db.raw_query(|conn| {
323            diesel::update(dsl::key_package_history.filter(dsl::id.eq(history.id)))
324                .set(dsl::delete_at_ns.eq(0))
325                .execute(conn)
326        })?;
327        let topic = StreamTopic {
328            entity_id: alix.context.installation_id().to_vec(),
329            kind: NetworkEntityKind::Welcome,
330        };
331        let limit = PendingBudget { rows: 1, bytes: 1 };
332        db.admit_ordered_batch(
333            &topic,
334            Cursor(0),
335            &[NewIncomingEnvelope {
336                sequence_id: Cursor(1),
337                envelope: vec![1],
338            }],
339            IncomingLimits {
340                batch: limit,
341                topic: limit,
342                kind: limit,
343            },
344        )?;
345
346        sweep_expired(&alix.context)?;
347        let retained: Option<openmls::prelude::KeyPackageBundle> =
348            alix.context.mls_storage().key_package(&hash)?;
349        assert!(retained.is_some());
350
351        db.complete_pending_envelope(&topic, Cursor(1))?;
352        sweep_expired(&alix.context)?;
353        let removed: Option<openmls::prelude::KeyPackageBundle> =
354            alix.context.mls_storage().key_package(&hash)?;
355        assert!(removed.is_none());
356        assert!(
357            db.find_key_package_history_entry_by_hash_ref(history.key_package_hash_ref)
358                .is_err()
359        );
360    }
361
362    /// The welcome nudge must self-heal a missing rotation seed (e.g. startup
363    /// seeding never ran) instead of enqueuing a dropped-on-miss pull-in.
364    #[xmtp_common::test(unwrap_try = true)]
365    async fn welcome_nudge_selfheals_missing_rotation_seed() {
366        let alix = IdentityTestPeer::new().await; // no TaskRunner -> no seeds
367        let db = alix.context.db();
368        assert!(row_by_hash(&db, kp_rotation_hash()).is_none());
369
370        queue_key_rotation(&alix.context)?;
371
372        assert!(
373            row_by_hash(&db, kp_rotation_hash()).is_some(),
374            "nudge must recreate the missing KpRotation singleton"
375        );
376        let has_pull_in = db.get_tasks()?.iter().any(|t| {
377            matches!(
378                TaskProtoDecode::decode(t.data.as_slice()).ok().and_then(|p| p.task),
379                Some(TaskKind::PullInDeadline(p)) if p.target_data_hash == kp_rotation_hash().as_ref()
380            )
381        });
382        assert!(has_pull_in, "nudge must enqueue a rotation pull-in");
383    }
384
385    #[cfg(not(target_arch = "wasm32"))]
386    #[xmtp_common::test(unwrap_try = true)]
387    async fn queued_rotation_is_visible_before_its_task_fact() {
388        use xmtp_events::EventFilter;
389        use xmtp_mls_core::context::CoreContext;
390        use xmtp_mls_core::events::InternalEvent;
391
392        let alix = IdentityTestPeer::new().await;
393        let subscription = alix.context.events().subscribe(
394            EventFilter::default()
395                .with_internal(|event| matches!(event, InternalEvent::TaskScheduled)),
396            Some(10),
397        );
398        let before = alix.context.events().subscribe(
399            EventFilter::default()
400                .with_internal(|event| matches!(event, InternalEvent::TaskScheduled)),
401            Some(10),
402        );
403        let fired = std::rc::Rc::new(std::cell::Cell::new(0));
404        let hook = xmtp_mls_core::precommit_test_hook::install({
405            let fired = fired.clone();
406            move |_| {
407                fired.set(fired.get() + 1);
408                assert!(before.drain().is_empty());
409            }
410        });
411        queue_key_rotation(&alix.context)?;
412        drop(hook);
413        assert!(fired.get() > 0);
414        assert!(row_by_hash(&alix.context.db(), kp_rotation_hash()).is_some());
415        assert!(matches!(
416            subscription.next().await.unwrap().internal,
417            Some(InternalEvent::TaskScheduled)
418        ));
419    }
420}
421
422#[cfg(test)]
423mod deletion_tests;
424
425// These fixtures use native pool errors. Browser connection lifecycle tests
426// cover closed persistent handles in the WASM database module.
427#[cfg(all(test, not(target_arch = "wasm32")))]
428mod disconnect_propagation_tests {
429    //! Pins that a dropped-pool signal survives this owner's error types, so
430    //! `needs_db_reconnect()` stays `true` for the worker that surfaces them.
431    use super::KeyPackageMaintenanceError;
432    use xmtp_db::{ConnectionError, PlatformStorageError, StorageError};
433    use xmtp_mls_core::reconnect::NeedsDbReconnect;
434
435    /// A `StorageError` that signals the connection pool was dropped.
436    fn disconnect_storage() -> StorageError {
437        StorageError::Platform(PlatformStorageError::PoolNeedsConnection)
438    }
439
440    /// A `ConnectionError` that signals the connection pool was dropped.
441    fn disconnect_connection() -> ConnectionError {
442        ConnectionError::Platform(PlatformStorageError::PoolNeedsConnection)
443    }
444
445    /// A storage error that is NOT a disconnect — must never trip the contract.
446    fn benign_storage() -> StorageError {
447        StorageError::InvalidHmacLength
448    }
449
450    #[xmtp_common::test]
451    fn key_package_maintenance_error_forwards_disconnect() {
452        use xmtp_mls_core::identity::IdentityError;
453        assert!(KeyPackageMaintenanceError::Storage(disconnect_storage()).needs_db_reconnect());
454        // Per-key-package delete returns an IdentityError; a disconnect must
455        // bubble whether it arrives as a StorageError or a bare ConnectionError.
456        assert!(
457            KeyPackageMaintenanceError::Identity(IdentityError::StorageError(disconnect_storage()))
458                .needs_db_reconnect()
459        );
460        assert!(
461            KeyPackageMaintenanceError::Identity(IdentityError::Db(disconnect_connection()))
462                .needs_db_reconnect()
463        );
464        assert!(!KeyPackageMaintenanceError::Storage(benign_storage()).needs_db_reconnect());
465    }
466}
xmtp_mls_identity/key_package_maintenance.rs

xmtp_mls_identity/
key_package_maintenance.rs
