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(self.context.mls_storage(), |tx| {
151            let storage = tx.storage();
152            let admitted = storage
153                .db()
154                .admit_ordered_batch(&topic, batch.after, &rows, limits)?;
155            if topic.kind == NetworkEntityKind::Welcome && admitted.inserted > 0 {
156                crate::worker::key_package_maintenance::queue_key_rotation_in(&storage)?;
157            }
158            Ok::<_, xmtp_db::StorageError>(xmtp_db::TransactionOutcome::Continue(admitted))
159        })?
160        .into_continued();
161        if topic.kind == NetworkEntityKind::Welcome && admitted.inserted > 0 {
162            self.context.task_channels().wake();
163        }
164        Ok(admitted)
165    }
166
167    /// Query from durable receipt positions and commit one bounded page.
168    pub async fn receive_topics_once(
169        &self,
170        topics: &[Topic],
171        limits: IncomingLimits,
172    ) -> Result<ReceivedPage, MlsStoreError> {
173        let cursors = self.received_cursors(topics)?;
174        let rows = limits
175            .batch
176            .rows
177            .min(limits.topic.rows)
178            .min(limits.kind.rows)
179            .min(u64::from(
180                self.context.incoming_runtime().policy().max_fetched_rows,
181            ));
182        let bytes = limits
183            .batch
184            .bytes
185            .min(limits.topic.bytes)
186            .min(limits.kind.bytes)
187            .min(self.context.incoming_runtime().policy().max_fetched_bytes);
188        let page = self
189            .context
190            .api()
191            .query_ordered_page(
192                cursors,
193                rows.min(self.context.api().limits().max_query_limit as u64) as u32,
194                IncomingBatchLimits {
195                    max_rows: usize::try_from(rows).unwrap_or(usize::MAX),
196                    max_bytes: usize::try_from(bytes).unwrap_or(usize::MAX),
197                },
198            )
199            .await?;
200        let mut admissions = Vec::with_capacity(page.batches.len());
201        for batch in page.batches {
202            let admitted = self.admit_incoming_batch(&batch, limits)?;
203            admissions.push((batch.topic, admitted));
204        }
205        Ok(ReceivedPage {
206            admissions,
207            has_more: page.has_more,
208        })
209    }
210
211    /// Fetches the current key package from the network for each of the `installation_id`s specified
212    #[tracing::instrument(level = "trace", skip_all)]
213    pub async fn get_key_packages_for_installation_ids(
214        &self,
215        installation_ids: Vec<Vec<u8>>,
216    ) -> Result<
217        HashMap<Vec<u8>, Result<VerifiedKeyPackageV2, KeyPackageVerificationError>>,
218        MlsStoreError,
219    > {
220        let installation_ids = installation_ids
221            .into_iter()
222            .map(InstallationId::try_from)
223            .collect::<Result<Vec<_>, _>>()
224            .map_err(ApiError::from)?;
225        let key_package_results = self
226            .context
227            .api()
228            .fetch_key_packages(&installation_ids)
229            .await?;
230
231        let crypto_provider = XmtpOpenMlsProvider::<()>::new_crypto();
232
233        let results: HashMap<Vec<u8>, Result<VerifiedKeyPackageV2, KeyPackageVerificationError>> =
234            key_package_results
235                .iter()
236                .filter_map(|(id, package)| {
237                    let package = package.as_ref()?;
238                    Some((
239                        id.to_vec(),
240                        VerifiedKeyPackageV2::from_bytes(
241                            &crypto_provider,
242                            &package.key_package_tls_serialized,
243                        ),
244                    ))
245                })
246                .collect();
247
248        Ok(results)
249    }
250
251    /// Query for groups with optional filters
252    ///
253    /// Filters:
254    /// - allowed_states: only return groups with the given membership states
255    /// - created_after_ns: only return groups created after the given timestamp (in nanoseconds)
256    /// - created_before_ns: only return groups created before the given timestamp (in nanoseconds)
257    /// - limit: only return the first `limit` groups
258    pub fn find_groups(
259        &self,
260        args: GroupQueryArgs,
261    ) -> Result<Vec<MlsGroup<Context>>, MlsStoreError> {
262        Ok(self
263            .context
264            .db()
265            .find_groups(args)?
266            .into_iter()
267            .map(|stored_group| {
268                MlsGroup::new(
269                    self.context.clone(),
270                    stored_group.id,
271                    stored_group.dm_id,
272                    stored_group.conversation_type,
273                    stored_group.created_at_ns,
274                )
275            })
276            .collect())
277    }
278
279    /// Look up a group by its ID
280    ///
281    /// Returns a [`MlsGroup`] if the group exists, or an error if it does not
282    ///
283    pub fn group(&self, group_id: &GroupId) -> Result<MlsGroup<Context>, MlsStoreError> {
284        let conn = self.context.db();
285        let stored_group: Option<StoredGroup> = conn.fetch(group_id)?;
286        stored_group
287            .map(|g| {
288                MlsGroup::new(
289                    self.context.clone(),
290                    g.id,
291                    g.dm_id,
292                    g.conversation_type,
293                    g.created_at_ns,
294                )
295            })
296            .ok_or(NotFound::GroupById(*group_id))
297            .map_err(Into::into)
298    }
299}
xmtp_mls/mls_store.rs

xmtp_mls/
mls_store.rs
