matrix-rust-sdk

public

Star0

Watch0

36 commits 4 branches 0 tags 0 releases16.6 MiB

use std::{
    future::Future,
    ops::Deref,
    option_env,
    path::{Path, PathBuf},
    sync::{Arc, LazyLock, Mutex as StdMutex},
    time::Duration,
};

use anyhow::Result;
use assign::assign;
use matrix_sdk::{
    Client, ClientBuilder, Room, ThreadingSupport,
    config::{RequestConfig, SyncSettings},
    encryption::EncryptionSettings,
    ruma::{
        RoomId,
        api::client::{account::register::v3::Request as RegistrationRequest, uiaa},
        time::Instant,
    },
    sliding_sync::VersionBuilder,
    sync::SyncResponse,
    timeout::ElapsedError,
};
use matrix_sdk_base::crypto::{CollectStrategy, DecryptionSettings, TrustRequirement};
use matrix_sdk_common::cross_process_lock::CrossProcessLockConfig;
use rand::RngExt;
use tempfile::{TempDir, tempdir};
use tokio::{sync::Mutex, time::sleep};

/// This global maintains temp directories alive for the whole lifetime of the
/// process.
static TMP_DIRS: LazyLock<Mutex<Vec<TempDir>>> = LazyLock::new(Mutex::default);

enum SqlitePath {
    Random,
    Path(PathBuf),
}

pub struct TestClientBuilder {
    username: String,
    use_sqlite_dir: Option<SqlitePath>,
    decryption_settings: Option<DecryptionSettings>,
    room_key_recipient_strategy: CollectStrategy,
    encryption_settings: EncryptionSettings,
    enable_share_history_on_invite: bool,
    threading_support: ThreadingSupport,
    http_proxy: Option<String>,
    cross_process_lock_config: CrossProcessLockConfig,
}

impl TestClientBuilder {
    pub fn new(username: impl AsRef<str>) -> Self {
        let suffix: u128 = rand::rng().random();
        let randomized_username = format!("{}{}", username.as_ref(), suffix);
        Self::with_exact_username(randomized_username)
    }

    pub fn with_exact_username(username: String) -> Self {
        Self {
            username,
            use_sqlite_dir: None,
            decryption_settings: None,
            encryption_settings: Default::default(),
            room_key_recipient_strategy: Default::default(),
            enable_share_history_on_invite: false,
            threading_support: ThreadingSupport::Disabled,
            http_proxy: None,
            cross_process_lock_config: CrossProcessLockConfig::SingleProcess,
        }
    }

    pub fn use_sqlite(mut self) -> Self {
        self.use_sqlite_dir = Some(SqlitePath::Random);
        self
    }

    /// Create or re-use a Sqlite store (with no passphrase) in the supplied
    /// directory. Note: this path must remain valid throughout the use of
    /// the constructed Client, so if you created a TempDir you must hang on
    /// to a reference to it throughout the test.
    pub fn use_sqlite_dir(mut self, path: &Path) -> Self {
        self.use_sqlite_dir = Some(SqlitePath::Path(path.to_owned()));
        self
    }

    pub fn encryption_settings(mut self, encryption_settings: EncryptionSettings) -> Self {
        self.encryption_settings = encryption_settings;
        self
    }

    pub fn enable_share_history_on_invite(mut self, enable_share_history_on_invite: bool) -> Self {
        self.enable_share_history_on_invite = enable_share_history_on_invite;
        self
    }

    pub fn enable_threading_support(mut self, thread_support: ThreadingSupport) -> Self {
        self.threading_support = thread_support;
        self
    }

    /// Simulate the behaviour of the clients when the "exclude insecure
    /// devices" (MSC4153) labs flag is enabled.
    pub fn exclude_insecure_devices(mut self, exclude_insecure_devices: bool) -> Self {
        let (sender_device_trust_requirement, room_key_recipient_strategy) =
            if exclude_insecure_devices {
                (TrustRequirement::CrossSignedOrLegacy, CollectStrategy::IdentityBasedStrategy)
            } else {
                (TrustRequirement::Untrusted, CollectStrategy::AllDevices)
            };
        self.decryption_settings = Some(DecryptionSettings { sender_device_trust_requirement });
        self.room_key_recipient_strategy = room_key_recipient_strategy;

        self
    }

