xmtp_mls/identity_updates/
dependencies.rs

1//! Exact, verified identity proofs for synchronous MLS state changes.
2
3#[cfg(test)]
4use super::load_identity_updates;
5use super::{get_association_state_for_client, load_identity_updates_for_client};
6use crate::{client::ClientError, context::XmtpSharedContext};
7use futures::{StreamExt, stream};
8use std::{
9    collections::{HashMap, HashSet},
10    sync::{Arc, OnceLock, Weak},
11};
12use thiserror::Error;
13use xmtp_common::{
14    RetryableError, retryable,
15    time::{Duration, Instant, sleep},
16};
17use xmtp_configuration::IDENTITY_REFERENCE_RETRY_INTERVAL;
18use xmtp_db::{DbQuery, StorageError, prelude::*};
19use xmtp_id::associations::{AssociationError, AssociationState};
20
21/// A proof at one exact identity sequence. Zero is not an identity sequence.
22#[derive(Clone, Debug, Eq, Hash, PartialEq)]
23pub struct IdentityRequirement {
24    /// Inbox whose verified history supplies this proof.
25    pub inbox_id: String,
26    /// Exact historical identity update, never a request for the latest state.
27    pub sequence_id: u64,
28}
29
30/// Separate a missing cached proof from invalid or unavailable identity history.
31#[derive(Debug, Error)]
32pub enum IdentityDependencyError {
33    /// Retry outside the state transaction after obtaining this exact proof.
34    #[error("A verified identity proof is required at sequence {}", .0.sequence_id)]
35    Need(IdentityRequirement),
36    /// The reference cannot name a stored identity update. Not retryable.
37    #[error("Invalid identity sequence {0}")]
38    InvalidSequence(u64),
39    /// The reference stayed absent for the healthy-primary wait. Not retryable.
40    #[error("Identity reference is absent at sequence {}", .0.sequence_id)]
41    MissingReference(IdentityRequirement),
42    /// A proof or its retrieval failed. Invalid identity history blocks dependents.
43    #[error(transparent)]
44    Client(#[from] Box<ClientError>),
45    /// A local proof could not be read. Retryability follows the storage error.
46    #[error(transparent)]
47    Storage(#[from] StorageError),
48    /// Concurrent callers observe the same failed proof attempt.
49    #[error(transparent)]
50    Shared(Arc<IdentityDependencyError>),
51}
52
53impl From<ClientError> for IdentityDependencyError {
54    fn from(error: ClientError) -> Self {
55        Self::Client(Box::new(error))
56    }
57}
58
59impl RetryableError for IdentityDependencyError {
60    fn is_retryable(&self) -> bool {
61        match self {
62            Self::Need(_) => true,
63            Self::Client(error) => retryable!(error),
64            Self::Storage(error) => retryable!(error),
65            Self::Shared(error) => error.is_retryable(),
66            Self::InvalidSequence(_) | Self::MissingReference(_) => false,
67        }
68    }
69}
70
71type Resolution = tokio::sync::Mutex<Option<Result<(), Arc<IdentityDependencyError>>>>;
72
73impl crate::worker::NeedsDbReconnect for IdentityDependencyError {
74    fn needs_db_reconnect(&self) -> bool {
75        match self {
76            Self::Client(error) => error.db_needs_connection(),
77            Self::Storage(error) => error.db_needs_connection(),
78            Self::Shared(error) => error.needs_db_reconnect(),
79            Self::Need(_) | Self::InvalidSequence(_) | Self::MissingReference(_) => false,
80        }
81    }
82}
83
84/// Share concurrent proof attempts for one client and bound network requests.
85#[derive(Default)]
86pub struct IdentityResolutionRegistry {
87    /// Weak entries disappear after all callers release an attempt.
88    active: parking_lot::Mutex<HashMap<IdentityRequirement, Weak<Resolution>>>,
89    /// One request limit shared by single-proof and batch callers.
90    permits: OnceLock<Arc<tokio::sync::Semaphore>>,
91    /// Serializes own-inbox loads and their event emission for one client.
92    own_inbox_refresh: tokio::sync::Mutex<()>,
93}
94
95impl IdentityResolutionRegistry {
96    pub(crate) fn own_inbox_refresh(&self) -> &tokio::sync::Mutex<()> {
97        &self.own_inbox_refresh
98    }
99
100    fn permits(&self, limit: usize) -> Arc<tokio::sync::Semaphore> {
101        self.permits
102            .get_or_init(|| Arc::new(tokio::sync::Semaphore::new(limit)))
103            .clone()
104    }
105
106    /// Reuse the current exact-proof attempt without caching its failure forever.
107    fn acquire(&self, requirement: &IdentityRequirement) -> Arc<Resolution> {
108        let mut active = self.active.lock();
109        active.retain(|_, gate| gate.strong_count() != 0);
110        if let Some(gate) = active.get(requirement).and_then(Weak::upgrade) {
111            return gate;
112        }
113        let gate = Arc::new(Resolution::default());
114        active.insert(requirement.clone(), Arc::downgrade(&gate));
115        gate
116    }
117}
118
119impl IdentityRequirement {
120    fn checked_sequence(&self) -> Result<i64, IdentityDependencyError> {
121        i64::try_from(self.sequence_id)
122            .ok()
123            .filter(|sequence| *sequence > 0)
124            .ok_or(IdentityDependencyError::InvalidSequence(self.sequence_id))
125    }
126}
127
128/// Read only a verified cached snapshot. Never resolve to a newer snapshot.
129/// The caller supplies the state transaction's connection.
130// implements: IDENT-070
131pub(crate) fn require_association_state(
132    conn: &impl DbQuery,
133    requirement: &IdentityRequirement,
134) -> Result<AssociationState, IdentityDependencyError> {
135    let sequence = requirement.checked_sequence()?;
136    let state = conn
137        .read_from_cache(&requirement.inbox_id, sequence)?
138        .ok_or_else(|| IdentityDependencyError::Need(requirement.clone()))?;
139    if state.inbox_id != requirement.inbox_id {
140        return Err(StorageError::DbDeserialize.into());
141    }
142    state
143        .try_into()
144        .map_err(StorageError::from)
145        .map_err(Into::into)
146}
147
148/// Fetch and verify an exact identity proof without a state transaction.
149/// Successful primary reads must cover the full absence interval. A transport
150/// or verification error stops the interval and leaves dependent work pending.
151/// Concurrent callers share one attempt; cancellation permits a later retry.
152pub(crate) async fn resolve_identity_requirement(
153    context: &impl XmtpSharedContext,
154    requirement: &IdentityRequirement,
155) -> Result<(), IdentityDependencyError> {
156    match require_association_state(&context.db(), requirement) {
157        Ok(_) => return Ok(()),
158        Err(IdentityDependencyError::Need(_)) => {}
159        Err(error) => return Err(error),
160    }
161    let gate = context.identity_resolution_registry().acquire(requirement);
162    let mut resolution = gate.lock().await;
163    if resolution.is_none() {
164        let permits = context
165            .identity_resolution_registry()
166            .permits(context.incoming_runtime().policy().max_dependency_requests);
167        let _permit = permits
168            .acquire()
169            .await
170            .expect("identity request semaphore stays open");
171        *resolution = Some(
172            resolve_identity_requirement_with_wait(
173                context,
174                requirement,
175                context.incoming_runtime().policy().identity_reference_wait,
176            )
177            .await
178            .map_err(Arc::new),
179        );
180    }
181    match resolution.as_ref().expect("resolution is assigned") {
182        Ok(()) => Ok(()),
183        Err(error) => Err(match error.as_ref() {
184            IdentityDependencyError::MissingReference(requirement) => {
185                IdentityDependencyError::MissingReference(requirement.clone())
186            }
187            IdentityDependencyError::InvalidSequence(sequence) => {
188                IdentityDependencyError::InvalidSequence(*sequence)
189            }
190            _ => IdentityDependencyError::Shared(error.clone()),
191        }),
192    }
193}
194
195/// Declare absence only after healthy primary reads span the configured wait.
196/// Verify the available prefix first; invalid history is not an absent reference.
197async fn resolve_identity_requirement_with_wait(
198    context: &impl XmtpSharedContext,
199    requirement: &IdentityRequirement,
200    absence_wait: Duration,
201) -> Result<(), IdentityDependencyError> {
202    let sequence = requirement.checked_sequence()?;
203    let mut healthy_since = None;
204    loop {
205        let conn = context.db();
206        match require_association_state(&conn, requirement) {
207            Ok(_) => return Ok(()),
208            Err(IdentityDependencyError::Need(_)) => {}
209            Err(error) => return Err(error),
210        }
211        // Query uses the primary and returns the complete retained prefix.
212        load_identity_updates_for_client(&context, &conn, &[requirement.inbox_id.as_str()]).await?;
213        match get_association_state_for_client(
214            context,
215            &conn,
216            &requirement.inbox_id,
217            Some(sequence),
218        )
219        .await
220        {
221            Ok(_) => return Ok(()),
222            Err(ClientError::Association(AssociationError::MissingIdentityUpdate)) => {
223                // An invalid earlier update blocks this dependency. Absence
224                // alone must not bypass verification of the available prefix.
225                let last = conn
226                    .get_identity_updates(&requirement.inbox_id, None, Some(sequence))
227                    .map_err(ClientError::from)?
228                    .last()
229                    .map(|update| update.sequence_id);
230                if let Some(last) = last {
231                    get_association_state_for_client(
232                        context,
233                        &conn,
234                        &requirement.inbox_id,
235                        Some(last),
236                    )
237                    .await?;
238                }
239            }
240            Err(error) => return Err(error.into()),
241        }
242        let start = healthy_since.get_or_insert_with(Instant::now);
243        if start.elapsed() >= absence_wait {
244            return Err(IdentityDependencyError::MissingReference(
245                requirement.clone(),
246            ));
247        }
248        drop(conn);
249        sleep(IDENTITY_REFERENCE_RETRY_INTERVAL).await;
250    }
251}
252
253/// Coalesce exact requirements and keep every completed result. One failed
254/// proof does not cancel other fetches or remove their verified cache entries.
255// implements: PROC-013
256pub(crate) async fn resolve_identity_requirements(
257    context: &impl XmtpSharedContext,
258    requirements: impl IntoIterator<Item = IdentityRequirement>,
259) -> Vec<(IdentityRequirement, Result<(), IdentityDependencyError>)> {
260    let requirements: HashSet<_> = requirements.into_iter().collect();
261    stream::iter(requirements.into_iter().map(|requirement| async move {
262        let result = resolve_identity_requirement(context, &requirement).await;
263        (requirement, result)
264    }))
265    .buffer_unordered(context.incoming_runtime().policy().max_dependency_requests)
266    .collect()
267    .await
268}
269
270#[cfg(test)]
271mod tests {
272    use super::*;
273    use crate::tester;
274
275    #[cfg(not(target_arch = "wasm32"))]
276    #[xmtp_common::test(unwrap_try = true)]
277    fn identity_and_query_futures_are_send() {
278        fn assert_send<T: Send>(_: T) {}
279        let context = crate::test::mock::context();
280        let requirement = IdentityRequirement {
281            inbox_id: "inbox".into(),
282            sequence_id: 1,
283        };
284        assert_send(resolve_identity_requirement(&context, &requirement));
285        assert_send(resolve_identity_requirements(
286            &context,
287            [requirement.clone()],
288        ));
289        assert_send(context.api().newest_topic_cursors(vec![]));
290        assert_send(context.api().query_ordered_page(
291            Default::default(),
292            1,
293            xmtp_proto::types::IncomingBatchLimits {
294                max_rows: 1,
295                max_bytes: 1024,
296            },
297        ));
298        // Check the production transport type as well as the mock type.
299        let _: fn(&crate::Client<crate::MlsContext>) = |client| {
300            assert_send(async move { client.inbox_state(true).await });
301            assert_send(xmtp_common::bind_task_hub(async move {
302                client.inbox_state(true).await
303            }));
304            assert_send(async move {
305                client
306                    .wait_for_registration_visible(Default::default())
307                    .await
308            });
309            assert_send(async move { client.context.api().get_inbox_ids(vec![]).await });
310            assert_send(async move { client.context.api().newest_topic_cursors(vec![]).await });
311        };
312    }
313
314    #[xmtp_common::test(unwrap_try = true)]
315    async fn identity_requests_share_one_context_limit() {
316        let registry = IdentityResolutionRegistry::default();
317        let first = registry.permits(1);
318        let second = registry.permits(1);
319        assert!(Arc::ptr_eq(&first, &second));
320        let permit = first.acquire().await?;
321        assert!(second.try_acquire().is_err());
322        drop(permit);
323        assert!(second.try_acquire().is_ok());
324    }
325
326    #[xmtp_common::test(unwrap_try = true)]
327    async fn concurrent_requirements_share_one_attempt_but_later_calls_retry() {
328        let registry = IdentityResolutionRegistry::default();
329        let requirement = IdentityRequirement {
330            inbox_id: "inbox".into(),
331            sequence_id: 5,
332        };
333        let first = registry.acquire(&requirement);
334        let second = registry.acquire(&requirement);
335        assert!(Arc::ptr_eq(&first, &second));
336
337        let mut resolution = first.lock().await;
338        assert!(second.try_lock().is_err());
339        *resolution = Some(Err(Arc::new(IdentityDependencyError::InvalidSequence(5))));
340        drop(resolution);
341        assert!(matches!(second.lock().await.as_ref(), Some(Err(_))));
342
343        let completed = Arc::downgrade(&first);
344        drop(first);
345        drop(second);
346        let next = registry.acquire(&requirement);
347        assert!(completed.upgrade().is_none());
348        assert!(next.lock().await.is_none());
349        assert_eq!(registry.active.lock().len(), 1);
350    }
351
352    // verifies: IDENT-070
353    #[xmtp_common::test(unwrap_try = true)]
354    async fn historical_proof_does_not_use_a_newer_cached_snapshot() {
355        tester!(alix, disable_workers);
356        tester!(bo, disable_workers);
357        let first_sequence =
358            alix.context
359                .db()
360                .get_latest_sequence_id(&[alix.inbox_id()])?[alix.inbox_id()] as u64;
361        tester!(alix2, from: alix, disable_workers);
362        load_identity_updates(alix.context.api(), &alix.context.db(), &[alix.inbox_id()]).await?;
363        let last_sequence =
364            alix.context
365                .db()
366                .get_latest_sequence_id(&[alix.inbox_id()])?[alix.inbox_id()] as u64;
367        assert!(last_sequence > first_sequence);
368        let latest = IdentityRequirement {
369            inbox_id: alix.inbox_id().to_owned(),
370            sequence_id: last_sequence,
371        };
372        resolve_identity_requirement(&bo.context, &latest).await?;
373        let historical = IdentityRequirement {
374            sequence_id: first_sequence,
375            ..latest
376        };
377        assert!(matches!(
378            require_association_state(&bo.context.db(), &historical),
379            Err(IdentityDependencyError::Need(_))
380        ));
381        resolve_identity_requirement(&bo.context, &historical).await?;
382        let state = require_association_state(&bo.context.db(), &historical)?;
383        assert!(
384            state
385                .installation_ids()
386                .contains(&alix.installation_public_key().to_vec())
387        );
388        assert!(
389            !state
390                .installation_ids()
391                .contains(&alix2.installation_public_key().to_vec())
392        );
393    }
394
395    // verifies: PROC-013
396    #[xmtp_common::test(unwrap_try = true)]
397    async fn dependency_batch_keeps_success_after_an_invalid_sibling() {
398        tester!(alix, disable_workers);
399        tester!(bo, disable_workers);
400        let sequence_id =
401            alix.context
402                .db()
403                .get_latest_sequence_id(&[alix.inbox_id()])?[alix.inbox_id()] as u64;
404        let valid = IdentityRequirement {
405            inbox_id: alix.inbox_id().to_owned(),
406            sequence_id,
407        };
408        let invalid = IdentityRequirement {
409            sequence_id: 0,
410            ..valid.clone()
411        };
412        let results =
413            resolve_identity_requirements(&bo.context, [valid.clone(), invalid, valid.clone()])
414                .await;
415        assert_eq!(results.len(), 2);
416        assert_eq!(
417            results.iter().filter(|(_, result)| result.is_ok()).count(),
418            1
419        );
420        assert!(require_association_state(&bo.context.db(), &valid).is_ok());
421        assert!(
422            results.iter().any(|(_, result)| matches!(
423                result,
424                Err(IdentityDependencyError::InvalidSequence(0))
425            ))
426        );
427    }
428
429    #[xmtp_common::test(unwrap_try = true)]
430    async fn primary_absence_is_terminal_only_after_successful_query() {
431        tester!(alix, disable_workers);
432        let requirement = IdentityRequirement {
433            inbox_id: hex::encode(xmtp_common::rand_vec::<32>()),
434            sequence_id: 1,
435        };
436        let result =
437            resolve_identity_requirement_with_wait(&alix.context, &requirement, Duration::ZERO)
438                .await;
439        assert!(
440            matches!(result, Err(IdentityDependencyError::MissingReference(found)) if found == requirement)
441        );
442    }
443}
xmtp_mls/identity_updates/dependencies.rs

xmtp_mls/identity_updates/
dependencies.rs
