Skip to main content

xmtp_mls/groups/
subscriptions.rs

1use super::MlsGroup;
2use crate::{
3    context::XmtpSharedContext,
4    subscriptions::{Result, stream_messages::StreamGroupMessages},
5};
6use futures::Stream;
7use prost::Message;
8use xmtp_proto::backend_v1::ServerEnvelope;
9
10use xmtp_db::group_message::StoredGroupMessage;
11use xmtp_proto::api_client::XmtpMlsStreams;
12
13impl<Context> MlsGroup<Context>
14where
15    Context: XmtpSharedContext + 'static,
16{
17    /// Use a push envelope only as a target for ordered receipt and processing.
18    pub async fn process_streamed_group_message(
19        &self,
20        envelope_bytes: Vec<u8>,
21    ) -> Result<Vec<StoredGroupMessage>> {
22        use xmtp_db::prelude::*;
23        let wire = ServerEnvelope::decode(envelope_bytes.as_slice())?;
24        let meta = wire
25            .meta
26            .as_ref()
27            .ok_or(xmtp_api::ApiError::InvalidResponse("group metadata"))?;
28        let (topic, cursor, _) = xmtp_api_backend::envelope::metadata(
29            meta,
30            xmtp_proto::types::TopicKind::GroupMessagesV1,
31        )?;
32        if topic != xmtp_proto::types::Topic::new_group_message(self.group_id) {
33            return Err(xmtp_api::ApiError::InvalidResponse("group topic").into());
34        }
35        crate::subscriptions::barrier::wait_through(&self.context, [(topic, cursor)].into(), None)
36            .await
37            .map_err(|error| {
38                self.context
39                    .server_configuration()
40                    .blocked_or(super::GroupError::from(error))
41            })?;
42        let message = self
43            .context
44            .db()
45            .get_group_message_by_cursor(self.group_id, cursor)?;
46        // implements: META-051
47        let current_time_ns = xmtp_common::time::now_ns();
48        Ok(message
49            .filter(|row| {
50                row.expire_at_ns
51                    .is_none_or(|deadline| deadline > current_time_ns)
52            })
53            .into_iter()
54            .collect())
55    }
56
57    #[tracing::instrument(err, skip_all, fields(operation = "stream.stream_group_messages"))]
58    pub async fn stream<'a>(
59        &'a self,
60    ) -> Result<impl Stream<Item = Result<StoredGroupMessage>> + use<'a, Context>>
61    where
62        Context::ApiClient: XmtpMlsStreams + 'a,
63    {
64        StreamGroupMessages::new(&self.context, vec![self.group_id]).await
65    }
66
67    /// create a stream that is not attached to any lifetime
68    #[tracing::instrument(
69        err,
70        skip_all,
71        fields(operation = "stream.stream_group_messages_owned")
72    )]
73    pub async fn stream_owned(
74        &self,
75    ) -> Result<impl Stream<Item = Result<StoredGroupMessage>> + 'static>
76    where
77        Context: 'static,
78        Context::ApiClient: XmtpMlsStreams + 'static,
79        Context::Db: 'static,
80    {
81        StreamGroupMessages::new_owned(self.context.clone(), vec![self.group_id]).await
82    }
83}
84
85#[cfg(test)]
86pub(crate) mod tests {
87    use crate::context::XmtpSharedContext;
88    use crate::groups::send_message_opts::SendMessageOpts;
89    use futures::StreamExt;
90    use std::sync::Arc;
91
92    use crate::builder::ClientBuilder;
93    use prost::Message as ProstMessage;
94    use std::time::Duration;
95    use xmtp_cryptography::utils::generate_local_wallet;
96    use xmtp_db::{
97        ConnectionExt,
98        group_message::{GroupMessageKind, QueryGroupMessage},
99    };
100
101    #[xmtp_common::timeout(Duration::from_secs(10))]
102    #[rstest::rstest]
103    #[xmtp_common::test(flavor = "current_thread")]
104    async fn test_subscribe_messages() {
105        let amal = ClientBuilder::new_test_client(&generate_local_wallet()).await;
106        let bola = Arc::new(ClientBuilder::new_test_client(&generate_local_wallet()).await);
107
108        let amal_group = amal.create_group(None, None).unwrap();
109        // Add bola
110        amal_group.add_members(&[bola.inbox_id()]).await.unwrap();
111
112        // Get bola's version of the same group
113        bola.sync_welcomes().await.unwrap();
114        // A background reader can store the Welcome before this call.
115        let bola_group = bola.group(&amal_group.group_id).unwrap();
116        bola_group.receive().await.unwrap();
117        let retained = bola_group.find_messages(&Default::default()).unwrap();
118
119        let stream = bola_group.stream().await.unwrap();
120        futures::pin_mut!(stream);
121        for expected in retained {
122            assert_eq!(stream.next().await.unwrap().unwrap(), expected);
123        }
124
125        amal_group
126            .send_message("hello".as_bytes(), SendMessageOpts::default())
127            .await
128            .unwrap();
129        let first_val = stream.next().await.unwrap().unwrap();
130        assert_eq!(first_val.decrypted_message_bytes, "hello".as_bytes());
131
132        amal_group
133            .send_message("goodbye".as_bytes(), SendMessageOpts::default())
134            .await
135            .unwrap();
136        let second_val = stream.next().await.unwrap().unwrap();
137        assert_eq!(second_val.decrypted_message_bytes, "goodbye".as_bytes());
138    }
139
140    // TODO: THIS TESTS ALSO LOSES MESSAGES
141    #[xmtp_common::timeout(Duration::from_secs(10))]
142    #[rstest::rstest]
143    #[xmtp_common::test(flavor = "multi_thread")]
144    #[cfg_attr(target_arch = "wasm32", ignore)]
145    async fn test_subscribe_multiple() {
146        let amal = Arc::new(ClientBuilder::new_test_client_vanilla(&generate_local_wallet()).await);
147        let group = amal.create_group(None, None).unwrap();
148
149        let stream = group.stream().await.unwrap();
150        futures::pin_mut!(stream);
151
152        for i in 0..10 {
153            group
154                .send_message(
155                    format!("hello {}", i).as_bytes(),
156                    SendMessageOpts::default(),
157                )
158                .await
159                .unwrap();
160        }
161
162        // Limit the stream so that it closes after 10 messages
163        let limited_stream = stream.take(10);
164        let values = limited_stream.collect::<Vec<_>>().await;
165        assert_eq!(values.len(), 10);
166        for value in values {
167            assert!(
168                value
169                    .unwrap()
170                    .decrypted_message_bytes
171                    .starts_with("hello".as_bytes())
172            );
173        }
174    }
175
176    #[xmtp_common::timeout(Duration::from_secs(5))]
177    #[rstest::rstest]
178    #[xmtp_common::test]
179    async fn test_subscribe_membership_changes() {
180        let amal = Arc::new(ClientBuilder::new_test_client(&generate_local_wallet()).await);
181        let bola = ClientBuilder::new_test_client(&generate_local_wallet()).await;
182
183        let amal_group = amal.create_group(None, None).unwrap();
184
185        let stream = amal_group.stream().await.unwrap();
186        futures::pin_mut!(stream);
187
188        amal_group.add_members(&[bola.inbox_id()]).await.unwrap();
189
190        let first_val = stream.next().await.unwrap().unwrap();
191        assert_eq!(first_val.kind, GroupMessageKind::MembershipChange);
192
193        amal_group
194            .send_message("hello".as_bytes(), SendMessageOpts::default())
195            .await
196            .unwrap();
197        let second_val = stream.next().await.unwrap().unwrap();
198        assert_eq!(second_val.decrypted_message_bytes, "hello".as_bytes());
199    }
200
201    // verifies: META-051
202    #[xmtp_common::test(unwrap_try = true)]
203    async fn test_process_streamed_group_message() {
204        use diesel::prelude::*;
205        use xmtp_db::schema::group_messages::dsl as message_dsl;
206
207        crate::tester!(alix, disable_workers);
208        crate::tester!(bo, disable_workers);
209        let group = alix.create_group(None, None)?;
210        group.add_members(&[bo.inbox_id()]).await?;
211        let bo_groups = bo.sync_welcomes().await?;
212        let bo_group = bo_groups.first().unwrap();
213        group
214            .send_message(b"test message", SendMessageOpts::default())
215            .await?;
216        let envelopes = alix
217            .context
218            .api()
219            .query_all(
220                std::collections::HashMap::from([(
221                    xmtp_proto::types::Topic::new_group_message(group.group_id),
222                    xmtp_proto::types::Cursor(0),
223                )]),
224                alix.context.api().limits().max_query_limit as u32,
225            )
226            .await?;
227        let envelope = envelopes.last().unwrap();
228        let result = bo_group
229            .process_streamed_group_message(envelope.encode_to_vec())
230            .await;
231        assert!(
232            result.is_ok(),
233            "Backend processing must succeed: {result:?}"
234        );
235        let messages = result?;
236        assert!(!messages.is_empty());
237        assert_eq!(messages.len(), 1);
238        assert_eq!(messages[0].decrypted_message_bytes, b"test message");
239
240        let message_id = messages[0].id.clone();
241        let cursor = xmtp_proto::types::Cursor(
242            envelope
243                .meta
244                .as_ref()
245                .unwrap()
246                .cursor
247                .as_ref()
248                .unwrap()
249                .sequence_id,
250        );
251        bo.context.db().raw_query(|conn| {
252            diesel::update(message_dsl::group_messages.find(&message_id))
253                .set(message_dsl::expire_at_ns.eq(Some(xmtp_common::time::now_ns() - 1)))
254                .execute(conn)
255        })?;
256        assert!(
257            bo.context
258                .db()
259                .get_group_message_by_cursor(bo_group.group_id, cursor)?
260                .is_some()
261        );
262        assert!(
263            bo_group
264                .process_streamed_group_message(envelope.encode_to_vec())
265                .await?
266                .is_empty()
267        );
268
269        bo.context.db().raw_query(|conn| {
270            diesel::update(message_dsl::group_messages.find(&message_id))
271                .set(message_dsl::expire_at_ns.eq(Some(i64::MAX)))
272                .execute(conn)
273        })?;
274        let visible = bo_group
275            .process_streamed_group_message(envelope.encode_to_vec())
276            .await?;
277        assert_eq!(visible.len(), 1);
278        assert_eq!(visible[0].id, message_id);
279    }
280}