diff --git a/crates/handlers/src/upstream_oauth2/authorize.rs b/crates/handlers/src/upstream_oauth2/authorize.rs index 72d65e896..5c925cfbf 100644 --- a/crates/handlers/src/upstream_oauth2/authorize.rs +++ b/crates/handlers/src/upstream_oauth2/authorize.rs @@ -28,7 +28,9 @@ use thiserror::Error; use ulid::Ulid; use url::Url; -use super::{UpstreamSessionsCookie, cache::LazyProviderInfos, template::environment}; +use super::{ + UpstreamSessionContext, UpstreamSessionsCookie, cache::LazyProviderInfos, template::environment, +}; use crate::{ impl_from_error_for_route, upstream_oauth2::cache::MetadataCache, views::shared::OptionalPostAuthAction, @@ -77,8 +79,9 @@ pub(crate) enum StartAuthorizationError { /// This discovers the provider metadata if needed, records an /// `upstream_oauth_authorization_sessions` row and stashes it in the browser's /// upstream sessions cookie, along with the action to perform once the user -/// comes back. It returns the authorization URL to redirect the browser to; it -/// is up to the caller to commit the repository. +/// comes back and the context carried from the page which started the flow. It +/// returns the authorization URL to redirect the browser to; it is up to the +/// caller to commit the repository. #[tracing::instrument( name = "handlers.upstream_oauth2.authorize.start", fields(upstream_oauth_provider.id = %provider.id), @@ -94,6 +97,7 @@ pub(crate) async fn start_authorization( cookie_jar: CookieJar, provider: &UpstreamOAuthProvider, post_auth_action: Option, + context: Option, ) -> Result<(CookieJar, Url), StartAuthorizationError> { // Load the session info from the cookie jar. We use this to know whether // the browser recently signed out, which we expose to the @@ -197,7 +201,13 @@ pub(crate) async fn start_authorization( .await?; let cookie_jar = UpstreamSessionsCookie::load(&cookie_jar) - .add(session.id, provider.id, data.state, post_auth_action) + .add( + session.id, + provider.id, + data.state, + post_auth_action, + context, + ) .save(cookie_jar, clock); Ok((cookie_jar, url)) @@ -236,6 +246,7 @@ pub(crate) async fn get( cookie_jar, &provider, query.post_auth_action, + None, ) .await?; diff --git a/crates/handlers/src/upstream_oauth2/cookie.rs b/crates/handlers/src/upstream_oauth2/cookie.rs index cc5ae2943..90b08eb19 100644 --- a/crates/handlers/src/upstream_oauth2/cookie.rs +++ b/crates/handlers/src/upstream_oauth2/cookie.rs @@ -1,3 +1,4 @@ +// Copyright 2025, 2026 Element Creations Ltd. // Copyright 2024, 2025 New Vector Ltd. // Copyright 2022-2024 The Matrix.org Foundation C.I.C. // @@ -20,6 +21,16 @@ static COOKIE_NAME: &str = "upstream-oauth2-sessions"; /// Sessions expire after 10 minutes static SESSION_MAX_TIME: Duration = Duration::microseconds(10 * 60 * 1000 * 1000); +/// Context gathered by the page which started the upstream OAuth 2.0 flow, +/// carried through it so that it can be used on the way back. +#[derive(Serialize, Deserialize, Default, Debug, Clone, PartialEq, Eq)] +pub struct UpstreamSessionContext { + /// The username the user typed on the registration page, used as a + /// candidate localpart if they end up registering an account + #[serde(default, skip_serializing_if = "Option::is_none")] + pub username: Option, +} + #[derive(Serialize, Deserialize, Debug)] pub struct Payload { session: Ulid, @@ -27,6 +38,11 @@ pub struct Payload { state: String, link: Option, post_auth_action: Option, + + /// Cookies saved by older versions of MAS don't have this field, hence the + /// `default` + #[serde(default, skip_serializing_if = "Option::is_none")] + context: Option, } impl Payload { @@ -87,6 +103,7 @@ impl UpstreamSessions { provider: Ulid, state: String, post_auth_action: Option, + context: Option, ) -> Self { self.0.push(Payload { session, @@ -94,6 +111,7 @@ impl UpstreamSessions { state, link: None, post_auth_action, + context, }); self } @@ -131,13 +149,20 @@ impl UpstreamSessions { pub fn lookup_link( &self, link_id: Ulid, - ) -> Result<(Ulid, Option<&PostAuthAction>), UpstreamSessionNotFound> { + ) -> Result< + ( + Ulid, + Option<&PostAuthAction>, + Option<&UpstreamSessionContext>, + ), + UpstreamSessionNotFound, + > { self.0 .iter() .filter(|p| p.link == Some(link_id)) // Find the session with the highest ID, aka. the most recent one .reduce(|a, b| if a.session > b.session { a } else { b }) - .map(|p| (p.session, p.post_auth_action.as_ref())) + .map(|p| (p.session, p.post_auth_action.as_ref(), p.context.as_ref())) .ok_or(UpstreamSessionNotFound) } @@ -178,13 +203,21 @@ mod tests { let first_session = Ulid::from_datetime_with_rng(now, &mut rng); let first_state = "first-state"; - let sessions = sessions.add(first_session, provider_a, first_state.into(), None); + let sessions = sessions.add(first_session, provider_a, first_state.into(), None, None); let now = now + Duration::microseconds(5 * 60 * 1000 * 1000); let second_session = Ulid::from_datetime_with_rng(now, &mut rng); let second_state = "second-state"; - let sessions = sessions.add(second_session, provider_b, second_state.into(), None); + let sessions = sessions.add( + second_session, + provider_b, + second_state.into(), + None, + Some(UpstreamSessionContext { + username: Some("john".to_owned()), + }), + ); let sessions = sessions.expire(now); assert_eq!( @@ -216,11 +249,31 @@ mod tests { // Now the session can't be found with its state assert!(sessions.find_session(provider_b, second_state).is_err()); - // But it can be looked up by its link - assert_eq!(sessions.lookup_link(second_link).unwrap().0, second_session); + // But it can be looked up by its link, along with the context it carries + let (session, _post_auth_action, context) = sessions.lookup_link(second_link).unwrap(); + assert_eq!(session, second_session); + assert_eq!(context.and_then(|c| c.username.as_deref()), Some("john")); + // And it can be consumed let sessions = sessions.consume_link(second_link).unwrap(); // But only once assert!(sessions.consume_link(second_link).is_err()); } + + /// Cookies saved before we started carrying a context should still parse + #[test] + fn test_payload_without_context() { + let sessions: UpstreamSessions = serde_json::from_value(serde_json::json!([{ + "session": "01FSHN9AG0AJ6AC5HQ9X6H4RP4", + "provider": "01FSHN9AG0MZAA6S4AF7CTV32E", + "state": "state", + "link": "01FSHN9AG09NMZYX8MFVH74RP4", + "post_auth_action": null, + }])) + .expect("payload without a context should deserialize"); + + let link = "01FSHN9AG09NMZYX8MFVH74RP4".parse().unwrap(); + let (_session, _post_auth_action, context) = sessions.lookup_link(link).unwrap(); + assert_eq!(context, None); + } } diff --git a/crates/handlers/src/upstream_oauth2/link.rs b/crates/handlers/src/upstream_oauth2/link.rs index 51de8ac5a..a877bcb65 100644 --- a/crates/handlers/src/upstream_oauth2/link.rs +++ b/crates/handlers/src/upstream_oauth2/link.rs @@ -245,10 +245,14 @@ pub(crate) async fn get( ) -> Result { let user_agent = user_agent.map(|ua| ua.as_str().to_owned()); let sessions_cookie = UpstreamSessionsCookie::load(&cookie_jar); - let (session_id, post_auth_action) = sessions_cookie + let (session_id, post_auth_action, context) = sessions_cookie .lookup_link(link_id) .map_err(|_| RouteError::MissingCookie)?; + // The username the user typed on the registration page before starting the + // flow, if any. Cloned out of the cookie so we can consume the cookie later. + let carried_username = context.and_then(|context| context.username.clone()); + let link = repo .upstream_oauth_link() .lookup(link_id) @@ -454,34 +458,50 @@ pub(crate) async fn get( )? }; + let forced_or_required = provider.claims_imports.localpart.is_forced_or_required(); + + // A username carried from the registration page is MAS-side user intent, + // so it takes precedence over the claim rendered from the upstream + // response, and applies even when the claim is ignored. Providers which + // force or require the claim ignore it entirely. + let carried_username = if forced_or_required { + None + } else { + carried_username + }; + // We do a bunch of checks for the localpart. Instead of using nested ifs all // the way, we use a labelled block, and use `break` for 'exiting' early when // needed let localpart = 'localpart: { - if provider.claims_imports.localpart.ignore() { - break 'localpart None; - } + let localpart = if let Some(username) = carried_username { + username + } else { + if provider.claims_imports.localpart.ignore() { + break 'localpart None; + } - let template = provider - .claims_imports - .localpart - .template - .as_deref() - .unwrap_or(DEFAULT_LOCALPART_TEMPLATE); + let template = provider + .claims_imports + .localpart + .template + .as_deref() + .unwrap_or(DEFAULT_LOCALPART_TEMPLATE); - let Some(localpart) = render_attribute_template( - &env, - template, - &context, - provider.claims_imports.localpart.is_required(), - )? - else { - break 'localpart None; + let Some(localpart) = render_attribute_template( + &env, + template, + &context, + provider.claims_imports.localpart.is_required(), + )? + else { + break 'localpart None; + }; + + localpart }; - let forced_or_required = provider.claims_imports.localpart.is_forced_or_required(); - - // We got a localpart from the template. We need to check if it's + // We got a localpart candidate. We need to check if it's // available, and if it's not apply the conflict resolution setup in // the config let maybe_existing_user = repo.user().find_by_username(&localpart).await?; @@ -869,7 +889,9 @@ pub(crate) async fn post( let form = cookie_jar.verify_form(&clock, form)?; let sessions_cookie = UpstreamSessionsCookie::load(&cookie_jar); - let (session_id, post_auth_action) = sessions_cookie + // The context carried through the flow doesn't matter here: the username the + // user picked, if they were allowed to pick one, comes from the form + let (session_id, post_auth_action, _context) = sessions_cookie .lookup_link(link_id) .map_err(|_| RouteError::MissingCookie)?; @@ -1274,7 +1296,10 @@ mod tests { use ulid::Ulid; use super::UpstreamSessionsCookie; - use crate::test_utils::{CookieHelper, RequestBuilderExt, ResponseExt, TestState, setup}; + use crate::{ + test_utils::{CookieHelper, RequestBuilderExt, ResponseExt, TestState, setup}, + upstream_oauth2::UpstreamSessionContext, + }; #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")] async fn test_register(pool: PgPool) { @@ -1399,7 +1424,7 @@ mod tests { let cookie_jar = state.cookie_jar(); let upstream_sessions = UpstreamSessionsCookie::default() - .add(session.id, provider.id, "state".to_owned(), None) + .add(session.id, provider.id, "state".to_owned(), None, None) .add_link_to_session(session.id, link.id) .unwrap(); let cookie_jar = upstream_sessions.save(cookie_jar, &state.clock); @@ -1597,7 +1622,7 @@ mod tests { let cookie_jar = state.cookie_jar(); let upstream_sessions = UpstreamSessionsCookie::default() - .add(session.id, provider.id, "state".to_owned(), None) + .add(session.id, provider.id, "state".to_owned(), None, None) .add_link_to_session(session.id, link.id) .unwrap(); let cookie_jar = upstream_sessions.save(cookie_jar, &state.clock); @@ -1731,7 +1756,7 @@ mod tests { let cookie_jar = state.cookie_jar(); let upstream_sessions = UpstreamSessionsCookie::default() - .add(session.id, provider.id, "state".to_owned(), None) + .add(session.id, provider.id, "state".to_owned(), None, None) .add_link_to_session(session.id, link.id) .unwrap(); let cookie_jar = upstream_sessions.save(cookie_jar, &state.clock); @@ -1855,7 +1880,7 @@ mod tests { let cookie_jar = state.cookie_jar(); let upstream_sessions = UpstreamSessionsCookie::default() - .add(session.id, provider.id, "state".to_owned(), None) + .add(session.id, provider.id, "state".to_owned(), None, None) .add_link_to_session(session.id, link.id) .unwrap(); let cookie_jar = upstream_sessions.save(cookie_jar, &state.clock); @@ -2045,7 +2070,7 @@ mod tests { let cookie_jar = state.cookie_jar(); let upstream_sessions = UpstreamSessionsCookie::default() - .add(session.id, provider.id, "state".to_owned(), None) + .add(session.id, provider.id, "state".to_owned(), None, None) .add_link_to_session(session.id, link.id) .unwrap(); let cookie_jar = upstream_sessions.save(cookie_jar, &state.clock); @@ -2176,7 +2201,7 @@ mod tests { let cookie_jar = state.cookie_jar(); let upstream_sessions = UpstreamSessionsCookie::default() - .add(session.id, provider.id, "state".to_owned(), None) + .add(session.id, provider.id, "state".to_owned(), None, None) .add_link_to_session(session.id, link.id) .unwrap(); let cookie_jar = upstream_sessions.save(cookie_jar, &state.clock); @@ -2314,7 +2339,7 @@ mod tests { let cookie_jar = state.cookie_jar(); let upstream_sessions = UpstreamSessionsCookie::default() - .add(session.id, provider.id, "state".to_owned(), None) + .add(session.id, provider.id, "state".to_owned(), None, None) .add_link_to_session(session.id, link.id) .unwrap(); let cookie_jar = upstream_sessions.save(cookie_jar, &state.clock); @@ -2357,4 +2382,201 @@ mod tests { assert!(old_link_result.is_some(), "Old link should still exist"); assert_eq!(old_link_result.unwrap().user_id, Some(user.id)); } + + /// Provision a provider with the given claims imports, along with a + /// completed upstream authorization session and its link, and load the + /// upstream sessions cookie carrying the given username in the given cookie + /// jar + async fn carried_username_setup( + state: &TestState, + cookies: &CookieHelper, + localpart: mas_data_model::UpstreamOAuthProviderImportAction, + carried_username: &str, + ) -> UpstreamOAuthLink { + let mut rng = state.rng(); + + let claims_imports = UpstreamOAuthProviderClaimsImports { + localpart: UpstreamOAuthProviderLocalpartPreference { + action: localpart, + template: None, + on_conflict: mas_data_model::UpstreamOAuthProviderOnConflict::default(), + }, + ..UpstreamOAuthProviderClaimsImports::default() + }; + + let id_token_claims = serde_json::json!({ + "preferred_username": "john", + }); + let id_token = sign_token(&mut rng, &state.key_store, id_token_claims.clone()).unwrap(); + + let mut repo = state.repository().await.unwrap(); + let provider = repo + .upstream_oauth_provider() + .add( + &mut rng, + &state.clock, + UpstreamOAuthProviderParams { + issuer: Some("https://example.com/".to_owned()), + human_name: Some("Example Ltd.".to_owned()), + brand_name: None, + scope: Scope::from_iter([OPENID]), + token_endpoint_auth_method: UpstreamOAuthProviderTokenAuthMethod::None, + token_endpoint_signing_alg: None, + id_token_signed_response_alg: JsonWebSignatureAlg::Rs256, + client_id: "client".to_owned(), + encrypted_client_secret: None, + claims_imports, + authorization_endpoint_override: None, + token_endpoint_override: None, + userinfo_endpoint_override: None, + fetch_userinfo: false, + userinfo_signed_response_alg: None, + jwks_uri_override: None, + discovery_mode: mas_data_model::UpstreamOAuthProviderDiscoveryMode::Oidc, + pkce_mode: mas_data_model::UpstreamOAuthProviderPkceMode::Auto, + response_mode: None, + additional_authorization_parameters: Vec::new(), + forward_login_hint: false, + ui_order: 0, + on_backchannel_logout: + mas_data_model::UpstreamOAuthProviderOnBackchannelLogout::DoNothing, + registration_token_required: false, + }, + ) + .await + .unwrap(); + + let (link, session) = add_linked_upstream_session( + &mut rng, + &state.clock, + &mut repo, + &provider, + "subject", + &id_token.into_string(), + id_token_claims, + ) + .await + .unwrap(); + + repo.save().await.unwrap(); + + let upstream_sessions = UpstreamSessionsCookie::default() + .add( + session.id, + provider.id, + "state".to_owned(), + None, + Some(UpstreamSessionContext { + username: Some(carried_username.to_owned()), + }), + ) + .add_link_to_session(session.id, link.id) + .unwrap(); + cookies.import(upstream_sessions.save(state.cookie_jar(), &state.clock)); + + link + } + + /// A username carried through the flow prefills the registration form, and + /// takes precedence over the username suggested by the provider + #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")] + async fn test_register_with_carried_username(pool: PgPool) { + setup(); + let state = TestState::from_pool(pool).await.unwrap(); + let cookies = CookieHelper::new(); + + let link = carried_username_setup( + &state, + &cookies, + mas_data_model::UpstreamOAuthProviderImportAction::Suggest, + "alice", + ) + .await; + + let request = Request::get(&*mas_router::UpstreamOAuth2Link::new(link.id).path()).empty(); + let response = state.request(cookies.with_cookies(request)).await; + response.assert_status(StatusCode::OK); + assert!(response.body().contains(r#"value="alice""#)); + assert!(!response.body().contains(r#"value="john""#)); + } + + /// A carried username prefills the form even when the provider is + /// configured to ignore the localpart claim: it isn't a claim, it is what + /// the user typed + #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")] + async fn test_register_with_carried_username_and_ignored_claim(pool: PgPool) { + setup(); + let state = TestState::from_pool(pool).await.unwrap(); + let cookies = CookieHelper::new(); + + let link = carried_username_setup( + &state, + &cookies, + mas_data_model::UpstreamOAuthProviderImportAction::Ignore, + "alice", + ) + .await; + + let request = Request::get(&*mas_router::UpstreamOAuth2Link::new(link.id).path()).empty(); + let response = state.request(cookies.with_cookies(request)).await; + response.assert_status(StatusCode::OK); + assert!(response.body().contains(r#"value="alice""#)); + } + + /// Providers which force the localpart ignore the carried username + #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")] + async fn test_carried_username_ignored_when_localpart_forced(pool: PgPool) { + setup(); + let state = TestState::from_pool(pool).await.unwrap(); + let cookies = CookieHelper::new(); + + let link = carried_username_setup( + &state, + &cookies, + mas_data_model::UpstreamOAuthProviderImportAction::Force, + "alice", + ) + .await; + + let request = Request::get(&*mas_router::UpstreamOAuth2Link::new(link.id).path()).empty(); + let response = state.request(cookies.with_cookies(request)).await; + response.assert_status(StatusCode::OK); + // The localpart is enforced by the provider, and displayed as a full MXID + assert!(response.body().contains(r#"value="@john:"#)); + assert!(!response.body().contains("alice")); + } + + /// A carried username which is already taken is dropped, and the user gets + /// to pick another one + #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")] + async fn test_carried_username_dropped_when_taken(pool: PgPool) { + setup(); + let state = TestState::from_pool(pool).await.unwrap(); + let cookies = CookieHelper::new(); + + let link = carried_username_setup( + &state, + &cookies, + mas_data_model::UpstreamOAuthProviderImportAction::Suggest, + "alice", + ) + .await; + + let mut rng = state.rng(); + let mut repo = state.repository().await.unwrap(); + repo.user() + .add(&mut rng, &state.clock, "alice".to_owned()) + .await + .unwrap(); + repo.save().await.unwrap(); + + let request = Request::get(&*mas_router::UpstreamOAuth2Link::new(link.id).path()).empty(); + let response = state.request(cookies.with_cookies(request)).await; + response.assert_status(StatusCode::OK); + // The username field is left empty: we don't fall back to the username + // suggested by the provider + assert!(response.body().contains(r#"value="""#)); + assert!(!response.body().contains(r#"value="alice""#)); + assert!(!response.body().contains(r#"value="john""#)); + } } diff --git a/crates/handlers/src/upstream_oauth2/mod.rs b/crates/handlers/src/upstream_oauth2/mod.rs index dab366094..f8adcb938 100644 --- a/crates/handlers/src/upstream_oauth2/mod.rs +++ b/crates/handlers/src/upstream_oauth2/mod.rs @@ -1,3 +1,4 @@ +// Copyright 2025, 2026 Element Creations Ltd. // Copyright 2024, 2025 New Vector Ltd. // Copyright 2022-2024 The Matrix.org Foundation C.I.C. // @@ -23,6 +24,7 @@ mod cookie; pub(crate) mod link; mod template; +pub(crate) use self::cookie::UpstreamSessionContext; use self::cookie::UpstreamSessions as UpstreamSessionsCookie; #[derive(Debug, Error)]