mas_matrix/
mock.rs

1// Copyright 2025, 2026 Element Creations Ltd.
2// Copyright 2024, 2025 New Vector Ltd.
3// Copyright 2023, 2024 The Matrix.org Foundation C.I.C.
4//
5// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
6// Please see LICENSE files in the repository root for full details.
7
8use std::{
9    collections::{HashMap, HashSet},
10    sync::atomic::{AtomicBool, Ordering},
11};
12
13use anyhow::Context;
14use async_trait::async_trait;
15use tokio::sync::RwLock;
16
17use crate::{HomeserverConnection as _, MatrixUser, ProvisionRequest};
18
19#[derive(Clone)]
20pub struct MockUser {
21    sub: String,
22    avatar_url: Option<String>,
23    displayname: Option<String>,
24    devices: HashSet<String>,
25    emails: Option<Vec<String>>,
26    cross_signing_reset_allowed: bool,
27    deactivated: bool,
28    pub locked: bool,
29}
30
31/// A mock implementation of a [`HomeserverConnection`], which doesn't do
32/// anything and only fails when explicitly told to.
33pub struct HomeserverConnection {
34    homeserver: String,
35    users: RwLock<HashMap<String, MockUser>>,
36    reserved_localparts: RwLock<HashSet<&'static str>>,
37    fail_localpart_availability: AtomicBool,
38}
39
40impl HomeserverConnection {
41    /// A valid bearer token that will be accepted by
42    /// [`crate::HomeserverConnection::verify_token`].
43    pub const VALID_BEARER_TOKEN: &str = "mock_homeserver_bearer_token";
44
45    /// Create a new mock connection.
46    pub fn new<H>(homeserver: H) -> Self
47    where
48        H: Into<String>,
49    {
50        Self {
51            homeserver: homeserver.into(),
52            users: RwLock::new(HashMap::new()),
53            reserved_localparts: RwLock::new(HashSet::new()),
54            fail_localpart_availability: AtomicBool::new(false),
55        }
56    }
57
58    pub async fn reserve_localpart(&self, localpart: &'static str) {
59        self.reserved_localparts.write().await.insert(localpart);
60    }
61
62    /// Make [`crate::HomeserverConnection::is_localpart_available`] fail, to
63    /// exercise homeserver failures.
64    pub fn fail_localpart_availability(&self) {
65        self.fail_localpart_availability
66            .store(true, Ordering::Relaxed);
67    }
68
69    /// Like `query_user` but get the raw test state of the user.
70    ///
71    /// # Errors
72    ///
73    /// Will fail if the user doesn't exist.
74    pub async fn query_user_raw(&self, localpart: &str) -> Result<MockUser, anyhow::Error> {
75        let mxid = self.mxid(localpart);
76        let users = self.users.read().await;
77        let user = users.get(&mxid).context("User not found")?;
78        Ok(user.clone())
79    }
80}
81
82#[async_trait]
83impl crate::HomeserverConnection for HomeserverConnection {
84    fn homeserver(&self) -> &str {
85        &self.homeserver
86    }
87
88    async fn verify_token(&self, token: &str) -> Result<bool, anyhow::Error> {
89        Ok(token == Self::VALID_BEARER_TOKEN)
90    }
91
92    async fn query_user(&self, localpart: &str) -> Result<MatrixUser, anyhow::Error> {
93        let mxid = self.mxid(localpart);
94        let users = self.users.read().await;
95        let user = users.get(&mxid).context("User not found")?;
96        Ok(MatrixUser {
97            displayname: user.displayname.clone(),
98            avatar_url: user.avatar_url.clone(),
99            deactivated: user.deactivated,
100        })
101    }
102
103    async fn provision_user(&self, request: &ProvisionRequest) -> Result<bool, anyhow::Error> {
104        let mut users = self.users.write().await;
105        let mxid = self.mxid(request.localpart());
106        let inserted = !users.contains_key(&mxid);
107        let user = users.entry(mxid).or_insert(MockUser {
108            sub: request.sub().to_owned(),
109            avatar_url: None,
110            displayname: None,
111            devices: HashSet::new(),
112            emails: None,
113            cross_signing_reset_allowed: false,
114            deactivated: false,
115            locked: false,
116        });
117
118        anyhow::ensure!(
119            user.sub == request.sub(),
120            "User already provisioned with different sub"
121        );
122
123        request.on_emails(|emails| {
124            user.emails = emails.map(ToOwned::to_owned);
125        });
126
127        request.on_displayname(|displayname| {
128            user.displayname = displayname.map(ToOwned::to_owned);
129        });
130
131        request.on_avatar_url(|avatar_url| {
132            user.avatar_url = avatar_url.map(ToOwned::to_owned);
133        });
134
135        user.locked = request.locked();
136
137        Ok(inserted)
138    }
139
140    async fn is_localpart_available(&self, localpart: &str) -> Result<bool, anyhow::Error> {
141        anyhow::ensure!(
142            !self.fail_localpart_availability.load(Ordering::Relaxed),
143            "Mock homeserver failure while checking localpart availability"
144        );
145
146        if self.reserved_localparts.read().await.contains(localpart) {
147            return Ok(false);
148        }
149
150        let mxid = self.mxid(localpart);
151        let users = self.users.read().await;
152        Ok(!users.contains_key(&mxid))
153    }
154
155    async fn upsert_device(
156        &self,
157        localpart: &str,
158        device_id: &str,
159        _initial_display_name: Option<&str>,
160    ) -> Result<(), anyhow::Error> {
161        let mxid = self.mxid(localpart);
162        let mut users = self.users.write().await;
163        let user = users.get_mut(&mxid).context("User not found")?;
164        user.devices.insert(device_id.to_owned());
165        Ok(())
166    }
167
168    async fn update_device_display_name(
169        &self,
170        localpart: &str,
171        device_id: &str,
172        _display_name: &str,
173    ) -> Result<(), anyhow::Error> {
174        let mxid = self.mxid(localpart);
175        let mut users = self.users.write().await;
176        let user = users.get_mut(&mxid).context("User not found")?;
177        user.devices.get(device_id).context("Device not found")?;
178        Ok(())
179    }
180
181    async fn delete_device(&self, localpart: &str, device_id: &str) -> Result<(), anyhow::Error> {
182        let mxid = self.mxid(localpart);
183        let mut users = self.users.write().await;
184        let user = users.get_mut(&mxid).context("User not found")?;
185        user.devices.remove(device_id);
186        Ok(())
187    }
188
189    async fn sync_devices(
190        &self,
191        localpart: &str,
192        devices: HashSet<String>,
193    ) -> Result<(), anyhow::Error> {
194        let mxid = self.mxid(localpart);
195        let mut users = self.users.write().await;
196        let user = users.get_mut(&mxid).context("User not found")?;
197        user.devices = devices;
198        Ok(())
199    }
200
201    async fn delete_user(&self, localpart: &str, erase: bool) -> Result<(), anyhow::Error> {
202        let mxid = self.mxid(localpart);
203        let mut users = self.users.write().await;
204        let user = users.get_mut(&mxid).context("User not found")?;
205        user.devices.clear();
206        user.emails = None;
207        user.deactivated = true;
208        if erase {
209            user.avatar_url = None;
210            user.displayname = None;
211        }
212
213        Ok(())
214    }
215
216    async fn reactivate_user(&self, localpart: &str) -> Result<(), anyhow::Error> {
217        let mxid = self.mxid(localpart);
218        let mut users = self.users.write().await;
219        let user = users.get_mut(&mxid).context("User not found")?;
220        user.deactivated = false;
221
222        Ok(())
223    }
224
225    async fn set_displayname(
226        &self,
227        localpart: &str,
228        displayname: &str,
229    ) -> Result<(), anyhow::Error> {
230        let mxid = self.mxid(localpart);
231        let mut users = self.users.write().await;
232        let user = users.get_mut(&mxid).context("User not found")?;
233        user.displayname = Some(displayname.to_owned());
234        Ok(())
235    }
236
237    async fn unset_displayname(&self, localpart: &str) -> Result<(), anyhow::Error> {
238        let mxid = self.mxid(localpart);
239        let mut users = self.users.write().await;
240        let user = users.get_mut(&mxid).context("User not found")?;
241        user.displayname = None;
242        Ok(())
243    }
244
245    async fn allow_cross_signing_reset(&self, localpart: &str) -> Result<(), anyhow::Error> {
246        let mxid = self.mxid(localpart);
247        let mut users = self.users.write().await;
248        let user = users.get_mut(&mxid).context("User not found")?;
249        user.cross_signing_reset_allowed = true;
250        Ok(())
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257
258    #[tokio::test]
259    async fn test_mock_connection() {
260        let conn = HomeserverConnection::new("example.org");
261
262        let mxid = "@test:example.org";
263        let device = "test";
264        assert_eq!(conn.homeserver(), "example.org");
265        assert_eq!(conn.mxid("test"), mxid);
266
267        assert!(conn.query_user("test").await.is_err());
268        assert!(conn.upsert_device("test", device, None).await.is_err());
269        assert!(conn.delete_device("test", device).await.is_err());
270
271        let request = ProvisionRequest::new("test", "test", false)
272            .set_displayname("Test User".into())
273            .set_avatar_url("mxc://example.org/1234567890".into())
274            .set_emails(vec!["test@example.org".to_owned()]);
275
276        let inserted = conn.provision_user(&request).await.unwrap();
277        assert!(inserted);
278
279        let user = conn.query_user("test").await.unwrap();
280        assert_eq!(user.displayname, Some("Test User".into()));
281        assert_eq!(user.avatar_url, Some("mxc://example.org/1234567890".into()));
282
283        // Set the displayname again
284        assert!(conn.set_displayname("test", "John").await.is_ok());
285
286        let user = conn.query_user("test").await.unwrap();
287        assert_eq!(user.displayname, Some("John".into()));
288
289        // Unset the displayname
290        assert!(conn.unset_displayname("test").await.is_ok());
291
292        let user = conn.query_user("test").await.unwrap();
293        assert_eq!(user.displayname, None);
294
295        // Deleting a non-existent device should not fail
296        assert!(conn.delete_device("test", device).await.is_ok());
297
298        // Create the device
299        assert!(conn.upsert_device("test", device, None).await.is_ok());
300        // Create the same device again
301        assert!(conn.upsert_device("test", device, None).await.is_ok());
302
303        // XXX: there is no API to query devices yet in the trait
304        // Delete the device
305        assert!(conn.delete_device("test", device).await.is_ok());
306
307        // The user we just created should be not available
308        assert!(!conn.is_localpart_available("test").await.unwrap());
309        // But another user should be
310        assert!(conn.is_localpart_available("alice").await.unwrap());
311
312        // Reserve the localpart, it should not be available anymore
313        conn.reserve_localpart("alice").await;
314        assert!(!conn.is_localpart_available("alice").await.unwrap());
315    }
316}
mas_matrix/mock.rs

mas_matrix/
mock.rs