    pub fn cross_process_lock_config(
        mut self,
        cross_process_lock_config: CrossProcessLockConfig,
    ) -> Self {
        self.cross_process_lock_config = cross_process_lock_config;
        self
    }

    fn common_client_builder(&self) -> ClientBuilder {
        let homeserver_url =
            option_env!("HOMESERVER_URL").unwrap_or("http://localhost:8228").to_owned();

        let mut client_builder = Client::builder()
            .user_agent("matrix-sdk-integration-tests")
            .homeserver_url(homeserver_url)
            .sliding_sync_version_builder(VersionBuilder::Native)
            .with_encryption_settings(self.encryption_settings)
            .with_room_key_recipient_strategy(self.room_key_recipient_strategy.clone())
            .with_enable_share_history_on_invite(self.enable_share_history_on_invite)
            .with_threading_support(self.threading_support)
            .request_config(RequestConfig::short_retry())
            .cross_process_store_config(self.cross_process_lock_config.clone());

        if let Some(decryption_settings) = &self.decryption_settings {
            client_builder = client_builder.with_decryption_settings(decryption_settings.clone())
        }

        if let Some(proxy) = &self.http_proxy {
            client_builder = client_builder.proxy(proxy);
        }

        client_builder
    }

    /// Create a new Client that is a copy of the supplied one, created using
    /// [`Client::restore_session`].
    pub async fn duplicate(self, other: &Client) -> Result<Client> {
        let client_builder = self.common_client_builder();

        let client = match self.use_sqlite_dir {
            Some(SqlitePath::Path(path_buf)) => {
                client_builder.sqlite_store(&path_buf, None).build().await?
            }
            _ => {
                panic!("You must call use_sqlite_dir for a duplicate client!");
            }
        };

        client
            .restore_session(
                other.session().expect("Session must be logged in before we can duplicate it"),
            )
            .await
            .expect("Failed to restore session");

        Ok(client)
    }

    pub async fn build(self) -> Result<Client> {
        let client_builder = self.common_client_builder();
        let client = match self.use_sqlite_dir {
            None => client_builder.build().await?,
            Some(SqlitePath::Random) => {
                let tmp_dir = tempdir()?;
                let client = client_builder.sqlite_store(tmp_dir.path(), None).build().await?;
                TMP_DIRS.lock().await.push(tmp_dir);
                client
            }
            Some(SqlitePath::Path(path_buf)) => {
                client_builder.sqlite_store(&path_buf, None).build().await?
            }
        };

        // safe to assume we have not registered this user yet, but ignore if we did

        let auth = client.matrix_auth();
        let mut try_login = true;
        if let Err(resp) = auth.register(RegistrationRequest::new()).await {
            // FIXME: do actually check the registration types...
            if let Some(_response) = resp.as_uiaa_response() {
                let request = assign!(RegistrationRequest::new(), {
                    username: Some(self.username.clone()),
                    password: Some(self.username.clone()),

                    auth: Some(uiaa::AuthData::Dummy(uiaa::Dummy::new())),
                });
                // if this failed, we will attempt to login after.
                try_login = auth.register(request).await.is_err();
            }
        }
        if try_login {
            auth.login_username(&self.username, &self.username).await?;
        }

        Ok(client)
    }
}

/// Client that correctly maintains and propagates sync token values.
#[derive(Clone)]
pub struct SyncTokenAwareClient {
    client: Client,
    token: Arc<StdMutex<Option<String>>>,
}

impl SyncTokenAwareClient {
    pub fn new(client: Client) -> Self {
        Self { client, token: Arc::new(None.into()) }
    }

    pub async fn sync_once(&self) -> Result<SyncResponse> {
        let mut settings = SyncSettings::default().timeout(Duration::from_secs(1));

        let token = { self.token.lock().unwrap().clone() };
        if let Some(token) = token {
            settings = settings.token(token);
        }

        let response = self.client.sync_once(settings).await?;

        let mut prev_token = self.token.lock().unwrap();
        if prev_token.as_ref() != Some(&response.next_batch) {
            *prev_token = Some(response.next_batch.clone());
        }
        Ok(response)
    }
}

