xmtp_mls/identity_updates/
dependencies.rs

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

xmtp_mls/identity_updates/
dependencies.rs
