Skip to main content

mas_storage_pg/personal/
mod.rs

1// Copyright 2025, 2026 Element Creations Ltd.
2// Copyright 2025 New Vector Ltd.
3//
4// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
5// Please see LICENSE files in the repository root for full details.
6
7//! A module containing the PostgreSQL implementations of the
8//! Personal Access Token / Personal Session repositories
9
10mod access_token;
11mod session;
12
13pub use access_token::PgPersonalAccessTokenRepository;
14pub use session::PgPersonalSessionRepository;
15
16#[cfg(test)]
17mod tests {
18    use chrono::Duration;
19    use mas_data_model::{
20        Clock, Device, clock::MockClock, personal::session::PersonalSessionOwner,
21    };
22    use mas_storage::{
23        Pagination, RepositoryAccess,
24        personal::{
25            PersonalAccessTokenRepository, PersonalSessionFilter, PersonalSessionRepository,
26        },
27        user::UserRepository,
28    };
29    use oauth2_types::scope::{OPENID, PROFILE, Scope};
30    use rand::SeedableRng;
31    use rand_chacha::ChaChaRng;
32    use sqlx::PgPool;
33
34    use crate::PgRepository;
35
36    #[sqlx::test(migrator = "crate::MIGRATOR")]
37    async fn test_session_repository(pool: PgPool) {
38        let mut rng = ChaChaRng::seed_from_u64(42);
39        let clock = MockClock::default();
40        let mut repo = PgRepository::from_pool(&pool).await.unwrap();
41
42        // Create a user
43        let admin_user = repo
44            .user()
45            .add(&mut rng, &clock, "john".to_owned())
46            .await
47            .unwrap();
48        let bot_user = repo
49            .user()
50            .add(&mut rng, &clock, "marvin".to_owned())
51            .await
52            .unwrap();
53
54        let all = PersonalSessionFilter::new().for_actor_user(&bot_user);
55        let active = all.active_only();
56        let finished = all.finished_only();
57        let pagination = Pagination::first(10);
58
59        assert_eq!(repo.personal_session().count(all).await.unwrap(), 0);
60        assert_eq!(repo.personal_session().count(active).await.unwrap(), 0);
61        assert_eq!(repo.personal_session().count(finished).await.unwrap(), 0);
62
63        // We start off with no sessions
64        let full_list = repo.personal_session().list(all, pagination).await.unwrap();
65        assert_eq!(full_list.edges, []);
66        let active_list = repo
67            .personal_session()
68            .list(active, pagination)
69            .await
70            .unwrap();
71        assert_eq!(active_list.edges, []);
72        let finished_list = repo
73            .personal_session()
74            .list(finished, pagination)
75            .await
76            .unwrap();
77        assert_eq!(finished_list.edges, []);
78
79        // Start a personal session for that user
80        let device = Device::generate(&mut rng);
81        let scope: Scope = [OPENID, PROFILE]
82            .into_iter()
83            .chain(device.to_scope_token().unwrap())
84            .collect();
85        let session = repo
86            .personal_session()
87            .add(
88                &mut rng,
89                &clock,
90                (&admin_user).into(),
91                &bot_user,
92                "Test Personal Session".to_owned(),
93                scope.clone(),
94            )
95            .await
96            .unwrap();
97        assert_eq!(session.owner, PersonalSessionOwner::User(admin_user.id));
98        assert_eq!(session.actor_user_id, bot_user.id);
99        assert!(session.is_valid());
100        assert!(!session.is_revoked());
101        assert_eq!(session.scope, scope);
102
103        assert_eq!(repo.personal_session().count(all).await.unwrap(), 1);
104        assert_eq!(repo.personal_session().count(active).await.unwrap(), 1);
105        assert_eq!(repo.personal_session().count(finished).await.unwrap(), 0);
106
107        let full_list = repo.personal_session().list(all, pagination).await.unwrap();
108        assert_eq!(full_list.edges.len(), 1);
109        assert_eq!(full_list.edges[0].node.0.id, session.id);
110        assert!(full_list.edges[0].node.0.is_valid());
111        let active_list = repo
112            .personal_session()
113            .list(active, pagination)
114            .await
115            .unwrap();
116        assert_eq!(active_list.edges.len(), 1);
117        assert_eq!(active_list.edges[0].node.0.id, session.id);
118        assert!(active_list.edges[0].node.0.is_valid());
119        let finished_list = repo
120            .personal_session()
121            .list(finished, pagination)
122            .await
123            .unwrap();
124        assert_eq!(finished_list.edges, []);
125
126        // Lookup the session and check it didn't change
127        let session_lookup = repo
128            .personal_session()
129            .lookup(session.id)
130            .await
131            .unwrap()
132            .expect("personal session not found");
133        assert_eq!(session_lookup.id, session.id);
134        assert_eq!(
135            session_lookup.owner,
136            PersonalSessionOwner::User(admin_user.id)
137        );
138        assert_eq!(session_lookup.actor_user_id, bot_user.id);
139        assert_eq!(session_lookup.scope, scope);
140        assert!(session_lookup.is_valid());
141        assert!(!session_lookup.is_revoked());
142
143        // Revoke the session
144        let session = repo
145            .personal_session()
146            .revoke(&clock, session)
147            .await
148            .unwrap();
149        assert!(!session.is_valid());
150        assert!(session.is_revoked());
151
152        assert_eq!(repo.personal_session().count(all).await.unwrap(), 1);
153        assert_eq!(repo.personal_session().count(active).await.unwrap(), 0);
154        assert_eq!(repo.personal_session().count(finished).await.unwrap(), 1);
155
156        let full_list = repo.personal_session().list(all, pagination).await.unwrap();
157        assert_eq!(full_list.edges.len(), 1);
158        assert_eq!(full_list.edges[0].node.0.id, session.id);
159        let active_list = repo
160            .personal_session()
161            .list(active, pagination)
162            .await
163            .unwrap();
164        assert_eq!(active_list.edges, []);
165        let finished_list = repo
166            .personal_session()
167            .list(finished, pagination)
168            .await
169            .unwrap();
170        assert_eq!(finished_list.edges.len(), 1);
171        assert_eq!(finished_list.edges[0].node.0.id, session.id);
172        assert!(finished_list.edges[0].node.0.is_revoked());
173
174        // Reload the session and check again
175        let session_lookup = repo
176            .personal_session()
177            .lookup(session.id)
178            .await
179            .unwrap()
180            .expect("personal session not found");
181        assert!(!session_lookup.is_valid());
182        assert!(session_lookup.is_revoked());
183    }
184
185    #[sqlx::test(migrator = "crate::MIGRATOR")]
186    async fn test_session_revoke_bulk(pool: PgPool) {
187        let mut rng = ChaChaRng::seed_from_u64(42);
188        let clock = MockClock::default();
189        let mut repo = PgRepository::from_pool(&pool).await.unwrap();
190
191        let alice_user = repo
192            .user()
193            .add(&mut rng, &clock, "alice".to_owned())
194            .await
195            .unwrap();
196        let bob_user = repo
197            .user()
198            .add(&mut rng, &clock, "bob".to_owned())
199            .await
200            .unwrap();
201
202        let session1 = repo
203            .personal_session()
204            .add(
205                &mut rng,
206                &clock,
207                (&alice_user).into(),
208                &bob_user,
209                "Test Personal Session".to_owned(),
210                "openid".parse().unwrap(),
211            )
212            .await
213            .unwrap();
214        repo.personal_access_token()
215            .add(
216                &mut rng,
217                &clock,
218                &session1,
219                "mpt_hiss",
220                Some(Duration::days(42)),
221            )
222            .await
223            .unwrap();
224
225        let session2 = repo
226            .personal_session()
227            .add(
228                &mut rng,
229                &clock,
230                (&bob_user).into(),
231                &bob_user,
232                "Test Personal Session".to_owned(),
233                "openid".parse().unwrap(),
234            )
235            .await
236            .unwrap();
237        repo.personal_access_token()
238            .add(
239                &mut rng, &clock, &session2, "mpt_meow", // No expiry
240                None,
241            )
242            .await
243            .unwrap();
244
245        // Just one session without a token expiry time
246        assert_eq!(
247            repo.personal_session()
248                .revoke_bulk(
249                    &clock,
250                    PersonalSessionFilter::new()
251                        .active_only()
252                        .with_expires(false)
253                )
254                .await
255                .unwrap(),
256            1
257        );
258
259        // Just one session with a token expiry time
260        assert_eq!(
261            repo.personal_session()
262                .revoke_bulk(
263                    &clock,
264                    PersonalSessionFilter::new()
265                        .active_only()
266                        .with_expires(true)
267                )
268                .await
269                .unwrap(),
270            1
271        );
272
273        // No active sessions left
274        assert_eq!(
275            repo.personal_session()
276                .revoke_bulk(&clock, PersonalSessionFilter::new().active_only())
277                .await
278                .unwrap(),
279            0
280        );
281    }
282
283    #[sqlx::test(migrator = "crate::MIGRATOR")]
284    async fn test_access_token_repository(pool: PgPool) {
285        const FIRST_TOKEN: &str = "first_access_token";
286        const SECOND_TOKEN: &str = "second_access_token";
287        let mut rng = ChaChaRng::seed_from_u64(42);
288        let clock = MockClock::default();
289        let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed();
290
291        // Create a user
292        let admin_user = repo
293            .user()
294            .add(&mut rng, &clock, "john".to_owned())
295            .await
296            .unwrap();
297        let bot_user = repo
298            .user()
299            .add(&mut rng, &clock, "marvin".to_owned())
300            .await
301            .unwrap();
302
303        // Start a personal session for that user
304        let device = Device::generate(&mut rng);
305        let scope: Scope = [OPENID, PROFILE]
306            .into_iter()
307            .chain(device.to_scope_token().unwrap())
308            .collect();
309        let session = repo
310            .personal_session()
311            .add(
312                &mut rng,
313                &clock,
314                (&admin_user).into(),
315                &bot_user,
316                "Test Personal Session".to_owned(),
317                scope,
318            )
319            .await
320            .unwrap();
321
322        // Add an access token to that session
323        let token = repo
324            .personal_access_token()
325            .add(
326                &mut rng,
327                &clock,
328                &session,
329                FIRST_TOKEN,
330                Some(Duration::try_minutes(1).unwrap()),
331            )
332            .await
333            .unwrap();
334        assert_eq!(token.session_id, session.id);
335
336        // Commit the txn and grab a new transaction, to test a conflict
337        repo.save().await.unwrap();
338
339        {
340            let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed();
341            // Adding the same token a second time should conflict
342            assert!(
343                repo.personal_access_token()
344                    .add(
345                        &mut rng,
346                        &clock,
347                        &session,
348                        FIRST_TOKEN,
349                        Some(Duration::try_minutes(1).unwrap()),
350                    )
351                    .await
352                    .is_err()
353            );
354            repo.cancel().await.unwrap();
355        }
356
357        // Grab a new repo
358        let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed();
359
360        // Looking up via ID works
361        let token_lookup = repo
362            .personal_access_token()
363            .lookup(token.id)
364            .await
365            .unwrap()
366            .expect("personal access token not found");
367        assert_eq!(token.id, token_lookup.id);
368        assert_eq!(token_lookup.session_id, session.id);
369
370        // Looking up via the token value works
371        let token_lookup = repo
372            .personal_access_token()
373            .find_by_token(FIRST_TOKEN)
374            .await
375            .unwrap()
376            .expect("personal access token not found");
377        assert_eq!(token.id, token_lookup.id);
378        assert_eq!(token_lookup.session_id, session.id);
379
380        // Token is currently valid
381        assert!(token.is_valid(clock.now()));
382
383        clock.advance(Duration::try_minutes(1).unwrap());
384        // Token should have expired
385        assert!(!token.is_valid(clock.now()));
386
387        // Add a second access token, this time without expiration
388        let _token = repo
389            .personal_access_token()
390            .revoke(&clock, token)
391            .await
392            .unwrap();
393        let token = repo
394            .personal_access_token()
395            .add(&mut rng, &clock, &session, SECOND_TOKEN, None)
396            .await
397            .unwrap();
398        assert_eq!(token.session_id, session.id);
399
400        // Token is currently valid
401        assert!(token.is_valid(clock.now()));
402
403        // Revoke it
404        let _token = repo
405            .personal_access_token()
406            .revoke(&clock, token)
407            .await
408            .unwrap();
409
410        // Reload it
411        let token = repo
412            .personal_access_token()
413            .find_by_token(SECOND_TOKEN)
414            .await
415            .unwrap()
416            .expect("personal access token not found");
417
418        // Token is not valid anymore
419        assert!(!token.is_valid(clock.now()));
420
421        repo.save().await.unwrap();
422    }
423}