impl Deref for SyncTokenAwareClient {
    type Target = Client;

    fn deref(&self) -> &Self::Target {
        &self.client
    }
}

/// Repeatedly run the given callback, at increasing intervals, until it returns
/// `Some`, or the given timeout has elapsed.
///
/// The callback `cb` is called with the attempt number as an argument.
///
/// # Returns
///
/// The result of `cb` if it returned `Some`. If the timeout elapses, an error
/// of type [`ElapsedError`].
pub async fn wait_until_some<F, Fut, R>(
    mut cb: F,
    timeout_duration: Duration,
) -> Result<R, ElapsedError>
where
    F: FnMut(u32) -> Fut,
    Fut: Future<Output = Option<R>>,
{
    let start = Instant::now();
    let mut attempt_count = 1;
    loop {
        if let Some(r) = cb(attempt_count).await {
            return Ok(r);
        }
        if start - Instant::now() > timeout_duration {
            return Err(ElapsedError());
        }
        sleep(Duration::from_millis(125) * attempt_count).await;
        attempt_count += 1;
    }
}

/// Waits for a room to arrive from a sync, for ~2 seconds.
///
/// Note: this doesn't run any sync, it assumes a sync has been running in the
/// background.
pub async fn wait_for_room(client: &Client, room_id: &RoomId) -> Room {
    let cb = async |i| {
        client.get_room(room_id).or_else(|| {
            tracing::warn!("attempt #{i}, alice couldn't see the room, retrying in a few…");
            None
        })
    };

    wait_until_some(cb, Duration::from_secs(2))
        .await
        .expect("room couldn't be fetched after all attempts")
}

