xmtp_mls_core/
mls_store.rs

1//! Higher level queries against the local database
2//! These queries return their mls-typed equivalents after converting
3//! from the data in DB/Api
4use std::collections::HashMap;
5
6use openmls_rust_crypto::RustCrypto;
7use xmtp_api::ApiError;
8use xmtp_common::RetryableError;
9use xmtp_db::incoming_envelope::{
10    AdmissionResult, IncomingLimits, NetworkEntityKind, NewIncomingEnvelope, StreamTopic,
11};
12use xmtp_db::{NotFound, XmtpOpenMlsProvider};
13use xmtp_proto::backend_v1 as wire;
14use xmtp_proto::types::{
15    IncomingBatchLimits, InstallationId, OrderedEnvelopeBatch, Topic, TopicCursor, TopicKind,
16};
17
18use crate::context::CoreContext;
19use xmtp_id::key_package::{KeyPackageVerificationError, VerifiedKeyPackageV2};
20
21use thiserror::Error;
22use xmtp_db::prelude::*;
23
24/// Verify the package served for `id`. A package whose leaf key is another
25/// installation's was not served for `id`, so it is left out.
26// implements: API-245
27pub fn verify_served_key_package(
28    crypto: &RustCrypto,
29    id: &InstallationId,
30    package: &wire::KeyPackage,
31) -> Option<Result<VerifiedKeyPackageV2, KeyPackageVerificationError>> {
32    let verified = VerifiedKeyPackageV2::from_bytes(crypto, &package.key_package_tls_serialized);
33    if verified
34        .as_ref()
35        .is_ok_and(|package| *id != package.installation_public_key)
36    {
37        tracing::warn!("dropping a key package whose leaf names another installation");
38        return None;
39    }
40    Some(verified)
41}
42
43#[derive(Error, Debug)]
44pub enum MlsStoreError {
45    #[error(transparent)]
46    Storage(#[from] xmtp_db::StorageError),
47    #[error(transparent)]
48    Api(#[from] ApiError),
49    #[error(transparent)]
50    Connection(#[from] xmtp_db::ConnectionError),
51    #[error(transparent)]
52    NotFound(#[from] NotFound),
53}
54
55impl RetryableError for MlsStoreError {
56    fn is_retryable(&self) -> bool {
57        match self {
58            Self::Storage(e) => e.is_retryable(),
59            Self::Api(e) => e.is_retryable(),
60            Self::Connection(e) => e.is_retryable(),
61            Self::NotFound(e) => e.is_retryable(),
62        }
63    }
64
65    fn rate_limit_delay(&self) -> Option<xmtp_common::time::Duration> {
66        match self {
67            Self::Api(e) => e.rate_limit_delay(),
68            _ => None,
69        }
70    }
71}
72
73impl crate::worker::NeedsDbReconnect for MlsStoreError {
74    /// Forwards a dropped-pool signal from the storage/connection variants so a
75    /// worker loading groups can stop on disconnect. `Api`/`NotFound` return `false`.
76    #[deny(clippy::wildcard_enum_match_arm)]
77    fn needs_db_reconnect(&self) -> bool {
78        match self {
79            Self::Storage(s) => s.db_needs_connection(),
80            Self::Connection(c) => c.db_needs_connection(),
81            Self::Api(_) | Self::NotFound(_) => false,
82        }
83    }
84}
85
86#[derive(Clone)]
87pub struct MlsStore<Context> {
88    context: Context,
89}
90
91/// Durable admission results from one bounded unary receipt page.
92#[derive(Debug)]
93pub struct ReceivedPage {
94    /// Each topic's committed receipt result; processing runs separately.
95    pub admissions: Vec<(Topic, AdmissionResult)>,
96    /// The backend reports another page after this response.
97    pub has_more: bool,
98}
99
100fn stream_topic(topic: &Topic) -> Result<StreamTopic, ApiError> {
101    let kind = match topic.kind() {
102        TopicKind::GroupMessagesV1 => NetworkEntityKind::Group,
103        TopicKind::WelcomeMessagesV1 => NetworkEntityKind::Welcome,
104        TopicKind::IdentityUpdatesV1 => NetworkEntityKind::Identity,
105        _ => return Err(ApiError::InvalidRequest("incoming topic kind")),
106    };
107    Ok(StreamTopic {
108        entity_id: topic.identifier().to_vec(),
109        kind,
110    })
111}
112
113impl<Context> MlsStore<Context> {
114    pub fn new(context: Context) -> Self {
115        Self { context }
116    }
117}
118
119impl<Context> MlsStore<Context>
120where
121    Context: CoreContext,
122{
123    /// Read durable receipt F. Transport delivery alone does not advance it.
124    pub fn received_cursors(&self, topics: &[Topic]) -> Result<TopicCursor, MlsStoreError> {
125        let conn = self.context.db();
126        topics
127            .iter()
128            .map(|topic| {
129                Ok((
130                    topic.clone(),
131                    conn.topic_progress(&stream_topic(topic)?)?.received,
132                ))
133            })
134            .collect()
135    }
136
137    /// Store a validated raw batch and its receipt position in one transaction.
138    /// Receipt queues no key rotation; only an installed Welcome does.
139    pub fn admit_incoming_batch(
140        &self,
141        batch: &OrderedEnvelopeBatch,
142        limits: IncomingLimits,
143    ) -> Result<AdmissionResult, MlsStoreError> {
144        let topic = stream_topic(&batch.topic)?;
145        let bytes = batch.envelopes.iter().try_fold(0u64, |bytes, envelope| {
146            bytes.checked_add(envelope.raw().len() as u64)
147        });
148        if batch.envelopes.len() as u64 > limits.batch.rows
149            || bytes.is_none_or(|bytes| bytes > limits.batch.bytes)
150        {
151            return Err(
152                ApiError::Envelope(xmtp_api_backend::envelope::EnvelopeError::Capacity).into(),
153            );
154        }
155        let mut previous = batch.after;
156        let mut rows = Vec::with_capacity(batch.envelopes.len());
157        for envelope in &batch.envelopes {
158            let meta = envelope
159                .meta()
160                .ok_or(ApiError::InvalidResponse("incoming metadata"))?;
161            let (envelope_topic, sequence_id, _) =
162                xmtp_api_backend::envelope::metadata(meta, batch.topic.kind())
163                    .map_err(ApiError::from)?;
164            if envelope_topic != batch.topic || sequence_id <= previous {
165                return Err(ApiError::InvalidResponse("incoming batch order").into());
166            }
167            previous = sequence_id;
168            rows.push(NewIncomingEnvelope {
169                sequence_id,
170                envelope: envelope.raw().to_vec(),
171            });
172        }
173        let admitted = crate::state_tx::state_write(self.context.mls_storage(), |tx| {
174            let admitted =
175                tx.storage()
176                    .db()
177                    .admit_ordered_batch(&topic, batch.after, &rows, limits)?;
178            Ok::<_, xmtp_db::StorageError>(xmtp_db::TransactionOutcome::Continue(admitted))
179        })?
180        .into_continued();
181        Ok(admitted)
182    }
183
184    /// Query from durable receipt positions and commit one bounded page.
185    pub async fn receive_topics_once(
186        &self,
187        topics: &[Topic],
188        limits: IncomingLimits,
189    ) -> Result<ReceivedPage, MlsStoreError> {
190        let cursors = self.received_cursors(topics)?;
191        self.receive_topics_once_from(cursors, limits).await
192    }
193
194    /// Use the captured durable cursors so the caller can identify a rejected
195    /// Query and avoid sending the same request again.
196    pub async fn receive_topics_once_from(
197        &self,
198        cursors: TopicCursor,
199        limits: IncomingLimits,
200    ) -> Result<ReceivedPage, MlsStoreError> {
201        let rows = limits
202            .batch
203            .rows
204            .min(limits.topic.rows)
205            .min(limits.kind.rows)
206            .min(u64::from(self.context.stream_policy().max_fetched_rows));
207        let bytes = limits
208            .batch
209            .bytes
210            .min(limits.topic.bytes)
211            .min(limits.kind.bytes)
212            .min(self.context.stream_policy().max_fetched_bytes);
213        let page = self
214            .context
215            .api()
216            .query_ordered_page(
217                cursors,
218                rows.min(self.context.api().limits().max_query_limit as u64) as u32,
219                IncomingBatchLimits {
220                    max_rows: usize::try_from(rows).unwrap_or(usize::MAX),
221                    max_bytes: usize::try_from(bytes).unwrap_or(usize::MAX),
222                },
223            )
224            .await?;
225        let mut admissions = Vec::with_capacity(page.batches.len());
226        for batch in page.batches {
227            let admitted = self.admit_incoming_batch(&batch, limits)?;
228            admissions.push((batch.topic, admitted));
229        }
230        Ok(ReceivedPage {
231            admissions,
232            has_more: page.has_more,
233        })
234    }
235
236    /// Fetches the current key package from the network for each of the `installation_id`s specified
237    #[tracing::instrument(level = "trace", skip_all)]
238    pub async fn get_key_packages_for_installation_ids(
239        &self,
240        installation_ids: Vec<Vec<u8>>,
241    ) -> Result<
242        HashMap<Vec<u8>, Result<VerifiedKeyPackageV2, KeyPackageVerificationError>>,
243        MlsStoreError,
244    > {
245        let installation_ids = installation_ids
246            .into_iter()
247            .map(InstallationId::try_from)
248            .collect::<Result<Vec<_>, _>>()
249            .map_err(ApiError::from)?;
250        let key_package_results = self
251            .context
252            .api()
253            .fetch_key_packages(&installation_ids)
254            .await?;
255
256        let crypto_provider = XmtpOpenMlsProvider::<()>::new_crypto();
257
258        let results: HashMap<Vec<u8>, Result<VerifiedKeyPackageV2, KeyPackageVerificationError>> =
259            key_package_results
260                .iter()
261                .filter_map(|(id, package)| {
262                    verify_served_key_package(&crypto_provider, id, package.as_ref()?)
263                        .map(|result| (id.to_vec(), result))
264                })
265                .collect();
266
267        Ok(results)
268    }
269}
270
271// These fixtures use native pool errors. Browser connection lifecycle tests
272// cover closed persistent handles in the WASM database module.
273#[cfg(all(test, not(target_arch = "wasm32")))]
274mod disconnect_propagation_tests {
275    //! Pins that a dropped-pool signal survives this owner's error types, so
276    //! `needs_db_reconnect()` stays `true` for the worker that surfaces them.
277    use super::MlsStoreError;
278    use crate::reconnect::NeedsDbReconnect;
279    use xmtp_db::{ConnectionError, PlatformStorageError, StorageError};
280
281    /// A `StorageError` that signals the connection pool was dropped.
282    fn disconnect_storage() -> StorageError {
283        StorageError::Platform(PlatformStorageError::PoolNeedsConnection)
284    }
285
286    /// A `ConnectionError` that signals the connection pool was dropped.
287    fn disconnect_connection() -> ConnectionError {
288        ConnectionError::Platform(PlatformStorageError::PoolNeedsConnection)
289    }
290
291    /// A storage error that is NOT a disconnect — must never trip the contract.
292    fn benign_storage() -> StorageError {
293        StorageError::InvalidHmacLength
294    }
295
296    #[xmtp_common::test]
297    fn mls_store_error_forwards_disconnect() {
298        assert!(MlsStoreError::Storage(disconnect_storage()).needs_db_reconnect());
299        assert!(MlsStoreError::Connection(disconnect_connection()).needs_db_reconnect());
300        assert!(!MlsStoreError::Storage(benign_storage()).needs_db_reconnect());
301    }
302}
xmtp_mls_core/mls_store.rs

xmtp_mls_core/
mls_store.rs
