xmtp_mls_identity/identity_updates/
dependencies.rs

1//! Exact, verified identity proofs for synchronous MLS state changes.
2
3use super::IdentityUpdateError;
4#[cfg(test)]
5use super::load_identity_updates;
6use super::{get_association_state_for_client, load_identity_updates_for_client};
7use crate::identity_updates::IdentityContext;
8use futures::{StreamExt, stream};
9use std::{
10    collections::{BTreeSet, HashSet},
11    sync::{Arc, OnceLock},
12};
13use thiserror::Error;
14use xmtp_common::{
15    RetryableError, retryable,
16    time::{Instant, sleep},
17};
18use xmtp_configuration::{
19    IDENTITY_REFERENCE_MAX_RETRY_INTERVAL, IDENTITY_REFERENCE_RETRY_INTERVAL,
20};
21use xmtp_db::{DbQuery, StorageError};
22use xmtp_id::associations::{AssociationError, AssociationState};
23
24/// A proof at one exact identity sequence. Zero is not an identity sequence.
25#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
26pub struct IdentityRequirement {
27    /// Inbox whose verified history supplies this proof.
28    pub inbox_id: String,
29    /// Exact historical identity update, never a request for the latest state.
30    pub sequence_id: u64,
31}
32
33/// Separate a missing cached proof from invalid or unavailable identity history.
34#[derive(Debug, Error)]
35pub enum IdentityDependencyError {
36    /// Retry outside the state transaction after obtaining these exact proofs.
37    /// The set holds every missing proof the attempt found, never zero proofs.
38    #[error("Verified identity proofs are required at {} references", .0.len())]
39    Need(BTreeSet<IdentityRequirement>),
40    /// The reference cannot name a stored identity update. Not retryable.
41    #[error("Invalid identity sequence {0}")]
42    InvalidSequence(u64),
43    /// The reference stayed absent for the healthy-primary wait. Not retryable.
44    #[error("Identity reference is absent at sequence {}", .0.sequence_id)]
45    MissingReference(IdentityRequirement),
46    /// A proof or its retrieval failed. Invalid identity history blocks dependents.
47    #[error(transparent)]
48    IdentityUpdate(#[from] Box<IdentityUpdateError>),
49    /// A local proof could not be read. Retryability follows the storage error.
50    #[error(transparent)]
51    Storage(#[from] StorageError),
52}
53
54impl From<IdentityUpdateError> for IdentityDependencyError {
55    fn from(error: IdentityUpdateError) -> Self {
56        Self::IdentityUpdate(Box::new(error))
57    }
58}
59
60impl RetryableError for IdentityDependencyError {
61    fn is_retryable(&self) -> bool {
62        match self {
63            Self::Need(_) => true,
64            Self::IdentityUpdate(error) => retryable!(error),
65            Self::Storage(error) => retryable!(error),
66            Self::InvalidSequence(_) | Self::MissingReference(_) => false,
67        }
68    }
69
70    fn rate_limit_delay(&self) -> Option<xmtp_common::time::Duration> {
71        match self {
72            Self::IdentityUpdate(error) => error.rate_limit_delay(),
73            _ => None,
74        }
75    }
76}
77
78impl xmtp_mls_core::worker::NeedsDbReconnect for IdentityDependencyError {
79    #[deny(clippy::wildcard_enum_match_arm)]
80    fn needs_db_reconnect(&self) -> bool {
81        match self {
82            Self::IdentityUpdate(error) => error.needs_db_reconnect(),
83            Self::Storage(error) => error.db_needs_connection(),
84            Self::Need(_) | Self::InvalidSequence(_) | Self::MissingReference(_) => false,
85        }
86    }
87}
88
89/// Bound identity proof requests for one client.
90#[derive(Default)]
91pub struct IdentityResolutionRegistry {
92    /// One request limit shared by every resolution round.
93    permits: OnceLock<Arc<tokio::sync::Semaphore>>,
94    /// Serializes own-inbox loads and their event emission for one client.
95    own_inbox_refresh: tokio::sync::Mutex<OwnInstallationEvents>,
96}
97
98/// How far this client has reported changes to its own installations.
99#[derive(Debug, Default)]
100pub struct OwnInstallationEvents {
101    /// The own-inbox identity update through which changes are emitted. A
102    /// built client starts at the latest one stored then. `None` while no
103    /// own-inbox update is stored; the first own-inbox load then starts at
104    /// the latest one stored without emitting for it.
105    pub(crate) emitted_through: Option<i64>,
106}
107
108impl IdentityResolutionRegistry {
109    /// A registry whose own-installation events start after `emitted_through`.
110    pub fn starting_after(emitted_through: Option<i64>) -> Self {
111        Self {
112            own_inbox_refresh: tokio::sync::Mutex::new(OwnInstallationEvents { emitted_through }),
113            ..Default::default()
114        }
115    }
116
117    pub(crate) fn own_inbox_refresh(&self) -> &tokio::sync::Mutex<OwnInstallationEvents> {
118        &self.own_inbox_refresh
119    }
120
121    fn permits(&self, limit: usize) -> Arc<tokio::sync::Semaphore> {
122        self.permits
123            .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(limit)))
124            .clone()
125    }
126}
127
128impl IdentityRequirement {
129    fn checked_sequence(&self) -> Result<i64, IdentityDependencyError> {
130        i64::try_from(self.sequence_id)
131            .ok()
132            .filter(|sequence| *sequence > 0)
133            .ok_or(IdentityDependencyError::InvalidSequence(self.sequence_id))
134    }
135}
136
137/// Read only a verified cached snapshot. Never resolve to a newer snapshot.
138/// The caller supplies the state transaction's connection.
139// implements: IDENT-070
140pub fn require_association_state(
141    conn: &impl DbQuery,
142    requirement: &IdentityRequirement,
143) -> Result<AssociationState, IdentityDependencyError> {
144    let sequence = requirement.checked_sequence()?;
145    let state = conn
146        .read_from_cache(&requirement.inbox_id, sequence)?
147        .ok_or_else(|| IdentityDependencyError::Need(BTreeSet::from([requirement.clone()])))?;
148    if state.inbox_id != requirement.inbox_id {
149        return Err(StorageError::DbDeserialize.into());
150    }
151    state
152        .try_into()
153        .map_err(StorageError::from)
154        .map_err(Into::into)
155}
156
157/// Read cached proofs and record each missing one, so that one resolution
158/// fetches every proof an attempt lacks instead of one proof per attempt.
159#[derive(Default)]
160pub struct MissingProofs(BTreeSet<IdentityRequirement>);
161
162impl MissingProofs {
163    /// Return the cached state, or record the requirement and return `None`.
164    /// Other failures do not depend on a fetch, so they return at once.
165    pub fn read(
166        &mut self,
167        conn: &impl DbQuery,
168        requirement: &IdentityRequirement,
169    ) -> Result<Option<AssociationState>, IdentityDependencyError> {
170        match require_association_state(conn, requirement) {
171            Ok(state) => Ok(Some(state)),
172            Err(IdentityDependencyError::Need(needed)) => {
173                self.0.extend(needed);
174                Ok(None)
175            }
176            Err(error) => Err(error),
177        }
178    }
179
180    /// Fail with every recorded requirement.
181    pub fn check(self) -> Result<(), IdentityDependencyError> {
182        if self.0.is_empty() {
183            Ok(())
184        } else {
185            Err(IdentityDependencyError::Need(self.0))
186        }
187    }
188}
189
190/// Resolve every requirement of one parent and return one result for it.
191/// Each round loads the logs of every still-missing inbox in one request and
192/// checks each proof once, so absent references share one window. Successful
193/// primary reads must cover the full absence interval. A transport or
194/// verification error stops the interval and leaves dependent work pending.
195/// One failed proof does not cancel other checks or remove their verified
196/// cache entries.
197// implements: PROC-013
198pub async fn resolve_identity_requirements(
199    context: &impl IdentityContext,
200    requirements: &BTreeSet<IdentityRequirement>,
201) -> Result<(), IdentityDependencyError> {
202    let policy = context.stream_policy();
203    let limit = policy.max_dependency_requests;
204    let permits = context.identity_resolution_registry().permits(limit);
205    let mut missing = requirements.clone();
206    let mut healthy_since = None;
207    let mut interval = IDENTITY_REFERENCE_RETRY_INTERVAL;
208    loop {
209        let conn = context.db();
210        let mut results = Vec::new();
211        let mut needed = Vec::new();
212        for requirement in &missing {
213            match require_association_state(&conn, requirement) {
214                Err(IdentityDependencyError::Need(_)) => needed.push(requirement.clone()),
215                result => results.push(result.map(drop)),
216            }
217        }
218        if !needed.is_empty() {
219            // The permit covers one load and its checks, never a sleep.
220            let _permit = permits
221                .acquire()
222                .await
223                .expect("identity request semaphore stays open");
224            let inbox_ids: Vec<_> = needed
225                .iter()
226                .map(|requirement| requirement.inbox_id.as_str())
227                .collect::<HashSet<_>>()
228                .into_iter()
229                .collect();
230            // Query uses the primary and returns the complete retained prefix.
231            load_identity_updates_for_client(context, &conn, &inbox_ids).await?;
232            let conn = &conn;
233            let checked: Vec<_> = stream::iter(needed.into_iter().map(|requirement| async move {
234                check_identity_requirement(context, conn, &requirement).await
235            }))
236            .buffer_unordered(limit)
237            .collect()
238            .await;
239            results.extend(checked);
240        }
241        missing = results
242            .iter()
243            .filter_map(|result| match result {
244                Err(IdentityDependencyError::MissingReference(requirement)) => {
245                    Some(requirement.clone())
246                }
247                _ => None,
248            })
249            .collect();
250        match parent_result(results) {
251            Err(IdentityDependencyError::MissingReference(requirement)) => {
252                let start = healthy_since.get_or_insert_with(Instant::now);
253                if start.elapsed() >= policy.identity_reference_wait {
254                    return Err(IdentityDependencyError::MissingReference(requirement));
255                }
256            }
257            result => return result,
258        }
259        drop(conn);
260        sleep(interval).await;
261        interval = (interval * 2).min(IDENTITY_REFERENCE_MAX_RETRY_INTERVAL);
262    }
263}
264
265/// Verify one exact proof after a primary load. Verify the available prefix
266/// first; invalid history is not an absent reference.
267async fn check_identity_requirement(
268    context: &impl IdentityContext,
269    conn: &impl DbQuery,
270    requirement: &IdentityRequirement,
271) -> Result<(), IdentityDependencyError> {
272    let sequence = requirement.checked_sequence()?;
273    match get_association_state_for_client(context, conn, &requirement.inbox_id, Some(sequence))
274        .await
275    {
276        Ok(_) => Ok(()),
277        Err(IdentityUpdateError::Association(AssociationError::MissingIdentityUpdate)) => {
278            // An invalid earlier update blocks this dependency. Absence
279            // alone must not bypass verification of the available prefix.
280            let last = conn
281                .get_identity_updates(&requirement.inbox_id, None, Some(sequence))
282                .map_err(IdentityUpdateError::from)?
283                .last()
284                .map(|update| update.sequence_id);
285            if let Some(last) = last {
286                get_association_state_for_client(context, conn, &requirement.inbox_id, Some(last))
287                    .await?;
288            }
289            Err(IdentityDependencyError::MissingReference(
290                requirement.clone(),
291            ))
292        }
293        Err(error) => Err(error.into()),
294    }
295}
296
297/// A failed proof takes precedence over an absent one, so a parent that
298/// needs unresolved identity history stays blocked, not rejected. Of the
299/// failed proofs, the rate-limited one with the longest delay is returned, so
300/// the parent's retry waits out every refused verification instead of
301/// resending one inside its own delay.
302// implements: PROC-012, RATE-041
303fn parent_result(
304    results: impl IntoIterator<Item = Result<(), IdentityDependencyError>>,
305) -> Result<(), IdentityDependencyError> {
306    let mut absent = None;
307    let mut failed: Option<IdentityDependencyError> = None;
308    for result in results {
309        match result {
310            Ok(()) => {}
311            Err(error @ IdentityDependencyError::MissingReference(_)) => absent = Some(error),
312            // `None` orders first, so any delay beats none and a longer one
313            // beats a shorter one; the first failure is kept on a tie.
314            Err(error) => {
315                if failed
316                    .as_ref()
317                    .is_none_or(|kept| error.rate_limit_delay() > kept.rate_limit_delay())
318                {
319                    failed = Some(error);
320                }
321            }
322        }
323    }
324    failed.or(absent).map_or(Ok(()), Err)
325}
326
327#[cfg(test)]
328mod tests {
329    use super::*;
330    use crate::identity_lifecycle::test_support::IdentityTestPeer;
331    use xmtp_common::time::Duration;
332    use xmtp_db::prelude::*;
333    use xmtp_mls_core::context::CoreContext;
334
335    #[xmtp_common::test(unwrap_try = true)]
336    async fn identity_requests_share_one_context_limit() {
337        let registry = IdentityResolutionRegistry::default();
338        let first = registry.permits(1);
339        let second = registry.permits(1);
340        assert!(Arc::ptr_eq(&first, &second));
341        let permit = first.acquire().await?;
342        assert!(second.try_acquire().is_err());
343        drop(permit);
344        assert!(second.try_acquire().is_ok());
345    }
346
347    // verifies: IDENT-070
348    #[xmtp_common::test(unwrap_try = true)]
349    async fn historical_proof_does_not_use_a_newer_cached_snapshot() {
350        let alix = IdentityTestPeer::new().await;
351        let bo = IdentityTestPeer::new().await;
352        let first_sequence =
353            alix.context
354                .db()
355                .get_latest_sequence_id(&[alix.inbox_id()])?[alix.inbox_id()] as u64;
356        let alix2 = alix.new_installation().await;
357        load_identity_updates(alix.context.api(), &alix.context.db(), &[alix.inbox_id()]).await?;
358        let last_sequence =
359            alix.context
360                .db()
361                .get_latest_sequence_id(&[alix.inbox_id()])?[alix.inbox_id()] as u64;
362        assert!(last_sequence > first_sequence);
363        let latest = IdentityRequirement {
364            inbox_id: alix.inbox_id().to_owned(),
365            sequence_id: last_sequence,
366        };
367        resolve_identity_requirements(&bo.context, &BTreeSet::from([latest.clone()])).await?;
368        let historical = IdentityRequirement {
369            sequence_id: first_sequence,
370            ..latest
371        };
372        assert!(matches!(
373            require_association_state(&bo.context.db(), &historical),
374            Err(IdentityDependencyError::Need(_))
375        ));
376        resolve_identity_requirements(&bo.context, &BTreeSet::from([historical.clone()])).await?;
377        let state = require_association_state(&bo.context.db(), &historical)?;
378        assert!(
379            state
380                .installation_ids()
381                .contains(&alix.installation_public_key().to_vec())
382        );
383        assert!(
384            !state
385                .installation_ids()
386                .contains(&alix2.installation_public_key().to_vec())
387        );
388    }
389
390    // verifies: PROC-013
391    #[xmtp_common::test(unwrap_try = true)]
392    async fn dependency_batch_keeps_success_after_an_invalid_sibling() {
393        let alix = IdentityTestPeer::new().await;
394        let bo = IdentityTestPeer::new().await;
395        let sequence_id =
396            alix.context
397                .db()
398                .get_latest_sequence_id(&[alix.inbox_id()])?[alix.inbox_id()] as u64;
399        let valid = IdentityRequirement {
400            inbox_id: alix.inbox_id().to_owned(),
401            sequence_id,
402        };
403        let invalid = IdentityRequirement {
404            sequence_id: 0,
405            ..valid.clone()
406        };
407        let result =
408            resolve_identity_requirements(&bo.context, &BTreeSet::from([valid.clone(), invalid]))
409                .await;
410        assert!(matches!(
411            result,
412            Err(IdentityDependencyError::InvalidSequence(0))
413        ));
414        assert!(require_association_state(&bo.context.db(), &valid).is_ok());
415    }
416
417    // verifies: PROC-012
418    #[xmtp_common::test(unwrap_try = true)]
419    fn a_failed_proof_blocks_before_an_absent_one() {
420        let absent = || {
421            Err(IdentityDependencyError::MissingReference(
422                IdentityRequirement {
423                    inbox_id: "inbox".into(),
424                    sequence_id: 1,
425                },
426            ))
427        };
428        let failed = || {
429            Err(IdentityDependencyError::Storage(
430                StorageError::DbDeserialize,
431            ))
432        };
433        assert!(parent_result([Ok(()), Ok(())]).is_ok());
434        assert!(matches!(
435            parent_result([Ok(()), absent()]),
436            Err(IdentityDependencyError::MissingReference(_))
437        ));
438        for results in [[absent(), failed()], [failed(), absent()]] {
439            assert!(matches!(
440                parent_result(results),
441                Err(IdentityDependencyError::Storage(_))
442            ));
443        }
444    }
445
446    // Absent references share one window that ends in a named absence after
447    // successful queries, and another parent's reference does not queue behind it.
448    #[xmtp_common::test(unwrap_try = true)]
449    async fn absent_references_share_one_window() {
450        let alix = IdentityTestPeer::new().await;
451        let bo = IdentityTestPeer::new().await;
452        let wait = Duration::from_secs(3);
453        let client = alix
454            .with_stream_policy(|settings| {
455                settings.identity_reference_wait = wait;
456                settings.max_dependency_requests = 1;
457            })
458            .await;
459        let absent: BTreeSet<_> = (0..16)
460            .map(|_| IdentityRequirement {
461                inbox_id: hex::encode(xmtp_common::rand_vec::<32>()),
462                sequence_id: 1,
463            })
464            .collect();
465        let latest = bo.context.db().get_latest_sequence_id(&[bo.inbox_id()])?[bo.inbox_id()];
466        let present = BTreeSet::from([IdentityRequirement {
467            inbox_id: bo.inbox_id().to_owned(),
468            sequence_id: latest as u64,
469        }]);
470        let queries = client.api_stats().query.get_count();
471        let ((absent_result, absent_elapsed), (present_result, present_elapsed)) =
472            futures::future::join(
473                async {
474                    let start = Instant::now();
475                    let result = resolve_identity_requirements(&client.context, &absent).await;
476                    (result, start.elapsed())
477                },
478                async {
479                    sleep(Duration::from_millis(400)).await;
480                    let start = Instant::now();
481                    let result = resolve_identity_requirements(&client.context, &present).await;
482                    (result, start.elapsed())
483                },
484            )
485            .await;
486        assert!(
487            matches!(absent_result, Err(IdentityDependencyError::MissingReference(found)) if absent.contains(&found))
488        );
489        // Every round reloads the absent inboxes from the primary.
490        assert!(client.api_stats().query.get_count() >= queries + 4);
491        assert!(
492            absent_elapsed >= wait && absent_elapsed < wait * 2,
493            "absent references should share one window, took {absent_elapsed:?}"
494        );
495        assert!(present_result.is_ok());
496        assert!(
497            present_elapsed < wait / 2,
498            "a present reference should not wait for a permit, took {present_elapsed:?}"
499        );
500    }
501
502    /// A requirement on an update its inbox has not published waits out the
503    /// absence window, then fails as a missing reference that is not retried
504    /// and needs no reconnect.
505    #[xmtp_common::test(unwrap_try = true)]
506    async fn dependency_resolution_waits_on_missing_identity_update() {
507        use xmtp_common::RetryableError;
508        use xmtp_mls_core::worker::NeedsDbReconnect;
509        let alix = IdentityTestPeer::new().await;
510        let bo = IdentityTestPeer::new().await;
511        let wait = Duration::from_secs(1);
512        let client = alix
513            .with_stream_policy(|settings| settings.identity_reference_wait = wait)
514            .await;
515        let latest = bo.context.db().get_latest_sequence_id(&[bo.inbox_id()])?[bo.inbox_id()];
516        let unpublished = IdentityRequirement {
517            inbox_id: bo.inbox_id().to_owned(),
518            sequence_id: latest as u64 + 1_000,
519        };
520        let start = Instant::now();
521        let error =
522            resolve_identity_requirements(&client.context, &BTreeSet::from([unpublished.clone()]))
523                .await
524                .unwrap_err();
525        assert!(
526            start.elapsed() >= wait,
527            "resolved after {:?}",
528            start.elapsed()
529        );
530        assert!(
531            matches!(&error, IdentityDependencyError::MissingReference(found) if *found == unpublished),
532            "{error:?}"
533        );
534        assert!(!error.is_retryable());
535        assert!(!error.needs_db_reconnect());
536    }
537}
xmtp_mls_identity/identity_updates/dependencies.rs

xmtp_mls_identity/identity_updates/
dependencies.rs
