dandri/matrix-authentication-service
1
0
Fork 0
mirror of https://github.com/element-hq/matrix-authentication-service.git synced 2026-08-14 15:40:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Carry the username typed on the registration page through the upstream OAuth 2.0 flow

Browse Source
This commit is contained in:
Quentin Gliech
2026-08-11 17:44:22 +02:00
parent b9552374df
commit d762d552ff
4 changed files with 328 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/handlers/src/upstream_oauth2/authorize.rs
+15 -4
View File
@@ -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<PostAuthAction>,
context: Option<UpstreamSessionContext>,
) -> 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?;
crates/handlers/src/upstream_oauth2/cookie.rs
+59 -6
View File
@@ -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<String>,
}
#[derive(Serialize, Deserialize, Debug)]
pub struct Payload {
session: Ulid,
@@ -27,6 +38,11 @@ pub struct Payload {
state: String,
link: Option<Ulid>,
post_auth_action: Option<PostAuthAction>,
/// 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<UpstreamSessionContext>,
}
impl Payload {
@@ -87,6 +103,7 @@ impl UpstreamSessions {
provider: Ulid,
state: String,
post_auth_action: Option<PostAuthAction>,
context: Option<UpstreamSessionContext>,
) -> 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);
}
}
crates/handlers/src/upstream_oauth2/link.rs
+252 -30
View File
@@ -245,10 +245,14 @@ pub(crate) async fn get(
) -> Result<impl IntoResponse, RouteError> {
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""#));
}
}
crates/handlers/src/upstream_oauth2/mod.rs
+2
View File
@@ -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)]