xmtp_mls/
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 prost::Message;
5use std::collections::HashMap;
6
7use xmtp_api::ApiError;
8use xmtp_common::RetryableError;
9use xmtp_db::incoming_envelope::{
10    AdmissionResult, IncomingLimits, NetworkEntityKind, NewIncomingEnvelope, StreamTopic,
11};
12use xmtp_db::{
13    Fetch, NotFound, XmtpOpenMlsProvider,
14    group::{GroupQueryArgs, StoredGroup},
15};
16use xmtp_proto::types::{
17    GroupId, IncomingBatchLimits, InstallationId, OrderedEnvelopeBatch, Topic, TopicCursor,
18    TopicKind,
19};
20
21use crate::{context::XmtpSharedContext, groups::MlsGroup};
22use xmtp_id::key_package::{KeyPackageVerificationError, VerifiedKeyPackageV2};
23
24use thiserror::Error;
25use xmtp_db::prelude::*;
26
27#[derive(Error, Debug)]
28pub enum MlsStoreError {
29    #[error(transparent)]
30    Storage(#[from] xmtp_db::StorageError),
31    #[error(transparent)]
32    Api(#[from] ApiError),
33    #[error(transparent)]
34    Connection(#[from] xmtp_db::ConnectionError),
35    #[error(transparent)]
36    NotFound(#[from] NotFound),
37}
38
39impl RetryableError for MlsStoreError {
40    fn is_retryable(&self) -> bool {
41        match self {
42            Self::Storage(e) => e.is_retryable(),
43            Self::Api(e) => e.is_retryable(),
44            Self::Connection(e) => e.is_retryable(),
45            Self::NotFound(e) => e.is_retryable(),
46        }
47    }
48}
49
50impl crate::worker::NeedsDbReconnect for MlsStoreError {
51    /// Forwards a dropped-pool signal from the storage/connection variants so a
52    /// worker loading groups can stop on disconnect. `Api`/`NotFound` return `false`.
53    fn needs_db_reconnect(&self) -> bool {
54        match self {
55            Self::Storage(s) => s.db_needs_connection(),
56            Self::Connection(c) => c.db_needs_connection(),
57            Self::Api(_) | Self::NotFound(_) => false,
58        }
59    }
60}
61
62#[derive(Clone)]
63pub struct MlsStore<Context> {
64    context: Context,
65}
66
67/// Durable admission results from one bounded unary receipt page.
68#[derive(Debug)]
69pub struct ReceivedPage {
70    /// Each topic's committed receipt result; processing runs separately.
71    pub admissions: Vec<(Topic, AdmissionResult)>,
72    /// The backend reports another page after this response.
73    pub has_more: bool,
74}
75
76fn stream_topic(topic: &Topic) -> Result<StreamTopic, ApiError> {
77    let kind = match topic.kind() {
78        TopicKind::GroupMessagesV1 => NetworkEntityKind::Group,
79        TopicKind::WelcomeMessagesV1 => NetworkEntityKind::Welcome,
80        TopicKind::IdentityUpdatesV1 => NetworkEntityKind::Identity,
81        _ => return Err(ApiError::InvalidRequest("incoming topic kind")),
82    };
83    Ok(StreamTopic {
84        entity_id: topic.identifier().to_vec(),
85        kind,
86    })
87}
88
89impl<Context> MlsStore<Context> {
90    pub fn new(context: Context) -> Self {
91        Self { context }
92    }
93}
94
95impl<Context> MlsStore<Context>
96where
97    Context: XmtpSharedContext,
98{
99    /// Read durable receipt F. Transport delivery alone does not advance it.
100    pub fn received_cursors(&self, topics: &[Topic]) -> Result<TopicCursor, MlsStoreError> {
101        let conn = self.context.db();
102        topics
103            .iter()
104            .map(|topic| {
105                Ok((
106                    topic.clone(),
107                    conn.topic_progress(&stream_topic(topic)?)?.received,
108                ))
109            })
110            .collect()
111    }
112
113    /// Store a validated raw batch and its receipt position in one transaction.
114    /// New Welcome rows queue key rotation in that same writer, before decoding.
115    pub fn admit_incoming_batch(
116        &self,
117        batch: &OrderedEnvelopeBatch,
118        limits: IncomingLimits,
119    ) -> Result<AdmissionResult, MlsStoreError> {
120        let topic = stream_topic(&batch.topic)?;
121        let bytes = batch.envelopes.iter().try_fold(0u64, |bytes, envelope| {
122            bytes.checked_add(envelope.encoded_len() as u64)
123        });
124        if batch.envelopes.len() as u64 > limits.batch.rows
125            || bytes.is_none_or(|bytes| bytes > limits.batch.bytes)
126        {
127            return Err(
128                ApiError::Envelope(xmtp_api_backend::envelope::EnvelopeError::Capacity).into(),
129            );
130        }
131        let mut previous = batch.after;
132        let mut rows = Vec::with_capacity(batch.envelopes.len());
133        for envelope in &batch.envelopes {
134            let meta = envelope
135                .meta
136                .as_ref()
137                .ok_or(ApiError::InvalidResponse("incoming metadata"))?;
138            let (envelope_topic, sequence_id, _) =
139                xmtp_api_backend::envelope::metadata(meta, batch.topic.kind())
140                    .map_err(ApiError::from)?;
141            if envelope_topic != batch.topic || sequence_id <= previous {
142                return Err(ApiError::InvalidResponse("incoming batch order").into());
143            }
144            previous = sequence_id;
145            rows.push(NewIncomingEnvelope {
146                sequence_id,
147                envelope: envelope.encode_to_vec(),
148            });
149        }
150        let admitted = crate::state_tx::state_write_with_events(
151            self.context.mls_storage(),
152            self.context.events(),
153            |tx, events| {
154                let storage = tx.storage();
155                let admitted =
156                    storage
157                        .db()
158                        .admit_ordered_batch(&topic, batch.after, &rows, limits)?;
159                if topic.kind == NetworkEntityKind::Welcome && admitted.inserted > 0 {
160                    crate::worker::key_package_maintenance::queue_key_rotation_in(&storage)?;
161                    xmtp_events::EventWriter::emit(
162                        events,
163                        None,
164                        Some(crate::subscriptions::internal::InternalEvent::TaskScheduled),
165                    );
166                }
167                Ok::<_, xmtp_db::StorageError>(xmtp_db::TransactionOutcome::Continue(admitted))
168            },
169        )?
170        .into_continued();
171        Ok(admitted)
172    }
173
174    /// Query from durable receipt positions and commit one bounded page.
175    pub async fn receive_topics_once(
176        &self,
177        topics: &[Topic],
178        limits: IncomingLimits,
179    ) -> Result<ReceivedPage, MlsStoreError> {
180        let cursors = self.received_cursors(topics)?;
181        self.receive_topics_once_from(cursors, limits).await
182    }
183
184    /// Use the captured durable cursors so the caller can identify a rejected
185    /// Query and avoid sending the same request again.
186    pub(crate) async fn receive_topics_once_from(
187        &self,
188        cursors: TopicCursor,
189        limits: IncomingLimits,
190    ) -> Result<ReceivedPage, MlsStoreError> {
191        let rows = limits
192            .batch
193            .rows
194            .min(limits.topic.rows)
195            .min(limits.kind.rows)
196            .min(u64::from(
197                self.context.incoming_runtime().policy().max_fetched_rows,
198            ));
199        let bytes = limits
200            .batch
201            .bytes
202            .min(limits.topic.bytes)
203            .min(limits.kind.bytes)
204            .min(self.context.incoming_runtime().policy().max_fetched_bytes);
205        let page = self
206            .context
207            .api()
208            .query_ordered_page(
209                cursors,
210                rows.min(self.context.api().limits().max_query_limit as u64) as u32,
211                IncomingBatchLimits {
212                    max_rows: usize::try_from(rows).unwrap_or(usize::MAX),
213                    max_bytes: usize::try_from(bytes).unwrap_or(usize::MAX),
214                },
215            )
216            .await?;
217        let mut admissions = Vec::with_capacity(page.batches.len());
218        for batch in page.batches {
219            let admitted = self.admit_incoming_batch(&batch, limits)?;
220            admissions.push((batch.topic, admitted));
221        }
222        Ok(ReceivedPage {
223            admissions,
224            has_more: page.has_more,
225        })
226    }
227
228    /// Fetches the current key package from the network for each of the `installation_id`s specified
229    #[tracing::instrument(level = "trace", skip_all)]
230    pub async fn get_key_packages_for_installation_ids(
231        &self,
232        installation_ids: Vec<Vec<u8>>,
233    ) -> Result<
234        HashMap<Vec<u8>, Result<VerifiedKeyPackageV2, KeyPackageVerificationError>>,
235        MlsStoreError,
236    > {
237        let installation_ids = installation_ids
238            .into_iter()
239            .map(InstallationId::try_from)
240            .collect::<Result<Vec<_>, _>>()
241            .map_err(ApiError::from)?;
242        let key_package_results = self
243            .context
244            .api()
245            .fetch_key_packages(&installation_ids)
246            .await?;
247
248        let crypto_provider = XmtpOpenMlsProvider::<()>::new_crypto();
249
250        let results: HashMap<Vec<u8>, Result<VerifiedKeyPackageV2, KeyPackageVerificationError>> =
251            key_package_results
252                .iter()
253                .filter_map(|(id, package)| {
254                    let package = package.as_ref()?;
255                    Some((
256                        id.to_vec(),
257                        VerifiedKeyPackageV2::from_bytes(
258                            &crypto_provider,
259                            &package.key_package_tls_serialized,
260                        ),
261                    ))
262                })
263                .collect();
264
265        Ok(results)
266    }
267
268    /// Query for groups with optional filters
269    ///
270    /// Filters:
271    /// - allowed_states: only return groups with the given membership states
272    /// - created_after_ns: only return groups created after the given timestamp (in nanoseconds)
273    /// - created_before_ns: only return groups created before the given timestamp (in nanoseconds)
274    /// - limit: only return the first `limit` groups
275    pub fn find_groups(
276        &self,
277        args: GroupQueryArgs,
278    ) -> Result<Vec<MlsGroup<Context>>, MlsStoreError> {
279        Ok(self
280            .context
281            .db()
282            .find_groups(args)?
283            .into_iter()
284            .map(|stored_group| {
285                MlsGroup::new(
286                    self.context.clone(),
287                    stored_group.id,
288                    stored_group.dm_id,
289                    stored_group.conversation_type,
290                    stored_group.created_at_ns,
291                )
292            })
293            .collect())
294    }
295
296    /// Look up a group by its ID
297    ///
298    /// Returns a [`MlsGroup`] if the group exists, or an error if it does not
299    ///
300    pub fn group(&self, group_id: &GroupId) -> Result<MlsGroup<Context>, MlsStoreError> {
301        let conn = self.context.db();
302        let stored_group: Option<StoredGroup> = conn.fetch(group_id)?;
303        stored_group
304            .map(|g| {
305                MlsGroup::new(
306                    self.context.clone(),
307                    g.id,
308                    g.dm_id,
309                    g.conversation_type,
310                    g.created_at_ns,
311                )
312            })
313            .ok_or(NotFound::GroupById(*group_id))
314            .map_err(Into::into)
315    }
316}
xmtp_mls/mls_store.rs

xmtp_mls/
mls_store.rs