1	use std::{
2	future::Future,
3	ops::Deref,
4	option_env,
5	path::{Path, PathBuf},
6	sync::{Arc, LazyLock, Mutex as StdMutex},
7	time::Duration,
8	};
9
10	use anyhow::Result;
11	use assign::assign;
12	use matrix_sdk::{
13	Client, ClientBuilder, Room, ThreadingSupport,
14	config::{RequestConfig, SyncSettings},
15	encryption::EncryptionSettings,
16	ruma::{
17	RoomId,
18	api::client::{account::register::v3::Request as RegistrationRequest, uiaa},
19	time::Instant,
20	},
21	sliding_sync::VersionBuilder,
22	sync::SyncResponse,
23	timeout::ElapsedError,
24	};
25	use matrix_sdk_base::crypto::{CollectStrategy, DecryptionSettings, TrustRequirement};
26	use matrix_sdk_common::cross_process_lock::CrossProcessLockConfig;
27	use rand::RngExt;
28	use tempfile::{TempDir, tempdir};
29	use tokio::{sync::Mutex, time::sleep};
30
31	/// This global maintains temp directories alive for the whole lifetime of the
32	/// process.
33	static TMP_DIRS: LazyLock<Mutex<Vec<TempDir>>> = LazyLock::new(Mutex::default);
34
35	enum SqlitePath {
36	Random,
37	Path(PathBuf),
38	}
39
40	pub struct TestClientBuilder {
41	username: String,
42	use_sqlite_dir: Option<SqlitePath>,
43	decryption_settings: Option<DecryptionSettings>,
44	room_key_recipient_strategy: CollectStrategy,
45	encryption_settings: EncryptionSettings,
46	enable_share_history_on_invite: bool,
47	threading_support: ThreadingSupport,
48	http_proxy: Option<String>,
49	cross_process_lock_config: CrossProcessLockConfig,
50	}
51
52	impl TestClientBuilder {
53	pub fn new(username: impl AsRef<str>) -> Self {
54	let suffix: u128 = rand::rng().random();
55	let randomized_username = format!("{}{}", username.as_ref(), suffix);
56	Self::with_exact_username(randomized_username)
57	}
58
59	pub fn with_exact_username(username: String) -> Self {
60	Self {
61	username,
62	use_sqlite_dir: None,
63	decryption_settings: None,
64	encryption_settings: Default::default(),
65	room_key_recipient_strategy: Default::default(),
66	enable_share_history_on_invite: false,
67	threading_support: ThreadingSupport::Disabled,
68	http_proxy: None,
69	cross_process_lock_config: CrossProcessLockConfig::SingleProcess,
70	}
71	}
72
73	pub fn use_sqlite(mut self) -> Self {
74	self.use_sqlite_dir = Some(SqlitePath::Random);
75	self
76	}
77
78	/// Create or re-use a Sqlite store (with no passphrase) in the supplied
79	/// directory. Note: this path must remain valid throughout the use of
80	/// the constructed Client, so if you created a TempDir you must hang on
81	/// to a reference to it throughout the test.
82	pub fn use_sqlite_dir(mut self, path: &Path) -> Self {
83	self.use_sqlite_dir = Some(SqlitePath::Path(path.to_owned()));
84	self
85	}
86
87	pub fn encryption_settings(mut self, encryption_settings: EncryptionSettings) -> Self {
88	self.encryption_settings = encryption_settings;
89	self
90	}
91
92	pub fn enable_share_history_on_invite(mut self, enable_share_history_on_invite: bool) -> Self {
93	self.enable_share_history_on_invite = enable_share_history_on_invite;
94	self
95	}
96
97	pub fn enable_threading_support(mut self, thread_support: ThreadingSupport) -> Self {
98	self.threading_support = thread_support;
99	self
100	}
101
102	/// Simulate the behaviour of the clients when the "exclude insecure
103	/// devices" (MSC4153) labs flag is enabled.
104	pub fn exclude_insecure_devices(mut self, exclude_insecure_devices: bool) -> Self {
105	let (sender_device_trust_requirement, room_key_recipient_strategy) =
106	if exclude_insecure_devices {
107	(TrustRequirement::CrossSignedOrLegacy, CollectStrategy::IdentityBasedStrategy)
108	} else {
109	(TrustRequirement::Untrusted, CollectStrategy::AllDevices)
110	};
111	self.decryption_settings = Some(DecryptionSettings { sender_device_trust_requirement });
112	self.room_key_recipient_strategy = room_key_recipient_strategy;
113
114	self
115	}
116
117	pub fn cross_process_lock_config(
118	mut self,
119	cross_process_lock_config: CrossProcessLockConfig,
120	) -> Self {
121	self.cross_process_lock_config = cross_process_lock_config;
122	self
123	}
124
125	fn common_client_builder(&self) -> ClientBuilder {
126	let homeserver_url =
127	option_env!("HOMESERVER_URL").unwrap_or("http://localhost:8228").to_owned();
128
129	let mut client_builder = Client::builder()
130	.user_agent("matrix-sdk-integration-tests")
131	.homeserver_url(homeserver_url)
132	.sliding_sync_version_builder(VersionBuilder::Native)
133	.with_encryption_settings(self.encryption_settings)
134	.with_room_key_recipient_strategy(self.room_key_recipient_strategy.clone())
135	.with_enable_share_history_on_invite(self.enable_share_history_on_invite)
136	.with_threading_support(self.threading_support)
137	.request_config(RequestConfig::short_retry())
138	.cross_process_store_config(self.cross_process_lock_config.clone());
139
140	if let Some(decryption_settings) = &self.decryption_settings {
141	client_builder = client_builder.with_decryption_settings(decryption_settings.clone())
142	}
143
144	if let Some(proxy) = &self.http_proxy {
145	client_builder = client_builder.proxy(proxy);
146	}
147
148	client_builder
149	}
150
151	/// Create a new Client that is a copy of the supplied one, created using
152	/// [`Client::restore_session`].
153	pub async fn duplicate(self, other: &Client) -> Result<Client> {
154	let client_builder = self.common_client_builder();
155
156	let client = match self.use_sqlite_dir {
157	Some(SqlitePath::Path(path_buf)) => {
158	client_builder.sqlite_store(&path_buf, None).build().await?
159	}
160	_ => {
161	panic!("You must call use_sqlite_dir for a duplicate client!");
162	}
163	};
164
165	client
166	.restore_session(
167	other.session().expect("Session must be logged in before we can duplicate it"),
168	)
169	.await
170	.expect("Failed to restore session");
171
172	Ok(client)
173	}
174
175	pub async fn build(self) -> Result<Client> {
176	let client_builder = self.common_client_builder();
177	let client = match self.use_sqlite_dir {
178	None => client_builder.build().await?,
179	Some(SqlitePath::Random) => {
180	let tmp_dir = tempdir()?;
181	let client = client_builder.sqlite_store(tmp_dir.path(), None).build().await?;
182	TMP_DIRS.lock().await.push(tmp_dir);
183	client
184	}
185	Some(SqlitePath::Path(path_buf)) => {
186	client_builder.sqlite_store(&path_buf, None).build().await?
187	}
188	};
189
190	// safe to assume we have not registered this user yet, but ignore if we did
191
192	let auth = client.matrix_auth();
193	let mut try_login = true;
194	if let Err(resp) = auth.register(RegistrationRequest::new()).await {
195	// FIXME: do actually check the registration types...
196	if let Some(_response) = resp.as_uiaa_response() {
197	let request = assign!(RegistrationRequest::new(), {
198	username: Some(self.username.clone()),
199	password: Some(self.username.clone()),
200
201	auth: Some(uiaa::AuthData::Dummy(uiaa::Dummy::new())),
202	});
203	// if this failed, we will attempt to login after.
204	try_login = auth.register(request).await.is_err();
205	}
206	}
207	if try_login {
208	auth.login_username(&self.username, &self.username).await?;
209	}
210
211	Ok(client)
212	}
213	}
214
215	/// Client that correctly maintains and propagates sync token values.
216	#[derive(Clone)]
217	pub struct SyncTokenAwareClient {
218	client: Client,
219	token: Arc<StdMutex<Option<String>>>,
220	}
221
222	impl SyncTokenAwareClient {
223	pub fn new(client: Client) -> Self {
224	Self { client, token: Arc::new(None.into()) }
225	}
226
227	pub async fn sync_once(&self) -> Result<SyncResponse> {
228	let mut settings = SyncSettings::default().timeout(Duration::from_secs(1));
229
230	let token = { self.token.lock().unwrap().clone() };
231	if let Some(token) = token {
232	settings = settings.token(token);
233	}
234
235	let response = self.client.sync_once(settings).await?;
236
237	let mut prev_token = self.token.lock().unwrap();
238	if prev_token.as_ref() != Some(&response.next_batch) {
239	*prev_token = Some(response.next_batch.clone());
240	}
241	Ok(response)
242	}
243	}
244
245	impl Deref for SyncTokenAwareClient {
246	type Target = Client;
247
248	fn deref(&self) -> &Self::Target {
249	&self.client
250	}
251	}
252
253	/// Repeatedly run the given callback, at increasing intervals, until it returns
254	/// `Some`, or the given timeout has elapsed.
255	///
256	/// The callback `cb` is called with the attempt number as an argument.
257	///
258	/// # Returns
259	///
260	/// The result of `cb` if it returned `Some`. If the timeout elapses, an error
261	/// of type [`ElapsedError`].
262	pub async fn wait_until_some<F, Fut, R>(
263	mut cb: F,
264	timeout_duration: Duration,
265	) -> Result<R, ElapsedError>
266	where
267	F: FnMut(u32) -> Fut,
268	Fut: Future<Output = Option<R>>,
269	{
270	let start = Instant::now();
271	let mut attempt_count = 1;
272	loop {
273	if let Some(r) = cb(attempt_count).await {
274	return Ok(r);
275	}
276	if start - Instant::now() > timeout_duration {
277	return Err(ElapsedError());
278	}
279	sleep(Duration::from_millis(125) * attempt_count).await;
280	attempt_count += 1;
281	}
282	}
283
284	/// Waits for a room to arrive from a sync, for ~2 seconds.
285	///
286	/// Note: this doesn't run any sync, it assumes a sync has been running in the
287	/// background.
288	pub async fn wait_for_room(client: &Client, room_id: &RoomId) -> Room {
289	let cb = async \|i\| {
290	client.get_room(room_id).or_else(\|\| {
291	tracing::warn!("attempt #{i}, alice couldn't see the room, retrying in a few…");
292	None
293	})
294	};
295
296	wait_until_some(cb, Duration::from_secs(2))
297	.await
298	.expect("room couldn't be fetched after all attempts")
299	}