mas_storage_pg/personal/
mod.rs

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

mas_storage_pg/personal/
mod.rs
