mas_handlers/upstream_oauth2/
cache.rs

1// Copyright 2024, 2025 New Vector Ltd.
2// Copyright 2023, 2024 The Matrix.org Foundation C.I.C.
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
7use std::{collections::HashMap, sync::Arc};
8
9use mas_context::LogContext;
10use mas_data_model::{
11    UpstreamOAuthProvider, UpstreamOAuthProviderDiscoveryMode, UpstreamOAuthProviderPkceMode,
12};
13use mas_iana::oauth::PkceCodeChallengeMethod;
14use mas_oidc_client::error::DiscoveryError;
15use mas_storage::{RepositoryAccess, upstream_oauth2::UpstreamOAuthProviderRepository};
16use oauth2_types::oidc::VerifiedProviderMetadata;
17use tokio::sync::RwLock;
18use url::Url;
19
20/// A high-level layer over metadata cache and provider configuration, which
21/// resolves endpoint overrides and discovery modes.
22pub struct LazyProviderInfos<'a> {
23    cache: &'a MetadataCache,
24    provider: &'a UpstreamOAuthProvider,
25    client: &'a reqwest::Client,
26    loaded_metadata: Option<Arc<VerifiedProviderMetadata>>,
27}
28
29impl<'a> LazyProviderInfos<'a> {
30    pub fn new(
31        cache: &'a MetadataCache,
32        provider: &'a UpstreamOAuthProvider,
33        client: &'a reqwest::Client,
34    ) -> Self {
35        Self {
36            cache,
37            provider,
38            client,
39            loaded_metadata: None,
40        }
41    }
42
43    /// Trigger the discovery process and return the metadata if discovery is
44    /// enabled.
45    pub async fn maybe_discover(
46        &mut self,
47    ) -> Result<Option<&VerifiedProviderMetadata>, DiscoveryError> {
48        match self.load().await {
49            Ok(metadata) => Ok(Some(metadata)),
50            Err(DiscoveryError::Disabled) => Ok(None),
51            Err(e) => Err(e),
52        }
53    }
54
55    async fn load(&mut self) -> Result<&VerifiedProviderMetadata, DiscoveryError> {
56        if self.loaded_metadata.is_none() {
57            let verify = match self.provider.discovery_mode {
58                UpstreamOAuthProviderDiscoveryMode::Oidc => true,
59                UpstreamOAuthProviderDiscoveryMode::Insecure => false,
60                UpstreamOAuthProviderDiscoveryMode::Disabled => {
61                    return Err(DiscoveryError::Disabled);
62                }
63            };
64
65            let Some(issuer) = &self.provider.issuer else {
66                return Err(DiscoveryError::MissingIssuer);
67            };
68
69            let metadata = self.cache.get(self.client, issuer, verify).await?;
70
71            self.loaded_metadata = Some(metadata);
72        }
73
74        Ok(self.loaded_metadata.as_ref().unwrap())
75    }
76
77    /// Get the JWKS URI for the provider.
78    ///
79    /// Uses [`UpstreamOAuthProvider.jwks_uri_override`] if set, otherwise uses
80    /// the one from discovery.
81    pub async fn jwks_uri(&mut self) -> Result<&Url, DiscoveryError> {
82        if let Some(jwks_uri) = &self.provider.jwks_uri_override {
83            return Ok(jwks_uri);
84        }
85
86        Ok(self.load().await?.jwks_uri())
87    }
88
89    /// Get the authorization endpoint for the provider.
90    ///
91    /// Uses [`UpstreamOAuthProvider.authorization_endpoint_override`] if set,
92    /// otherwise uses the one from discovery.
93    pub async fn authorization_endpoint(&mut self) -> Result<&Url, DiscoveryError> {
94        if let Some(authorization_endpoint) = &self.provider.authorization_endpoint_override {
95            return Ok(authorization_endpoint);
96        }
97
98        Ok(self.load().await?.authorization_endpoint())
99    }
100
101    /// Get the token endpoint for the provider.
102    ///
103    /// Uses [`UpstreamOAuthProvider.token_endpoint_override`] if set, otherwise
104    /// uses the one from discovery.
105    pub async fn token_endpoint(&mut self) -> Result<&Url, DiscoveryError> {
106        if let Some(token_endpoint) = &self.provider.token_endpoint_override {
107            return Ok(token_endpoint);
108        }
109
110        Ok(self.load().await?.token_endpoint())
111    }
112
113    /// Get the userinfo endpoint for the provider.
114    ///
115    /// Uses [`UpstreamOAuthProvider.userinfo_endpoint_override`] if set,
116    /// otherwise uses the one from discovery.
117    pub async fn userinfo_endpoint(&mut self) -> Result<&Url, DiscoveryError> {
118        if let Some(userinfo_endpoint) = &self.provider.userinfo_endpoint_override {
119            return Ok(userinfo_endpoint);
120        }
121
122        Ok(self.load().await?.userinfo_endpoint())
123    }
124
125    /// Get the PKCE methods supported by the provider.
126    ///
127    /// If the mode is set to auto, it will use the ones from discovery,
128    /// defaulting to none if discovery is disabled.
129    pub async fn pkce_methods(
130        &mut self,
131    ) -> Result<Option<Vec<PkceCodeChallengeMethod>>, DiscoveryError> {
132        let methods = match self.provider.pkce_mode {
133            UpstreamOAuthProviderPkceMode::Auto => self
134                .maybe_discover()
135                .await?
136                .and_then(|metadata| metadata.code_challenge_methods_supported.clone()),
137            UpstreamOAuthProviderPkceMode::S256 => Some(vec![PkceCodeChallengeMethod::S256]),
138            UpstreamOAuthProviderPkceMode::Disabled => None,
139        };
140
141        Ok(methods)
142    }
143}
144
145/// A simple OIDC metadata cache
146///
147/// It never evicts entries, does not cache failures and has no locking.
148/// It can also be refreshed in the background, and warmed up on startup.
149/// It is good enough for our use case.
150#[allow(clippy::module_name_repetitions)]
151#[derive(Debug, Clone, Default)]
152pub struct MetadataCache {
153    cache: Arc<RwLock<HashMap<String, Arc<VerifiedProviderMetadata>>>>,
154    insecure_cache: Arc<RwLock<HashMap<String, Arc<VerifiedProviderMetadata>>>>,
155}
156
157impl MetadataCache {
158    #[must_use]
159    pub fn new() -> Self {
160        Self::default()
161    }
162
163    /// Warm up the cache by fetching all the known providers from the database
164    /// and inserting them into the cache.
165    ///
166    /// This spawns a background task that will refresh the cache at the given
167    /// interval.
168    ///
169    /// # Errors
170    ///
171    /// Returns an error if the warm up task could not be started.
172    #[tracing::instrument(name = "metadata_cache.warm_up_and_run", skip_all)]
173    pub async fn warm_up_and_run<R: RepositoryAccess>(
174        &self,
175        client: &reqwest::Client,
176        interval: std::time::Duration,
177        repository: &mut R,
178    ) -> Result<tokio::task::JoinHandle<()>, R::Error> {
179        let providers = repository.upstream_oauth_provider().all_enabled().await?;
180
181        for provider in providers {
182            let verify = match provider.discovery_mode {
183                UpstreamOAuthProviderDiscoveryMode::Oidc => true,
184                UpstreamOAuthProviderDiscoveryMode::Insecure => false,
185                UpstreamOAuthProviderDiscoveryMode::Disabled => continue,
186            };
187
188            let Some(issuer) = &provider.issuer else {
189                tracing::error!(%provider.id, "Provider doesn't have an issuer set, but discovery is enabled!");
190                continue;
191            };
192
193            if let Err(e) = self.fetch(client, issuer, verify).await {
194                tracing::error!(%issuer, error = &e as &dyn std::error::Error, "Failed to fetch provider metadata");
195            }
196        }
197
198        // Spawn a background task to refresh the cache regularly
199        let cache = self.clone();
200        let client = client.clone();
201        Ok(tokio::spawn(async move {
202            loop {
203                // Re-fetch the known metadata at the given interval
204                tokio::time::sleep(interval).await;
205                LogContext::new("metadata-cache-refresh")
206                    .run(|| cache.refresh_all(&client))
207                    .await;
208            }
209        }))
210    }
211
212    #[tracing::instrument(name = "metadata_cache.fetch", fields(%issuer), skip_all)]
213    async fn fetch(
214        &self,
215        client: &reqwest::Client,
216        issuer: &str,
217        verify: bool,
218    ) -> Result<Arc<VerifiedProviderMetadata>, DiscoveryError> {
219        if verify {
220            let metadata = mas_oidc_client::requests::discovery::discover(client, issuer).await?;
221            let metadata = Arc::new(metadata);
222
223            self.cache
224                .write()
225                .await
226                .insert(issuer.to_owned(), metadata.clone());
227
228            Ok(metadata)
229        } else {
230            let metadata =
231                mas_oidc_client::requests::discovery::insecure_discover(client, issuer).await?;
232            let metadata = Arc::new(metadata);
233
234            self.insecure_cache
235                .write()
236                .await
237                .insert(issuer.to_owned(), metadata.clone());
238
239            Ok(metadata)
240        }
241    }
242
243    /// Get the metadata for the given issuer.
244    ///
245    /// # Errors
246    ///
247    /// Returns an error if the metadata could not be retrieved.
248    #[tracing::instrument(name = "metadata_cache.get", fields(%issuer), skip_all)]
249    pub async fn get(
250        &self,
251        client: &reqwest::Client,
252        issuer: &str,
253        verify: bool,
254    ) -> Result<Arc<VerifiedProviderMetadata>, DiscoveryError> {
255        let cache = if verify {
256            self.cache.read().await
257        } else {
258            self.insecure_cache.read().await
259        };
260
261        if let Some(metadata) = cache.get(issuer) {
262            return Ok(Arc::clone(metadata));
263        }
264        // Drop the cache guard so that we don't deadlock when we try to fetch
265        drop(cache);
266
267        let metadata = self.fetch(client, issuer, verify).await?;
268        Ok(metadata)
269    }
270
271    #[tracing::instrument(name = "metadata_cache.refresh_all", skip_all)]
272    async fn refresh_all(&self, client: &reqwest::Client) {
273        // Grab all the keys first to avoid locking the cache for too long
274        let keys: Vec<String> = {
275            let cache = self.cache.read().await;
276            cache.keys().cloned().collect()
277        };
278
279        for issuer in keys {
280            if let Err(e) = self.fetch(client, &issuer, true).await {
281                tracing::error!(issuer = %issuer, error = &e as &dyn std::error::Error, "Failed to refresh provider metadata");
282            }
283        }
284
285        // Do the same for the insecure cache
286        let keys: Vec<String> = {
287            let cache = self.insecure_cache.read().await;
288            cache.keys().cloned().collect()
289        };
290
291        for issuer in keys {
292            if let Err(e) = self.fetch(client, &issuer, false).await {
293                tracing::error!(issuer = %issuer, error = &e as &dyn std::error::Error, "Failed to refresh provider metadata");
294            }
295        }
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    // XXX: sadly, we can't test HTTPS requests with wiremock, so we can only test
302    // 'insecure' discovery
303
304    use mas_data_model::{
305        UpstreamOAuthProviderClaimsImports, UpstreamOAuthProviderOnBackchannelLogout,
306        UpstreamOAuthProviderTokenAuthMethod,
307    };
308    use mas_iana::jose::JsonWebSignatureAlg;
309    use mas_storage::{Clock, clock::MockClock};
310    use oauth2_types::scope::{OPENID, Scope};
311    use ulid::Ulid;
312    use wiremock::{
313        Mock, MockServer, ResponseTemplate,
314        matchers::{method, path},
315    };
316
317    use super::*;
318    use crate::test_utils::setup;
319
320    #[tokio::test]
321    async fn test_metadata_cache() {
322        setup();
323        let mock_server = MockServer::start().await;
324        let http_client = mas_http::reqwest_client();
325
326        let cache = MetadataCache::new();
327
328        // An inexistant issuer should fail
329        cache
330            .get(&http_client, &mock_server.uri(), false)
331            .await
332            .unwrap_err();
333
334        let expected_calls = 3;
335        let mut calls = 0;
336        let _mock_guard = Mock::given(method("GET"))
337            .and(path("/.well-known/openid-configuration"))
338            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
339                "issuer": mock_server.uri(),
340                "authorization_endpoint": "https://example.com/authorize",
341                "token_endpoint": "https://example.com/token",
342                "jwks_uri": "https://example.com/jwks",
343                "userinfo_endpoint": "https://example.com/userinfo",
344                "scopes_supported": ["openid"],
345                "response_types_supported": ["code"],
346                "response_modes_supported": ["query", "fragment"],
347                "grant_types_supported": ["authorization_code"],
348                "subject_types_supported": ["public"],
349                "id_token_signing_alg_values_supported": ["RS256"],
350            })))
351            .expect(expected_calls)
352            .mount(&mock_server)
353            .await;
354
355        // A valid issuer should succeed
356        cache
357            .get(&http_client, &mock_server.uri(), false)
358            .await
359            .unwrap();
360        calls += 1;
361
362        // Calling again should not trigger a new fetch
363        cache
364            .get(&http_client, &mock_server.uri(), false)
365            .await
366            .unwrap();
367        calls += 0;
368
369        // A secure discovery should call but fail because the issuer is insecure
370        cache
371            .get(&http_client, &mock_server.uri(), true)
372            .await
373            .unwrap_err();
374        calls += 1;
375
376        // Calling refresh should refresh all the known issuers
377        cache.refresh_all(&http_client).await;
378        calls += 1;
379
380        assert_eq!(calls, expected_calls);
381    }
382
383    #[tokio::test]
384    async fn test_lazy_provider_infos() {
385        setup();
386
387        let mock_server = MockServer::start().await;
388        let http_client = mas_http::reqwest_client();
389
390        let expected_calls = 2;
391        let mut calls = 0;
392        let _mock_guard = Mock::given(method("GET"))
393            .and(path("/.well-known/openid-configuration"))
394            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
395                "issuer": mock_server.uri(),
396                "authorization_endpoint": "https://example.com/authorize",
397                "token_endpoint": "https://example.com/token",
398                "jwks_uri": "https://example.com/jwks",
399                "userinfo_endpoint": "https://example.com/userinfo",
400                "scopes_supported": ["openid"],
401                "response_types_supported": ["code"],
402                "response_modes_supported": ["query", "fragment"],
403                "grant_types_supported": ["authorization_code"],
404                "subject_types_supported": ["public"],
405                "id_token_signing_alg_values_supported": ["RS256"],
406            })))
407            .expect(expected_calls)
408            .mount(&mock_server)
409            .await;
410
411        let clock = MockClock::default();
412        let provider = UpstreamOAuthProvider {
413            id: Ulid::nil(),
414            issuer: Some(mock_server.uri()),
415            human_name: Some("Example Ltd.".to_owned()),
416            brand_name: None,
417            discovery_mode: UpstreamOAuthProviderDiscoveryMode::Insecure,
418            pkce_mode: UpstreamOAuthProviderPkceMode::Auto,
419            fetch_userinfo: false,
420            userinfo_signed_response_alg: None,
421            jwks_uri_override: None,
422            authorization_endpoint_override: None,
423            scope: Scope::from_iter([OPENID]),
424            userinfo_endpoint_override: None,
425            token_endpoint_override: None,
426            client_id: "client_id".to_owned(),
427            encrypted_client_secret: None,
428            token_endpoint_signing_alg: None,
429            token_endpoint_auth_method: UpstreamOAuthProviderTokenAuthMethod::None,
430            id_token_signed_response_alg: JsonWebSignatureAlg::Rs256,
431            response_mode: None,
432            created_at: clock.now(),
433            disabled_at: None,
434            claims_imports: UpstreamOAuthProviderClaimsImports::default(),
435            additional_authorization_parameters: Vec::new(),
436            forward_login_hint: false,
437            on_backchannel_logout: UpstreamOAuthProviderOnBackchannelLogout::DoNothing,
438        };
439
440        // Without any override, it should just use discovery
441        {
442            let cache = MetadataCache::new();
443            let mut lazy_metadata = LazyProviderInfos::new(&cache, &provider, &http_client);
444            lazy_metadata.maybe_discover().await.unwrap();
445            assert_eq!(
446                lazy_metadata
447                    .authorization_endpoint()
448                    .await
449                    .unwrap()
450                    .as_str(),
451                "https://example.com/authorize"
452            );
453            calls += 1;
454        }
455
456        // Test overriding endpoints
457        {
458            let provider = UpstreamOAuthProvider {
459                jwks_uri_override: Some("https://example.com/jwks_override".parse().unwrap()),
460                authorization_endpoint_override: Some(
461                    "https://example.com/authorize_override".parse().unwrap(),
462                ),
463                token_endpoint_override: Some(
464                    "https://example.com/token_override".parse().unwrap(),
465                ),
466                ..provider.clone()
467            };
468            let cache = MetadataCache::new();
469            let mut lazy_metadata = LazyProviderInfos::new(&cache, &provider, &http_client);
470            assert_eq!(
471                lazy_metadata.jwks_uri().await.unwrap().as_str(),
472                "https://example.com/jwks_override"
473            );
474            assert_eq!(
475                lazy_metadata
476                    .authorization_endpoint()
477                    .await
478                    .unwrap()
479                    .as_str(),
480                "https://example.com/authorize_override"
481            );
482            assert_eq!(
483                lazy_metadata.token_endpoint().await.unwrap().as_str(),
484                "https://example.com/token_override"
485            );
486            // This shouldn't trigger a new fetch as the endpoint is overriden
487            calls += 0;
488        }
489
490        // Loading an insecure provider with secure discovery should fail
491        {
492            let provider = UpstreamOAuthProvider {
493                discovery_mode: UpstreamOAuthProviderDiscoveryMode::Oidc,
494                ..provider.clone()
495            };
496            let cache = MetadataCache::new();
497            let mut lazy_metadata = LazyProviderInfos::new(&cache, &provider, &http_client);
498            lazy_metadata.authorization_endpoint().await.unwrap_err();
499            // This triggered a fetch, even though it failed
500            calls += 1;
501        }
502
503        // Getting endpoints when discovery is disabled only works for overriden ones
504        {
505            let provider = UpstreamOAuthProvider {
506                discovery_mode: UpstreamOAuthProviderDiscoveryMode::Disabled,
507                authorization_endpoint_override: Some(
508                    Url::parse("https://example.com/authorize_override").unwrap(),
509                ),
510                token_endpoint_override: None,
511                ..provider.clone()
512            };
513            let cache = MetadataCache::new();
514            let mut lazy_metadata = LazyProviderInfos::new(&cache, &provider, &http_client);
515            // This should not fail, but also does nothing
516            assert!(lazy_metadata.maybe_discover().await.unwrap().is_none());
517            assert_eq!(
518                lazy_metadata
519                    .authorization_endpoint()
520                    .await
521                    .unwrap()
522                    .as_str(),
523                "https://example.com/authorize_override"
524            );
525            assert!(matches!(
526                lazy_metadata.token_endpoint().await,
527                Err(DiscoveryError::Disabled),
528            ));
529            // This did not trigger a fetch
530            calls += 0;
531        }
532
533        assert_eq!(calls, expected_calls);
534    }
535}
mas_handlers/upstream_oauth2/cache.rs

mas_handlers/upstream_oauth2/
cache.rs
