mas_storage_pg/upstream_oauth2/
mod.rs

1// Copyright 2025, 2026 Element Creations Ltd.
2// Copyright 2024, 2025 New Vector Ltd.
3// Copyright 2022-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
8//! A module containing the PostgreSQL implementation of the repositories
9//! related to the upstream OAuth 2.0 providers
10
11mod link;
12mod provider;
13mod session;
14
15pub use self::{
16    link::PgUpstreamOAuthLinkRepository, provider::PgUpstreamOAuthProviderRepository,
17    session::PgUpstreamOAuthSessionRepository,
18};
19
20#[cfg(test)]
21mod tests {
22    use chrono::Duration;
23    use mas_data_model::{
24        UpstreamOAuthProviderClaimsImports, UpstreamOAuthProviderOnBackchannelLogout,
25        UpstreamOAuthProviderTokenAuthMethod, clock::MockClock,
26    };
27    use mas_iana::jose::JsonWebSignatureAlg;
28    use mas_storage::{
29        Pagination, RepositoryAccess,
30        upstream_oauth2::{
31            UpstreamOAuthLinkFilter, UpstreamOAuthLinkRepository, UpstreamOAuthProviderFilter,
32            UpstreamOAuthProviderParams, UpstreamOAuthProviderRepository,
33            UpstreamOAuthSessionFilter, UpstreamOAuthSessionRepository,
34        },
35        user::UserRepository,
36    };
37    use oauth2_types::scope::{OPENID, Scope};
38    use rand::SeedableRng;
39    use sqlx::PgPool;
40
41    use crate::PgRepository;
42
43    #[sqlx::test(migrator = "crate::MIGRATOR")]
44    async fn test_repository(pool: PgPool) {
45        let mut rng = rand_chacha::ChaChaRng::seed_from_u64(42);
46        let clock = MockClock::default();
47        let mut repo = PgRepository::from_pool(&pool).await.unwrap();
48
49        // The provider list should be empty at the start
50        let all_providers = repo.upstream_oauth_provider().all_enabled().await.unwrap();
51        assert_eq!(all_providers, []);
52
53        // Let's add a provider
54        let provider = repo
55            .upstream_oauth_provider()
56            .add(
57                &mut rng,
58                &clock,
59                UpstreamOAuthProviderParams {
60                    issuer: Some("https://example.com/".to_owned()),
61                    human_name: None,
62                    brand_name: None,
63                    scope: Scope::from_iter([OPENID]),
64                    token_endpoint_auth_method: UpstreamOAuthProviderTokenAuthMethod::None,
65                    id_token_signed_response_alg: JsonWebSignatureAlg::Rs256,
66                    fetch_userinfo: false,
67                    userinfo_signed_response_alg: None,
68                    token_endpoint_signing_alg: None,
69                    client_id: "client-id".to_owned(),
70                    encrypted_client_secret: None,
71                    claims_imports: UpstreamOAuthProviderClaimsImports::default(),
72                    token_endpoint_override: None,
73                    authorization_endpoint_override: None,
74                    userinfo_endpoint_override: None,
75                    jwks_uri_override: None,
76                    discovery_mode: mas_data_model::UpstreamOAuthProviderDiscoveryMode::Oidc,
77                    pkce_mode: mas_data_model::UpstreamOAuthProviderPkceMode::Auto,
78                    response_mode: None,
79                    additional_authorization_parameters: Vec::new(),
80                    forward_login_hint: false,
81                    ui_order: 0,
82                    on_backchannel_logout: UpstreamOAuthProviderOnBackchannelLogout::DoNothing,
83                    registration_token_required: false,
84                },
85            )
86            .await
87            .unwrap();
88
89        // Look it up in the database
90        let provider = repo
91            .upstream_oauth_provider()
92            .lookup(provider.id)
93            .await
94            .unwrap()
95            .expect("provider to be found in the database");
96        assert_eq!(provider.issuer.as_deref(), Some("https://example.com/"));
97        assert_eq!(provider.client_id, "client-id");
98
99        // It should be in the list of all providers
100        let providers = repo.upstream_oauth_provider().all_enabled().await.unwrap();
101        assert_eq!(providers.len(), 1);
102        assert_eq!(providers[0].issuer.as_deref(), Some("https://example.com/"));
103        assert_eq!(providers[0].client_id, "client-id");
104
105        // Start a session
106        let session = repo
107            .upstream_oauth_session()
108            .add(
109                &mut rng,
110                &clock,
111                &provider,
112                "some-state".to_owned(),
113                None,
114                Some("some-nonce".to_owned()),
115            )
116            .await
117            .unwrap();
118
119        // Look it up in the database
120        let session = repo
121            .upstream_oauth_session()
122            .lookup(session.id)
123            .await
124            .unwrap()
125            .expect("session to be found in the database");
126        assert_eq!(session.provider_id, provider.id);
127        assert_eq!(session.link_id(), None);
128        assert!(session.is_pending());
129        assert!(!session.is_completed());
130        assert!(!session.is_consumed());
131
132        // Create a link
133        let link = repo
134            .upstream_oauth_link()
135            .add(
136                &mut rng,
137                &clock,
138                &provider,
139                "a-subject".to_owned(),
140                Some("alice@example.com".to_owned()),
141            )
142            .await
143            .unwrap();
144
145        // We can look it up by its ID
146        repo.upstream_oauth_link()
147            .lookup(link.id)
148            .await
149            .unwrap()
150            .expect("link to be found in database");
151
152        // or by its subject
153        let link = repo
154            .upstream_oauth_link()
155            .find_by_subject(&provider, "a-subject")
156            .await
157            .unwrap()
158            .expect("link to be found in database");
159        assert_eq!(link.subject, "a-subject");
160        assert_eq!(link.provider_id, provider.id);
161
162        let session = repo
163            .upstream_oauth_session()
164            .complete_with_link(&clock, session, &link, None, None, None, None)
165            .await
166            .unwrap();
167        // Reload the session
168        let session = repo
169            .upstream_oauth_session()
170            .lookup(session.id)
171            .await
172            .unwrap()
173            .expect("session to be found in the database");
174        assert!(session.is_completed());
175        assert!(!session.is_consumed());
176        assert_eq!(session.link_id(), Some(link.id));
177
178        // We need to create a user and start a browser session to consume the session
179        let user = repo
180            .user()
181            .add(&mut rng, &clock, "john".to_owned())
182            .await
183            .unwrap();
184        let browser_session = repo
185            .browser_session()
186            .add(&mut rng, &clock, &user, None)
187            .await
188            .unwrap();
189
190        let session = repo
191            .upstream_oauth_session()
192            .consume(&clock, session, &browser_session)
193            .await
194            .unwrap();
195
196        // Reload the session
197        let session = repo
198            .upstream_oauth_session()
199            .lookup(session.id)
200            .await
201            .unwrap()
202            .expect("session to be found in the database");
203        assert!(session.is_consumed());
204
205        repo.upstream_oauth_link()
206            .associate_to_user(&link, &user)
207            .await
208            .unwrap();
209
210        // XXX: we should also try other combinations of the filter
211        let filter = UpstreamOAuthLinkFilter::new()
212            .for_user(&user)
213            .for_provider(&provider)
214            .for_subject("a-subject")
215            .enabled_providers_only();
216
217        let links = repo
218            .upstream_oauth_link()
219            .list(filter, Pagination::first(10))
220            .await
221            .unwrap();
222        assert!(!links.has_previous_page);
223        assert!(!links.has_next_page);
224        assert_eq!(links.edges.len(), 1);
225        assert_eq!(links.edges[0].node.id, link.id);
226        assert_eq!(links.edges[0].node.user_id, Some(user.id));
227
228        assert_eq!(repo.upstream_oauth_link().count(filter).await.unwrap(), 1);
229
230        // Filtering on a case-insensitive partial human_account_name should match
231        let matching_filter = UpstreamOAuthLinkFilter::new().matching_human_account_name("LICE");
232        let links = repo
233            .upstream_oauth_link()
234            .list(matching_filter, Pagination::first(10))
235            .await
236            .unwrap();
237        assert_eq!(links.edges.len(), 1);
238        assert_eq!(links.edges[0].node.id, link.id);
239        assert_eq!(
240            repo.upstream_oauth_link()
241                .count(matching_filter)
242                .await
243                .unwrap(),
244            1
245        );
246
247        // A non-matching human_account_name should return nothing
248        let non_matching_filter =
249            UpstreamOAuthLinkFilter::new().matching_human_account_name("nope");
250        let links = repo
251            .upstream_oauth_link()
252            .list(non_matching_filter, Pagination::first(10))
253            .await
254            .unwrap();
255        assert_eq!(links.edges.len(), 0);
256        assert_eq!(
257            repo.upstream_oauth_link()
258                .count(non_matching_filter)
259                .await
260                .unwrap(),
261            0
262        );
263
264        // There should be exactly one enabled provider
265        assert_eq!(
266            repo.upstream_oauth_provider()
267                .count(UpstreamOAuthProviderFilter::new())
268                .await
269                .unwrap(),
270            1
271        );
272        assert_eq!(
273            repo.upstream_oauth_provider()
274                .count(UpstreamOAuthProviderFilter::new().enabled_only())
275                .await
276                .unwrap(),
277            1
278        );
279        assert_eq!(
280            repo.upstream_oauth_provider()
281                .count(UpstreamOAuthProviderFilter::new().disabled_only())
282                .await
283                .unwrap(),
284            0
285        );
286
287        // Disable the provider
288        repo.upstream_oauth_provider()
289            .disable(&clock, provider.clone())
290            .await
291            .unwrap();
292
293        // There should be exactly one disabled provider
294        assert_eq!(
295            repo.upstream_oauth_provider()
296                .count(UpstreamOAuthProviderFilter::new())
297                .await
298                .unwrap(),
299            1
300        );
301        assert_eq!(
302            repo.upstream_oauth_provider()
303                .count(UpstreamOAuthProviderFilter::new().enabled_only())
304                .await
305                .unwrap(),
306            0
307        );
308        assert_eq!(
309            repo.upstream_oauth_provider()
310                .count(UpstreamOAuthProviderFilter::new().disabled_only())
311                .await
312                .unwrap(),
313            1
314        );
315
316        // Test listing and counting sessions
317        let session_filter = UpstreamOAuthSessionFilter::new().for_provider(&provider);
318
319        // Count the sessions for the provider
320        let session_count = repo
321            .upstream_oauth_session()
322            .count(session_filter)
323            .await
324            .unwrap();
325        assert_eq!(session_count, 1);
326
327        // List the sessions for the provider
328        let session_page = repo
329            .upstream_oauth_session()
330            .list(session_filter, Pagination::first(10))
331            .await
332            .unwrap();
333
334        assert_eq!(session_page.edges.len(), 1);
335        assert_eq!(session_page.edges[0].node.id, session.id);
336        assert!(!session_page.has_next_page);
337        assert!(!session_page.has_previous_page);
338
339        // Try deleting the provider
340        repo.upstream_oauth_provider()
341            .delete(provider)
342            .await
343            .unwrap();
344        assert_eq!(
345            repo.upstream_oauth_provider()
346                .count(UpstreamOAuthProviderFilter::new())
347                .await
348                .unwrap(),
349            0
350        );
351    }
352
353    /// Test that the pagination works as expected in the upstream OAuth
354    /// provider repository
355    #[sqlx::test(migrator = "crate::MIGRATOR")]
356    async fn test_provider_repository_pagination(pool: PgPool) {
357        let scope = Scope::from_iter([OPENID]);
358
359        let mut rng = rand_chacha::ChaChaRng::seed_from_u64(42);
360        let clock = MockClock::default();
361        let mut repo = PgRepository::from_pool(&pool).await.unwrap();
362
363        let filter = UpstreamOAuthProviderFilter::new();
364
365        // Count the number of providers before we start
366        assert_eq!(
367            repo.upstream_oauth_provider().count(filter).await.unwrap(),
368            0
369        );
370
371        let mut ids = Vec::with_capacity(20);
372        // Create 20 providers
373        for idx in 0..20 {
374            let client_id = format!("client-{idx}");
375            let provider = repo
376                .upstream_oauth_provider()
377                .add(
378                    &mut rng,
379                    &clock,
380                    UpstreamOAuthProviderParams {
381                        issuer: None,
382                        human_name: None,
383                        brand_name: None,
384                        scope: scope.clone(),
385                        token_endpoint_auth_method: UpstreamOAuthProviderTokenAuthMethod::None,
386                        fetch_userinfo: false,
387                        userinfo_signed_response_alg: None,
388                        token_endpoint_signing_alg: None,
389                        id_token_signed_response_alg: JsonWebSignatureAlg::Rs256,
390                        client_id,
391                        encrypted_client_secret: None,
392                        claims_imports: UpstreamOAuthProviderClaimsImports::default(),
393                        token_endpoint_override: None,
394                        authorization_endpoint_override: None,
395                        userinfo_endpoint_override: None,
396                        jwks_uri_override: None,
397                        discovery_mode: mas_data_model::UpstreamOAuthProviderDiscoveryMode::Oidc,
398                        pkce_mode: mas_data_model::UpstreamOAuthProviderPkceMode::Auto,
399                        response_mode: None,
400                        additional_authorization_parameters: Vec::new(),
401                        forward_login_hint: false,
402                        ui_order: 0,
403                        on_backchannel_logout: UpstreamOAuthProviderOnBackchannelLogout::DoNothing,
404                        registration_token_required: false,
405                    },
406                )
407                .await
408                .unwrap();
409            ids.push(provider.id);
410            clock.advance(Duration::microseconds(10 * 1000 * 1000));
411        }
412
413        // Now we have 20 providers
414        assert_eq!(
415            repo.upstream_oauth_provider().count(filter).await.unwrap(),
416            20
417        );
418
419        // Lookup the first 10 items
420        let page = repo
421            .upstream_oauth_provider()
422            .list(filter, Pagination::first(10))
423            .await
424            .unwrap();
425
426        // It returned the first 10 items
427        assert!(page.has_next_page);
428        let edge_ids: Vec<_> = page.edges.iter().map(|p| p.node.id).collect();
429        assert_eq!(&edge_ids, &ids[..10]);
430
431        // Getting the same page with the "enabled only" filter should return the same
432        // results
433        let other_page = repo
434            .upstream_oauth_provider()
435            .list(filter.enabled_only(), Pagination::first(10))
436            .await
437            .unwrap();
438
439        assert_eq!(page, other_page);
440
441        // Lookup the next 10 items
442        let page = repo
443            .upstream_oauth_provider()
444            .list(filter, Pagination::first(10).after(ids[9]))
445            .await
446            .unwrap();
447
448        // It returned the next 10 items
449        assert!(!page.has_next_page);
450        let edge_ids: Vec<_> = page.edges.iter().map(|p| p.node.id).collect();
451        assert_eq!(&edge_ids, &ids[10..]);
452
453        // Lookup the last 10 items
454        let page = repo
455            .upstream_oauth_provider()
456            .list(filter, Pagination::last(10))
457            .await
458            .unwrap();
459
460        // It returned the last 10 items
461        assert!(page.has_previous_page);
462        let edge_ids: Vec<_> = page.edges.iter().map(|p| p.node.id).collect();
463        assert_eq!(&edge_ids, &ids[10..]);
464
465        // Lookup the previous 10 items
466        let page = repo
467            .upstream_oauth_provider()
468            .list(filter, Pagination::last(10).before(ids[10]))
469            .await
470            .unwrap();
471
472        // It returned the previous 10 items
473        assert!(!page.has_previous_page);
474        let edge_ids: Vec<_> = page.edges.iter().map(|p| p.node.id).collect();
475        assert_eq!(&edge_ids, &ids[..10]);
476
477        // Lookup 10 items between two IDs
478        let page = repo
479            .upstream_oauth_provider()
480            .list(filter, Pagination::first(10).after(ids[5]).before(ids[8]))
481            .await
482            .unwrap();
483
484        // It returned the items in between
485        assert!(!page.has_next_page);
486        let edge_ids: Vec<_> = page.edges.iter().map(|p| p.node.id).collect();
487        assert_eq!(&edge_ids, &ids[6..8]);
488
489        // There should not be any disabled providers
490        assert_eq!(
491            repo.upstream_oauth_provider()
492                .list(
493                    UpstreamOAuthProviderFilter::new().disabled_only(),
494                    Pagination::first(1)
495                )
496                .await
497                .unwrap()
498                .edges,
499            []
500        );
501    }
502
503    /// Test that the pagination works as expected in the upstream OAuth
504    /// session repository
505    #[sqlx::test(migrator = "crate::MIGRATOR")]
506    async fn test_session_repository_pagination(pool: PgPool) {
507        let scope = Scope::from_iter([OPENID]);
508
509        let mut rng = rand_chacha::ChaChaRng::seed_from_u64(42);
510        let clock = MockClock::default();
511        let mut repo = PgRepository::from_pool(&pool).await.unwrap();
512
513        // Create a provider
514        let provider = repo
515            .upstream_oauth_provider()
516            .add(
517                &mut rng,
518                &clock,
519                UpstreamOAuthProviderParams {
520                    issuer: Some("https://example.com/".to_owned()),
521                    human_name: None,
522                    brand_name: None,
523                    scope,
524                    token_endpoint_auth_method: UpstreamOAuthProviderTokenAuthMethod::None,
525                    id_token_signed_response_alg: JsonWebSignatureAlg::Rs256,
526                    fetch_userinfo: false,
527                    userinfo_signed_response_alg: None,
528                    token_endpoint_signing_alg: None,
529                    client_id: "client-id".to_owned(),
530                    encrypted_client_secret: None,
531                    claims_imports: UpstreamOAuthProviderClaimsImports::default(),
532                    token_endpoint_override: None,
533                    authorization_endpoint_override: None,
534                    userinfo_endpoint_override: None,
535                    jwks_uri_override: None,
536                    discovery_mode: mas_data_model::UpstreamOAuthProviderDiscoveryMode::Oidc,
537                    pkce_mode: mas_data_model::UpstreamOAuthProviderPkceMode::Auto,
538                    response_mode: None,
539                    additional_authorization_parameters: Vec::new(),
540                    forward_login_hint: false,
541                    ui_order: 0,
542                    on_backchannel_logout: UpstreamOAuthProviderOnBackchannelLogout::DoNothing,
543                    registration_token_required: false,
544                },
545            )
546            .await
547            .unwrap();
548
549        let filter = UpstreamOAuthSessionFilter::new().for_provider(&provider);
550
551        // Count the number of sessions before we start
552        assert_eq!(
553            repo.upstream_oauth_session().count(filter).await.unwrap(),
554            0
555        );
556
557        let mut links = Vec::with_capacity(3);
558        for subject in ["alice", "bob", "charlie"] {
559            let link = repo
560                .upstream_oauth_link()
561                .add(&mut rng, &clock, &provider, subject.to_owned(), None)
562                .await
563                .unwrap();
564            links.push(link);
565        }
566
567        let mut ids = Vec::with_capacity(20);
568        let sids = ["one", "two"].into_iter().cycle();
569        // Create 20 sessions
570        for (idx, (link, sid)) in links.iter().cycle().zip(sids).enumerate().take(20) {
571            let state = format!("state-{idx}");
572            let session = repo
573                .upstream_oauth_session()
574                .add(&mut rng, &clock, &provider, state, None, None)
575                .await
576                .unwrap();
577            let id_token_claims = serde_json::json!({
578                "sub": link.subject,
579                "sid": sid,
580                "aud": provider.client_id,
581                "iss": "https://example.com/",
582            });
583            let session = repo
584                .upstream_oauth_session()
585                .complete_with_link(
586                    &clock,
587                    session,
588                    link,
589                    None,
590                    Some(id_token_claims),
591                    None,
592                    None,
593                )
594                .await
595                .unwrap();
596            ids.push(session.id);
597            clock.advance(Duration::microseconds(10 * 1000 * 1000));
598        }
599
600        // Now we have 20 sessions
601        assert_eq!(
602            repo.upstream_oauth_session().count(filter).await.unwrap(),
603            20
604        );
605
606        // Lookup the first 10 items
607        let page = repo
608            .upstream_oauth_session()
609            .list(filter, Pagination::first(10))
610            .await
611            .unwrap();
612
613        // It returned the first 10 items
614        assert!(page.has_next_page);
615        let edge_ids: Vec<_> = page.edges.iter().map(|s| s.node.id).collect();
616        assert_eq!(&edge_ids, &ids[..10]);
617
618        // Lookup the next 10 items
619        let page = repo
620            .upstream_oauth_session()
621            .list(filter, Pagination::first(10).after(ids[9]))
622            .await
623            .unwrap();
624
625        // It returned the next 10 items
626        assert!(!page.has_next_page);
627        let edge_ids: Vec<_> = page.edges.iter().map(|s| s.node.id).collect();
628        assert_eq!(&edge_ids, &ids[10..]);
629
630        // Lookup the last 10 items
631        let page = repo
632            .upstream_oauth_session()
633            .list(filter, Pagination::last(10))
634            .await
635            .unwrap();
636
637        // It returned the last 10 items
638        assert!(page.has_previous_page);
639        let edge_ids: Vec<_> = page.edges.iter().map(|s| s.node.id).collect();
640        assert_eq!(&edge_ids, &ids[10..]);
641
642        // Lookup the previous 10 items
643        let page = repo
644            .upstream_oauth_session()
645            .list(filter, Pagination::last(10).before(ids[10]))
646            .await
647            .unwrap();
648
649        // It returned the previous 10 items
650        assert!(!page.has_previous_page);
651        let edge_ids: Vec<_> = page.edges.iter().map(|s| s.node.id).collect();
652        assert_eq!(&edge_ids, &ids[..10]);
653
654        // Lookup 5 items between two IDs
655        let page = repo
656            .upstream_oauth_session()
657            .list(filter, Pagination::first(10).after(ids[5]).before(ids[11]))
658            .await
659            .unwrap();
660
661        // It returned the items in between
662        assert!(!page.has_next_page);
663        let edge_ids: Vec<_> = page.edges.iter().map(|s| s.node.id).collect();
664        assert_eq!(&edge_ids, &ids[6..11]);
665
666        // Check the sub/sid filters
667        assert_eq!(
668            repo.upstream_oauth_session()
669                .count(filter.with_sub_claim("alice").with_sid_claim("one"))
670                .await
671                .unwrap(),
672            4
673        );
674        assert_eq!(
675            repo.upstream_oauth_session()
676                .count(filter.with_sub_claim("bob").with_sid_claim("two"))
677                .await
678                .unwrap(),
679            4
680        );
681
682        let page = repo
683            .upstream_oauth_session()
684            .list(
685                filter.with_sub_claim("alice").with_sid_claim("one"),
686                Pagination::first(10),
687            )
688            .await
689            .unwrap();
690        assert_eq!(page.edges.len(), 4);
691        for edge in page.edges {
692            assert_eq!(
693                edge.node
694                    .id_token_claims()
695                    .unwrap()
696                    .get("sub")
697                    .unwrap()
698                    .as_str(),
699                Some("alice")
700            );
701            assert_eq!(
702                edge.node
703                    .id_token_claims()
704                    .unwrap()
705                    .get("sid")
706                    .unwrap()
707                    .as_str(),
708                Some("one")
709            );
710        }
711    }
712}
mas_storage_pg/upstream_oauth2/mod.rs

mas_storage_pg/upstream_oauth2/
mod.rs
