feat: Refactor UIAA service, add support for email stage

src/api/client/account.rs

@@ -1,9 +1,9 @@
-use std::fmt::Write;
+use std::{collections::HashMap, fmt::Write};
 
 use axum::extract::State;
 use axum_client_ip::InsecureClientIp;
 use conduwuit::{
-	Err, Error, Event, Result, debug_info, err, error, info,
+	Err, Event, Result, debug_info, err, error, info,
 	matrix::pdu::PduBuilder,
 	utils::{self, ReadyExt, stream::BroadbandExt},
 	warn,
@@ -12,7 +12,7 @@ use conduwuit_service::Services;
 use futures::{FutureExt, StreamExt};
 use register::RegistrationKind;
 use ruma::{
-	OwnedRoomId, UserId,
+	OwnedRoomId, OwnedUserId, UserId,
 	api::client::{
 		account::{
 			ThirdPartyIdRemovalStatus, change_password, check_registration_token_validity,
@@ -21,7 +21,7 @@ use ruma::{
 			request_3pid_management_token_via_email, request_3pid_management_token_via_msisdn,
 			whoami,
 		},
-		uiaa::{AuthFlow, AuthType, UiaaInfo},
+		uiaa::{AuthFlow, AuthType},
 	},
 	events::{
 		GlobalAccountDataEventType, StateEventType,
@@ -33,8 +33,10 @@ use ruma::{
 	},
 	push,
 };
+use serde_json::value::RawValue;
+use service::uiaa::Identity;
 
-use super::{DEVICE_ID_LENGTH, SESSION_ID_LENGTH, TOKEN_LENGTH, join_room_by_id_helper};
+use super::{DEVICE_ID_LENGTH, TOKEN_LENGTH, join_room_by_id_helper};
 use crate::Ruma;
 
 const RANDOM_USER_ID_LENGTH: usize = 10;
@@ -201,7 +203,7 @@ pub(crate) async fn register_route(
 
 	// forbid guests from registering if there is not a real admin user yet. give
 	// generic user error.
-	if is_guest && services.users.count().await < 2 {
+	if is_guest && services.firstrun.is_first_run() {
 		warn!(
 			"Guest account attempted to register before a real admin user has been registered, \
 			 rejecting registration. Guest's initial device name: \"{}\"",
@@ -212,81 +214,19 @@ pub(crate) async fn register_route(
 		)));
 	}
 
-	let user_id = match (body.username.as_ref(), is_guest) {
-		| (Some(username), false) => {
-			// workaround for https://github.com/matrix-org/matrix-appservice-irc/issues/1780 due to inactivity of fixing the issue
-			let is_matrix_appservice_irc =
-				body.appservice_info.as_ref().is_some_and(|appservice| {
-					appservice.registration.id == "irc"
-						|| appservice.registration.id.contains("matrix-appservice-irc")
-						|| appservice.registration.id.contains("matrix_appservice_irc")
-				});
-
-			if services.globals.forbidden_usernames().is_match(username)
-				&& !emergency_mode_enabled
-			{
-				return Err!(Request(Forbidden("Username is forbidden")));
-			}
-
-			// don't force the username lowercase if it's from matrix-appservice-irc
-			let body_username = if is_matrix_appservice_irc {
-				username.clone()
-			} else {
-				username.to_lowercase()
-			};
-
-			let proposed_user_id = match UserId::parse_with_server_name(
-				&body_username,
-				services.globals.server_name(),
-			) {
-				| Ok(user_id) => {
-					if let Err(e) = user_id.validate_strict() {
-						// unless the username is from the broken matrix appservice IRC bridge, or
-						// we are in emergency mode, we should follow synapse's behaviour on
-						// not allowing things like spaces and UTF-8 characters in usernames
-						if !is_matrix_appservice_irc && !emergency_mode_enabled {
-							return Err!(Request(InvalidUsername(debug_warn!(
-								"Username {body_username} contains disallowed characters or \
-								 spaces: {e}"
-							))));
-						}
-					}
-
-					// Don't allow registration with user IDs that aren't local
-					if !services.globals.user_is_local(&user_id) {
-						return Err!(Request(InvalidUsername(
-							"Username {body_username} is not local to this server"
-						)));
-					}
-
-					user_id
-				},
-				| Err(e) => {
-					return Err!(Request(InvalidUsername(debug_warn!(
-						"Username {body_username} is not valid: {e}"
-					))));
-				},
-			};
-
-			if services.users.exists(&proposed_user_id).await {
-				return Err!(Request(UserInUse("User ID is not available.")));
-			}
-
-			proposed_user_id
-		},
-		| _ => loop {
-			let proposed_user_id = UserId::parse_with_server_name(
-				utils::random_string(RANDOM_USER_ID_LENGTH).to_lowercase(),
-				services.globals.server_name(),
-			)
-			.unwrap();
-			if !services.users.exists(&proposed_user_id).await {
-				break proposed_user_id;
-			}
-		},
-	};
+	let user_id = determine_registration_user_id(
+		&services,
+		body.username.clone(),
+		body.appservice_info.as_ref(),
+		is_guest,
+		emergency_mode_enabled,
+	)
+	.await?;
 
 	if body.body.login_type == Some(LoginType::ApplicationService) {
+		// For appservice logins, make sure that the user ID is in the appservice's
+		// namespace
+
 		match body.appservice_info {
 			| Some(ref info) =>
 				if !info.is_user_match(&user_id) && !emergency_mode_enabled {
@@ -300,126 +240,48 @@ pub(crate) async fn register_route(
 		}
 	} else if services.appservice.is_exclusive_user_id(&user_id).await && !emergency_mode_enabled
 	{
+		// For non-appservice logins, ban user IDs which are in an appservice's
+		// namespace (unless emergency mode is enabled)
 		return Err!(Request(Exclusive("Username is reserved by an appservice.")));
 	}
 
-	// UIAA
-	let mut uiaainfo = UiaaInfo {
-		flows: Vec::new(),
-		completed: Vec::new(),
-		params: Box::default(),
-		session: None,
-		auth_error: None,
-	};
+	// Appeservices and guests get to skip auth
 	let skip_auth = body.appservice_info.is_some() || is_guest;
 
-	// Populate required UIAA flows
-
-	if services.firstrun.is_first_run() {
-		// Registration token forced while in first-run mode
-		uiaainfo.flows.push(AuthFlow {
-			stages: vec![AuthType::RegistrationToken],
-		});
+	let identity = if skip_auth {
+		// Appservices and guests have no identity
+		None
 	} else {
-		if services
-			.registration_tokens
-			.iterate_tokens()
-			.next()
-			.await
-			.is_some()
-		{
-			// Registration token required
-			uiaainfo.flows.push(AuthFlow {
-				stages: vec![AuthType::RegistrationToken],
-			});
-		}
+		// Perform UIAA to determine the user's identity
+		let (flows, params) = create_registration_uiaa_session(&services).await?;
 
-		if services.config.recaptcha_private_site_key.is_some() {
-			if let Some(pubkey) = &services.config.recaptcha_site_key {
-				// ReCaptcha required
-				uiaainfo
-					.flows
-					.push(AuthFlow { stages: vec![AuthType::ReCaptcha] });
-				uiaainfo.params = serde_json::value::to_raw_value(&serde_json::json!({
-					"m.login.recaptcha": {
-						"public_key": pubkey,
-					},
-				}))
-				.expect("Failed to serialize recaptcha params");
-			}
-		}
-
-		if uiaainfo.flows.is_empty() && !skip_auth {
-			// Registration isn't _disabled_, but there's no captcha configured and no
-			// registration tokens currently set. Bail out by default unless open
-			// registration was explicitly enabled.
-			if !services
-				.config
-				.yes_i_am_very_very_sure_i_want_an_open_registration_server_prone_to_abuse
-			{
-				return Err!(Request(Forbidden(
-					"This server is not accepting registrations at this time."
-				)));
-			}
-
-			// We have open registration enabled (😧), provide a dummy stage
-			uiaainfo = UiaaInfo {
-				flows: vec![AuthFlow { stages: vec![AuthType::Dummy] }],
-				completed: Vec::new(),
-				params: Box::default(),
-				session: None,
-				auth_error: None,
-			};
-		}
-	}
-
-	if !skip_auth {
-		match &body.auth {
-			| Some(auth) => {
-				let (worked, uiaainfo) = services
-					.uiaa
-					.try_auth(
-						&UserId::parse_with_server_name("", services.globals.server_name())
-							.unwrap(),
-						"".into(),
-						auth,
-						&uiaainfo,
-					)
-					.await?;
-				if !worked {
-					return Err(Error::Uiaa(uiaainfo));
-				}
-				// Success!
-			},
-			| _ => match body.json_body {
-				| Some(ref json) => {
-					uiaainfo.session = Some(utils::random_string(SESSION_ID_LENGTH));
-					services.uiaa.create(
-						&UserId::parse_with_server_name("", services.globals.server_name())
-							.unwrap(),
-						"".into(),
-						&uiaainfo,
-						json,
-					);
-					return Err(Error::Uiaa(uiaainfo));
-				},
-				| _ => {
-					return Err!(Request(NotJson("JSON body is not valid")));
-				},
-			},
-		}
-	}
+		Some(
+			services
+				.uiaa
+				.authenticate(&body.auth, flows, params, None)
+				.await?,
+		)
+	};
 
 	let password = if is_guest { None } else { body.password.as_deref() };
 
 	// Create user
 	services.users.create(&user_id, password, None).await?;
 
-	// Default to pretty displayname
+	// If the user registered with an email, associate it with their account
+	#[allow(clippy::collapsible_if)]
+	if let Some(identity) = identity {
+		if let Some(email) = identity.email {
+			services
+				.threepid
+				.associate_localpart_email(user_id.localpart(), email);
+		}
+	}
+
+	// Set an initial display name
 	let mut displayname = user_id.localpart().to_owned();
 
-	// If `new_user_displayname_suffix` is set, registration will push whatever
-	// content is set to the user's display name with a space before it
+	// Apply the new user displayname suffix, if it's set
 	if !services.globals.new_user_displayname_suffix().is_empty()
 		&& body.appservice_info.is_none()
 	{
@@ -615,6 +477,157 @@ pub(crate) async fn register_route(
 	})
 }
 
+/// Determine which flows and parameters should be presented when
+/// registering a new account.
+async fn create_registration_uiaa_session(
+	services: &Services,
+) -> Result<(Vec<AuthFlow>, Box<RawValue>)> {
+	let mut flows = vec![];
+	let mut params = HashMap::<String, serde_json::Value>::new();
+
+	if services.firstrun.is_first_run() {
+		// Registration token forced while in first-run mode
+		flows.push(AuthFlow::new(vec![AuthType::RegistrationToken]));
+	} else {
+		if services
+			.registration_tokens
+			.iterate_tokens()
+			.next()
+			.await
+			.is_some()
+		{
+			// Registration token flow is available
+			flows.push(AuthFlow::new(vec![AuthType::RegistrationToken]));
+		}
+
+		if services.config.recaptcha_private_site_key.is_some() {
+			if let Some(pubkey) = &services.config.recaptcha_site_key {
+				// ReCaptcha flow is available
+				flows.push(AuthFlow::new(vec![AuthType::ReCaptcha]));
+
+				params.insert(
+					AuthType::ReCaptcha.as_str().to_owned(),
+					serde_json::json!({
+						"public_key": pubkey,
+					}),
+				);
+			}
+		}
+	}
+
+	if flows.is_empty() {
+		// Registration is enabled, but no flows are configured. Bail out by default
+		// unless open registration was explicitly enabled.
+		if !services
+			.config
+			.yes_i_am_very_very_sure_i_want_an_open_registration_server_prone_to_abuse
+		{
+			return Err!(Request(Forbidden(
+				"This server is not accepting registrations at this time."
+			)));
+		}
+
+		// We have open registration enabled (😧), provide a dummy flow
+		flows.push(AuthFlow { stages: vec![AuthType::Dummy] });
+	}
+
+	let params = serde_json::value::to_raw_value(&params).expect("params should be valid JSON");
+
+	Ok((flows, params))
+}
+
+async fn determine_registration_user_id(
+	services: &Services,
+	supplied_username: Option<String>,
+	appservice_info: Option<&service::appservice::RegistrationInfo>,
+	is_guest: bool,
+	emergency_mode_enabled: bool,
+) -> Result<OwnedUserId> {
+	if let Some(mut supplied_username) = supplied_username
+		&& !is_guest
+	{
+		// The user gets to pick their username. Do some validation to make sure it's
+		// acceptable.
+
+		// Don't allow registration with forbidden usernames.
+		if services
+			.globals
+			.forbidden_usernames()
+			.is_match(&supplied_username)
+			&& !emergency_mode_enabled
+		{
+			return Err!(Request(Forbidden("Username is forbidden")));
+		}
+
+		// Workaround for https://github.com/matrix-org/matrix-appservice-irc/issues/1780 due to inactivity of fixing the issue
+		let is_matrix_appservice_irc = appservice_info.is_some_and(|appservice| {
+			appservice.registration.id == "irc"
+				|| appservice.registration.id.contains("matrix-appservice-irc")
+				|| appservice.registration.id.contains("matrix_appservice_irc")
+		});
+
+		// Don't force the username lowercase if it's from matrix-appservice-irc.
+		if !is_matrix_appservice_irc {
+			supplied_username = supplied_username.to_lowercase();
+		}
+
+		// Create and validate the user ID
+		let user_id = match UserId::parse_with_server_name(
+			&supplied_username,
+			services.globals.server_name(),
+		) {
+			| Ok(user_id) => {
+				if let Err(e) = user_id.validate_strict() {
+					// unless the username is from the broken matrix appservice IRC bridge, or
+					// we are in emergency mode, we should follow synapse's behaviour on
+					// not allowing things like spaces and UTF-8 characters in usernames
+					if !is_matrix_appservice_irc && !emergency_mode_enabled {
+						return Err!(Request(InvalidUsername(debug_warn!(
+							"Username {supplied_username} contains disallowed characters or \
+							 spaces: {e}"
+						))));
+					}
+				}
+
+				// Don't allow registration with user IDs that aren't local
+				if !services.globals.user_is_local(&user_id) {
+					return Err!(Request(InvalidUsername(
+						"Username {supplied_username} is not local to this server"
+					)));
+				}
+
+				user_id
+			},
+			| Err(e) => {
+				return Err!(Request(InvalidUsername(debug_warn!(
+					"Username {supplied_username} is not valid: {e}"
+				))));
+			},
+		};
+
+		if services.users.exists(&user_id).await {
+			return Err!(Request(UserInUse("User ID is not available.")));
+		}
+
+		Ok(user_id)
+	} else {
+		// The user is a guest or is lacking in creativity. Generate a username for
+		// them.
+
+		loop {
+			let user_id = UserId::parse_with_server_name(
+				utils::random_string(RANDOM_USER_ID_LENGTH).to_lowercase(),
+				services.globals.server_name(),
+			)
+			.unwrap();
+
+			if !services.users.exists(&user_id).await {
+				break Ok(user_id);
+			}
+		}
+	}
+}
+
 /// # `POST /_matrix/client/r0/account/password`
 ///
 /// Changes the password of this account.
@@ -638,67 +651,61 @@ pub(crate) async fn change_password_route(
 	InsecureClientIp(client): InsecureClientIp,
 	body: Ruma<change_password::v3::Request>,
 ) -> Result<change_password::v3::Response> {
-	// Authentication for this endpoint was made optional, but we need
-	// authentication currently
-	let sender_user = body
-		.sender_user
-		.as_ref()
-		.ok_or_else(|| err!(Request(MissingToken("Missing access token."))))?;
+	let identity = if let Some(ref user_id) = body.sender_user {
+		// A signed-in user is trying to change their password, prompt them for their
+		// existing one
 
-	let mut uiaainfo = UiaaInfo {
-		flows: vec![AuthFlow { stages: vec![AuthType::Password] }],
-		completed: Vec::new(),
-		params: Box::default(),
-		session: None,
-		auth_error: None,
+		services
+			.uiaa
+			.authenticate(
+				&body.auth,
+				vec![AuthFlow::new(vec![AuthType::Password])],
+				Box::default(),
+				Some(Identity::from_user_id(user_id)),
+			)
+			.await?
+	} else {
+		// A signed-out user is trying to reset their password, prompt them for email
+		// confirmation Note that we do not _send_ an email here, their client should
+		// have already hit `/account/password/requestToken` to send the email. We
+		// just validate it.
+
+		services
+			.uiaa
+			.authenticate(
+				&body.auth,
+				vec![AuthFlow::new(vec![AuthType::EmailIdentity])],
+				Box::default(),
+				None,
+			)
+			.await?
 	};
 
-	match &body.auth {
-		| Some(auth) => {
-			let (worked, uiaainfo) = services
-				.uiaa
-				.try_auth(sender_user, body.sender_device(), auth, &uiaainfo)
-				.await?;
-
-			if !worked {
-				return Err(Error::Uiaa(uiaainfo));
-			}
-
-			// Success!
-		},
-		| _ => match body.json_body {
-			| Some(ref json) => {
-				uiaainfo.session = Some(utils::random_string(SESSION_ID_LENGTH));
-				services
-					.uiaa
-					.create(sender_user, body.sender_device(), &uiaainfo, json);
-
-				return Err(Error::Uiaa(uiaainfo));
-			},
-			| _ => {
-				return Err!(Request(NotJson("JSON body is not valid")));
-			},
-		},
-	}
+	let sender_user = OwnedUserId::parse(format!(
+		"@{}:{}",
+		identity.localpart.expect("localpart should be known"),
+		services.globals.server_name()
+	))
+	.expect("user ID should be valid");
 
 	services
 		.users
-		.set_password(sender_user, Some(&body.new_password))
+		.set_password(&sender_user, Some(&body.new_password))
 		.await?;
 
 	if body.logout_devices {
 		// Logout all devices except the current one
 		services
 			.users
-			.all_device_ids(sender_user)
+			.all_device_ids(&sender_user)
 			.ready_filter(|id| *id != body.sender_device())
-			.for_each(|id| services.users.remove_device(sender_user, id))
+			.for_each(|id| services.users.remove_device(&sender_user, id))
 			.await;
 
 		// Remove all pushers except the ones associated with this session
 		services
 			.pusher
-			.get_pushkeys(sender_user)
+			.get_pushkeys(&sender_user)
 			.map(ToOwned::to_owned)
 			.broad_filter_map(async |pushkey| {
 				services
@@ -711,17 +718,17 @@ pub(crate) async fn change_password_route(
 					.then_some(pushkey)
 			})
 			.for_each(async |pushkey| {
-				services.pusher.delete_pusher(sender_user, &pushkey).await;
+				services.pusher.delete_pusher(&sender_user, &pushkey).await;
 			})
 			.await;
 	}
 
-	info!("User {sender_user} changed their password.");
+	info!("User {} changed their password.", &sender_user);
 
 	if services.server.config.admin_room_notices {
 		services
 			.admin
-			.notice(&format!("User {sender_user} changed their password."))
+			.notice(&format!("User {} changed their password.", &sender_user))
 			.await;
 	}
 
@@ -768,47 +775,18 @@ pub(crate) async fn deactivate_route(
 	InsecureClientIp(client): InsecureClientIp,
 	body: Ruma<deactivate::v3::Request>,
 ) -> Result<deactivate::v3::Response> {
-	// Authentication for this endpoint was made optional, but we need
-	// authentication currently
+	// Authentication for this endpoint is technically optional,
+	// but we require the user to be logged in
 	let sender_user = body
 		.sender_user
 		.as_ref()
 		.ok_or_else(|| err!(Request(MissingToken("Missing access token."))))?;
 
-	let mut uiaainfo = UiaaInfo {
-		flows: vec![AuthFlow { stages: vec![AuthType::Password] }],
-		completed: Vec::new(),
-		params: Box::default(),
-		session: None,
-		auth_error: None,
-	};
-
-	match &body.auth {
-		| Some(auth) => {
-			let (worked, uiaainfo) = services
-				.uiaa
-				.try_auth(sender_user, body.sender_device(), auth, &uiaainfo)
-				.await?;
-
-			if !worked {
-				return Err(Error::Uiaa(uiaainfo));
-			}
-			// Success!
-		},
-		| _ => match body.json_body {
-			| Some(ref json) => {
-				uiaainfo.session = Some(utils::random_string(SESSION_ID_LENGTH));
-				services
-					.uiaa
-					.create(sender_user, body.sender_device(), &uiaainfo, json);
-
-				return Err(Error::Uiaa(uiaainfo));
-			},
-			| _ => {
-				return Err!(Request(NotJson("JSON body is not valid")));
-			},
-		},
-	}
+	// Prompt the user to confirm with their password using UIAA
+	let _ = services
+		.uiaa
+		.authenticate_password(&body.auth, Some(Identity::from_user_id(sender_user)))
+		.await?;
 
 	// Remove profile pictures and display name
 	let all_joined_rooms: Vec<OwnedRoomId> = services

src/api/client/device.rs

@@ -1,17 +1,15 @@
 use axum::extract::State;
 use axum_client_ip::InsecureClientIp;
-use conduwuit::{Err, Error, Result, debug, err, utils};
+use conduwuit::{Err, Result, debug, err, utils};
 use futures::StreamExt;
 use ruma::{
 	MilliSecondsSinceUnixEpoch, OwnedDeviceId,
-	api::client::{
-		device::{self, delete_device, delete_devices, get_device, get_devices, update_device},
-		error::ErrorKind,
-		uiaa::{AuthFlow, AuthType, UiaaInfo},
+	api::client::device::{
+		self, delete_device, delete_devices, get_device, get_devices, update_device,
 	},
 };
+use service::uiaa::Identity;
 
-use super::SESSION_ID_LENGTH;
 use crate::{Ruma, client::DEVICE_ID_LENGTH};
 
 /// # `GET /_matrix/client/r0/devices`
@@ -123,7 +121,7 @@ pub(crate) async fn delete_device_route(
 	State(services): State<crate::State>,
 	body: Ruma<delete_device::v3::Request>,
 ) -> Result<delete_device::v3::Response> {
-	let (sender_user, sender_device) = body.sender();
+	let sender_user = body.sender_user();
 	let appservice = body.appservice_info.as_ref();
 
 	if appservice.is_some_and(|appservice| appservice.registration.device_management) {
@@ -139,41 +137,11 @@ pub(crate) async fn delete_device_route(
 		return Ok(delete_device::v3::Response {});
 	}
 
-	// UIAA
-	let mut uiaainfo = UiaaInfo {
-		flows: vec![AuthFlow { stages: vec![AuthType::Password] }],
-		completed: Vec::new(),
-		params: Box::default(),
-		session: None,
-		auth_error: None,
-	};
-
-	match &body.auth {
-		| Some(auth) => {
-			let (worked, uiaainfo) = services
-				.uiaa
-				.try_auth(sender_user, sender_device, auth, &uiaainfo)
-				.await?;
-
-			if !worked {
-				return Err!(Uiaa(uiaainfo));
-			}
-			// Success!
-		},
-		| _ => match body.json_body {
-			| Some(ref json) => {
-				uiaainfo.session = Some(utils::random_string(SESSION_ID_LENGTH));
-				services
-					.uiaa
-					.create(sender_user, sender_device, &uiaainfo, json);
-
-				return Err!(Uiaa(uiaainfo));
-			},
-			| _ => {
-				return Err!(Request(NotJson("Not json.")));
-			},
-		},
-	}
+	// Prompt the user to confirm with their password using UIAA
+	let _ = services
+		.uiaa
+		.authenticate_password(&body.auth, Some(Identity::from_user_id(sender_user)))
+		.await?;
 
 	services
 		.users
@@ -200,7 +168,7 @@ pub(crate) async fn delete_devices_route(
 	State(services): State<crate::State>,
 	body: Ruma<delete_devices::v3::Request>,
 ) -> Result<delete_devices::v3::Response> {
-	let (sender_user, sender_device) = body.sender();
+	let sender_user = body.sender_user();
 	let appservice = body.appservice_info.as_ref();
 
 	if appservice.is_some_and(|appservice| appservice.registration.device_management) {
@@ -215,41 +183,11 @@ pub(crate) async fn delete_devices_route(
 		return Ok(delete_devices::v3::Response {});
 	}
 
-	// UIAA
-	let mut uiaainfo = UiaaInfo {
-		flows: vec![AuthFlow { stages: vec![AuthType::Password] }],
-		completed: Vec::new(),
-		params: Box::default(),
-		session: None,
-		auth_error: None,
-	};
-
-	match &body.auth {
-		| Some(auth) => {
-			let (worked, uiaainfo) = services
-				.uiaa
-				.try_auth(sender_user, sender_device, auth, &uiaainfo)
-				.await?;
-
-			if !worked {
-				return Err(Error::Uiaa(uiaainfo));
-			}
-			// Success!
-		},
-		| _ => match body.json_body {
-			| Some(ref json) => {
-				uiaainfo.session = Some(utils::random_string(SESSION_ID_LENGTH));
-				services
-					.uiaa
-					.create(sender_user, sender_device, &uiaainfo, json);
-
-				return Err(Error::Uiaa(uiaainfo));
-			},
-			| _ => {
-				return Err(Error::BadRequest(ErrorKind::NotJson, "Not json."));
-			},
-		},
-	}
+	// Prompt the user to confirm with their password using UIAA
+	let _ = services
+		.uiaa
+		.authenticate_password(&body.auth, Some(Identity::from_user_id(sender_user)))
+		.await?;
 
 	for device_id in &body.devices {
 		services.users.remove_device(sender_user, device_id).await;

src/api/client/keys.rs

@@ -7,7 +7,6 @@ use axum::extract::State;
 use conduwuit::{
 	Err, Error, Result, debug, debug_warn, err,
 	result::NotFound,
-	utils,
 	utils::{IterStream, stream::WidebandExt},
 };
 use conduwuit_service::{Services, users::parse_master_key};
@@ -22,7 +21,6 @@ use ruma::{
 				upload_signatures::{self},
 				upload_signing_keys,
 			},
-			uiaa::{AuthFlow, AuthType, UiaaInfo},
 		},
 		federation,
 	},
@@ -30,8 +28,8 @@ use ruma::{
 	serde::Raw,
 };
 use serde_json::json;
+use service::uiaa::Identity;
 
-use super::SESSION_ID_LENGTH;
 use crate::Ruma;
 
 /// # `POST /_matrix/client/r0/keys/upload`
@@ -174,16 +172,7 @@ pub(crate) async fn upload_signing_keys_route(
 	State(services): State<crate::State>,
 	body: Ruma<upload_signing_keys::v3::Request>,
 ) -> Result<upload_signing_keys::v3::Response> {
-	let (sender_user, sender_device) = body.sender();
-
-	// UIAA
-	let mut uiaainfo = UiaaInfo {
-		flows: vec![AuthFlow { stages: vec![AuthType::Password] }],
-		completed: Vec::new(),
-		params: Box::default(),
-		session: None,
-		auth_error: None,
-	};
+	let sender_user = body.sender_user();
 
 	match check_for_new_keys(
 		services,
@@ -207,32 +196,10 @@ pub(crate) async fn upload_signing_keys_route(
 			// Some of the keys weren't found, so we let them upload
 		},
 		| _ => {
-			match &body.auth {
-				| Some(auth) => {
-					let (worked, uiaainfo) = services
-						.uiaa
-						.try_auth(sender_user, sender_device, auth, &uiaainfo)
-						.await?;
-
-					if !worked {
-						return Err(Error::Uiaa(uiaainfo));
-					}
-					// Success!
-				},
-				| _ => match body.json_body.as_ref() {
-					| Some(json) => {
-						uiaainfo.session = Some(utils::random_string(SESSION_ID_LENGTH));
-						services
-							.uiaa
-							.create(sender_user, sender_device, &uiaainfo, json);
-
-						return Err(Error::Uiaa(uiaainfo));
-					},
-					| _ => {
-						return Err(Error::BadRequest(ErrorKind::NotJson, "Not json."));
-					},
-				},
-			}
+			let _ = services
+				.uiaa
+				.authenticate_password(&body.auth, Some(Identity::from_user_id(sender_user)))
+				.await?;
 		},
 	}
 

src/api/client/mod.rs

@@ -92,6 +92,3 @@ const DEVICE_ID_LENGTH: usize = 10;
 
 /// generated user access token length
 const TOKEN_LENGTH: usize = 32;
-
-/// generated user session ID length
-const SESSION_ID_LENGTH: usize = service::uiaa::SESSION_ID_LENGTH;

src/api/client/session.rs

@@ -8,7 +8,7 @@ use conduwuit::{
 	warn,
 };
 use conduwuit_core::{debug_error, debug_warn};
-use conduwuit_service::{Services, uiaa::SESSION_ID_LENGTH};
+use conduwuit_service::Services;
 use futures::StreamExt;
 use ruma::{
 	OwnedUserId, UserId,
@@ -29,6 +29,7 @@ use ruma::{
 		uiaa,
 	},
 };
+use service::uiaa::Identity;
 
 use super::{DEVICE_ID_LENGTH, TOKEN_LENGTH};
 use crate::Ruma;
@@ -370,45 +371,13 @@ pub(crate) async fn login_token_route(
 		return Err!(Request(Forbidden("Login via an existing session is not enabled")));
 	}
 
-	// This route SHOULD have UIA
-	// TODO: How do we make only UIA sessions that have not been used before valid?
-	let (sender_user, sender_device) = body.sender();
+	let sender_user = body.sender_user();
 
-	let mut uiaainfo = uiaa::UiaaInfo {
-		flows: vec![uiaa::AuthFlow { stages: vec![uiaa::AuthType::Password] }],
-		completed: Vec::new(),
-		params: Box::default(),
-		session: None,
-		auth_error: None,
-	};
-
-	match &body.auth {
-		| Some(auth) => {
-			let (worked, uiaainfo) = services
-				.uiaa
-				.try_auth(sender_user, sender_device, auth, &uiaainfo)
-				.await?;
-
-			if !worked {
-				return Err(Error::Uiaa(uiaainfo));
-			}
-
-			// Success!
-		},
-		| _ => match body.json_body.as_ref() {
-			| Some(json) => {
-				uiaainfo.session = Some(utils::random_string(SESSION_ID_LENGTH));
-				services
-					.uiaa
-					.create(sender_user, sender_device, &uiaainfo, json);
-
-				return Err(Error::Uiaa(uiaainfo));
-			},
-			| _ => {
-				return Err!(Request(NotJson("No JSON body was sent when required.")));
-			},
-		},
-	}
+	// Prompt the user to confirm with their password using UIAA
+	let _ = services
+		.uiaa
+		.authenticate_password(&body.auth, Some(Identity::from_user_id(sender_user)))
+		.await?;
 
 	let login_token = utils::random_string(TOKEN_LENGTH);
 	let expires_in = services.users.create_login_token(sender_user, &login_token);

src/api/router/args.rs

@@ -2,14 +2,13 @@ use std::{mem, ops::Deref};
 
 use axum::{body::Body, extract::FromRequest};
 use bytes::{BufMut, Bytes, BytesMut};
-use conduwuit::{Error, Result, debug, debug_warn, err, trace, utils::string::EMPTY};
+use conduwuit::{Error, Result, debug, debug_warn, err, trace};
 use ruma::{
 	CanonicalJsonObject, CanonicalJsonValue, DeviceId, OwnedDeviceId, OwnedServerName,
 	OwnedUserId, ServerName, UserId, api::IncomingRequest,
 };
-use service::Services;
 
-use super::{auth, auth::Auth, request, request::Request};
+use super::{auth, request, request::Request};
 use crate::{State, service::appservice::RegistrationInfo};
 
 /// Extractor for Ruma request structs
@@ -108,7 +107,7 @@ where
 		}
 		let auth = auth::auth(services, &mut request, json_body.as_ref(), &T::METADATA).await?;
 		Ok(Self {
-			body: make_body::<T>(services, &mut request, json_body.as_mut(), &auth)?,
+			body: make_body::<T>(&mut request, json_body.as_mut())?,
 			origin: auth.origin,
 			sender_user: auth.sender_user,
 			sender_device: auth.sender_device,
@@ -118,16 +117,11 @@ where
 	}
 }
 
-fn make_body<T>(
-	services: &Services,
-	request: &mut Request,
-	json_body: Option<&mut CanonicalJsonValue>,
-	auth: &Auth,
-) -> Result<T>
+fn make_body<T>(request: &mut Request, json_body: Option<&mut CanonicalJsonValue>) -> Result<T>
 where
 	T: IncomingRequest,
 {
-	let body = take_body(services, request, json_body, auth);
+	let body = take_body(request, json_body);
 	let http_request = into_http_request(request, body);
 	T::try_from_http_request(http_request, &request.path)
 		.map_err(|e| err!(Request(BadJson(debug_warn!("{e}")))))
@@ -151,38 +145,11 @@ fn into_http_request(request: &Request, body: Bytes) -> hyper::Request<Bytes> {
 }
 
 #[allow(clippy::needless_pass_by_value)]
-fn take_body(
-	services: &Services,
-	request: &mut Request,
-	json_body: Option<&mut CanonicalJsonValue>,
-	auth: &Auth,
-) -> Bytes {
+fn take_body(request: &mut Request, json_body: Option<&mut CanonicalJsonValue>) -> Bytes {
 	let Some(CanonicalJsonValue::Object(json_body)) = json_body else {
 		return mem::take(&mut request.body);
 	};
 
-	let user_id = auth.sender_user.clone().unwrap_or_else(|| {
-		let server_name = services.globals.server_name();
-		UserId::parse_with_server_name(EMPTY, server_name).expect("valid user_id")
-	});
-
-	let uiaa_request = json_body
-		.get("auth")
-		.and_then(CanonicalJsonValue::as_object)
-		.and_then(|auth| auth.get("session"))
-		.and_then(CanonicalJsonValue::as_str)
-		.and_then(|session| {
-			services
-				.uiaa
-				.get_uiaa_request(&user_id, auth.sender_device.as_deref(), session)
-		});
-
-	if let Some(CanonicalJsonValue::Object(initial_request)) = uiaa_request {
-		for (key, value) in initial_request {
-			json_body.entry(key).or_insert(value);
-		}
-	}
-
 	let mut buf = BytesMut::new().writer();
 	serde_json::to_writer(&mut buf, &json_body).expect("value serialization can't fail");
 	buf.into_inner().freeze()

src/service/uiaa/mod.rs

@@ -1,24 +1,28 @@
-use std::{collections::BTreeMap, sync::Arc};
-
-use conduwuit::{
-	Err, Error, Result, SyncRwLock, err, error, implement, utils,
-	utils::{hash, string::EMPTY},
+use std::{
+	collections::{HashMap, HashSet, hash_map::Entry},
+	sync::Arc,
 };
-use database::{Deserialized, Json, Map};
+
+use conduwuit::{Err, Error, Result, error, utils, utils::hash};
+use lettre::Address;
 use ruma::{
-	CanonicalJsonValue, DeviceId, OwnedDeviceId, OwnedUserId, UserId,
+	UserId,
 	api::client::{
 		error::{ErrorKind, StandardErrorBody},
-		uiaa::{AuthData, AuthType, Password, UiaaInfo, UserIdentifier},
+		uiaa::{
+			AuthData, AuthFlow, AuthType, EmailIdentity, Password, ReCaptcha, RegistrationToken,
+			ThirdpartyIdCredentials, UiaaInfo, UserIdentifier,
+		},
 	},
 };
+use serde_json::value::RawValue;
+use tokio::sync::Mutex;
 
-use crate::{Dep, config, globals, registration_tokens, users};
+use crate::{Dep, config, globals, registration_tokens, threepid, users};
 
 pub struct Service {
-	userdevicesessionid_uiaarequest: SyncRwLock<RequestMap>,
-	db: Data,
 	services: Services,
+	uiaa_sessions: Mutex<HashMap<String, UiaaSession>>,
 }
 
 struct Services {
@@ -26,205 +30,144 @@ struct Services {
 	users: Dep<users::Service>,
 	config: Dep<config::Service>,
 	registration_tokens: Dep<registration_tokens::Service>,
+	threepid: Dep<threepid::Service>,
 }
 
-struct Data {
-	userdevicesessionid_uiaainfo: Arc<Map>,
-}
-
-type RequestMap = BTreeMap<RequestKey, CanonicalJsonValue>;
-type RequestKey = (OwnedUserId, OwnedDeviceId, String);
-
-pub const SESSION_ID_LENGTH: usize = 32;
-
 impl crate::Service for Service {
 	fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {
 		Ok(Arc::new(Self {
-			userdevicesessionid_uiaarequest: SyncRwLock::new(RequestMap::new()),
-			db: Data {
-				userdevicesessionid_uiaainfo: args.db["userdevicesessionid_uiaainfo"].clone(),
-			},
 			services: Services {
 				globals: args.depend::<globals::Service>("globals"),
 				users: args.depend::<users::Service>("users"),
 				config: args.depend::<config::Service>("config"),
 				registration_tokens: args
 					.depend::<registration_tokens::Service>("registration_tokens"),
+				threepid: args.depend::<threepid::Service>("threepid"),
 			},
+			uiaa_sessions: Mutex::new(HashMap::new()),
 		}))
 	}
 
 	fn name(&self) -> &str { crate::service::make_name(std::module_path!()) }
 }
 
-/// Creates a new Uiaa session. Make sure the session token is unique.
-#[implement(Service)]
-pub fn create(
-	&self,
-	user_id: &UserId,
-	device_id: &DeviceId,
-	uiaainfo: &UiaaInfo,
-	json_body: &CanonicalJsonValue,
-) {
-	// TODO: better session error handling (why is uiaainfo.session optional in
-	// ruma?)
-	self.set_uiaa_request(
-		user_id,
-		device_id,
-		uiaainfo.session.as_ref().expect("session should be set"),
-		json_body,
-	);
-
-	self.update_uiaa_session(
-		user_id,
-		device_id,
-		uiaainfo.session.as_ref().expect("session should be set"),
-		Some(uiaainfo),
-	);
+struct UiaaSession {
+	info: UiaaInfo,
+	identity: Identity,
 }
 
-#[implement(Service)]
-#[allow(clippy::useless_let_if_seq)]
-pub async fn try_auth(
-	&self,
-	user_id: &UserId,
-	device_id: &DeviceId,
-	auth: &AuthData,
-	uiaainfo: &UiaaInfo,
-) -> Result<(bool, UiaaInfo)> {
-	let mut uiaainfo = if let Some(session) = auth.session() {
-		self.get_uiaa_session(user_id, device_id, session).await?
-	} else {
-		uiaainfo.clone()
-	};
+pub enum UiaaStatus {
+	/// The UIAA session succeeded and the request should be completed as
+	/// normal.
+	Success(Identity),
+	/// More UIAA stages need to be completed, or the current stage failed.
+	Retry(UiaaInfo),
+}
 
-	if uiaainfo.session.is_none() {
-		uiaainfo.session = Some(utils::random_string(SESSION_ID_LENGTH));
+/// Information about the authenticated user's identity.
+///
+/// A field of this struct will only be Some if the user completed
+/// a stage which provided that information. If multiple stages provide
+/// the same field, authentication will fail if they do not all provide
+/// _identical_ values for that field.
+#[derive(Default, Clone)]
+pub struct Identity {
+	/// The authenticated user's user ID, if it could be determined.
+	///
+	/// This will be Some if:
+	/// - The user completed a m.login.password stage
+	/// - The user completed a m.login.email.identity stage, and their email has
+	///   an associated user ID
+	pub localpart: Option<String>,
+
+	/// The authenticated user's email address, if it could be determined.
+	///
+	/// This will be Some if:
+	/// - The user completed a m.login.email.identity stage
+	/// - The user completed a m.login.password stage, and their user ID has an
+	///   associated email
+	pub email: Option<Address>,
+}
+
+macro_rules! identity_update_fn {
+	(fn $method:ident($field:ident : $type:ty)else $error:literal) => {
+		fn $method(&mut self, $field: $type) -> Result<(), StandardErrorBody> {
+			if self.$field.is_none() {
+				self.$field = Some($field);
+				Ok(())
+			} else if self.$field == Some($field) {
+				Ok(())
+			} else {
+				Err(StandardErrorBody {
+					kind: ErrorKind::InvalidParam,
+					message: $error.to_owned(),
+				})
+			}
+		}
+	};
+}
+
+impl Identity {
+	identity_update_fn!(fn try_set_localpart(localpart: String) else "User ID mismatch");
+
+	identity_update_fn!(fn try_set_email(email: Address) else "Email mismatch");
+
+	/// Create an Identity with the localpart of the provided user ID
+	/// and all other fields set to None.
+	#[must_use]
+	pub fn from_user_id(user_id: &UserId) -> Self {
+		Self {
+			localpart: Some(user_id.localpart().to_owned()),
+			..Default::default()
+		}
+	}
+}
+
+impl Service {
+	const SESSION_ID_LENGTH: usize = 32;
+
+	/// Create a new UIAA session with a random session ID.
+	///
+	/// If information about the user's identity is already known, it may be
+	/// supplied with the `identity` parameter. Authentication will fail if
+	/// flows provide different values for known identity information.
+	///
+	/// Returns the info of the newly created session.
+	pub async fn create_session(
+		&self,
+		flows: Vec<AuthFlow>,
+		params: Box<RawValue>,
+		identity: Option<Identity>,
+	) -> UiaaInfo {
+		let mut uiaa_sessions = self.uiaa_sessions.lock().await;
+
+		let session_id = utils::random_string(Self::SESSION_ID_LENGTH);
+		let mut info = UiaaInfo::new(flows, params);
+		info.session = Some(session_id.clone());
+
+		uiaa_sessions.insert(session_id, UiaaSession {
+			info: info.clone(),
+			identity: identity.unwrap_or_default(),
+		});
+
+		info
 	}
 
-	match auth {
-		// Find out what the user completed
-		| AuthData::Password(Password {
-			identifier,
-			password,
-			#[cfg(feature = "element_hacks")]
-			user,
-			..
-		}) => {
-			#[cfg(feature = "element_hacks")]
-			let username = if let Some(UserIdentifier::UserIdOrLocalpart(username)) = identifier {
-				username
-			} else if let Some(username) = user {
-				username
-			} else {
-				return Err(Error::BadRequest(
-					ErrorKind::Unrecognized,
-					"Identifier type not recognized.",
-				));
-			};
+	/// Proceed with UIAA authentication given a client's authorization data.
+	pub async fn continue_session(&self, auth: &AuthData) -> Result<UiaaStatus> {
+		let Some(session) = auth.session() else {
+			return Err!(Request(MissingParam("No session provided")));
+		};
 
-			#[cfg(not(feature = "element_hacks"))]
-			let Some(UserIdentifier::UserIdOrLocalpart(username)) = identifier else {
-				return Err(Error::BadRequest(
-					ErrorKind::Unrecognized,
-					"Identifier type not recognized.",
-				));
-			};
+		// Hold this lock for the entire function to make sure that, if try_auth()
+		// is called concurrently with the same session, only one call will succeed
+		let mut uiaa_sessions = self.uiaa_sessions.lock().await;
 
-			let user_id_from_username = UserId::parse_with_server_name(
-				username.clone(),
-				self.services.globals.server_name(),
-			)
-			.map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "User ID is invalid."))?;
+		let Entry::Occupied(mut session) = uiaa_sessions.entry(session.to_owned()) else {
+			return Err!(Request(InvalidParam("Invalid session")));
+		};
 
-			// Check if the access token being used matches the credentials used for UIAA
-			if user_id.localpart() != user_id_from_username.localpart() {
-				return Err!(Request(Forbidden("User ID and access token mismatch.")));
-			}
-			let user_id = user_id_from_username;
-
-			// Check if password is correct
-			let mut password_verified = false;
-
-			// First try local password hash verification
-			if let Ok(hash) = self.services.users.password_hash(&user_id).await {
-				password_verified = hash::verify_password(password, &hash).is_ok();
-			}
-
-			// If local password verification failed, try LDAP authentication
-			#[cfg(feature = "ldap")]
-			if !password_verified && self.services.config.ldap.enable {
-				// Search for user in LDAP to get their DN
-				if let Ok(dns) = self.services.users.search_ldap(&user_id).await {
-					if let Some((user_dn, _is_admin)) = dns.first() {
-						// Try to authenticate with LDAP
-						password_verified = self
-							.services
-							.users
-							.auth_ldap(user_dn, password)
-							.await
-							.is_ok();
-					}
-				}
-			}
-
-			if !password_verified {
-				uiaainfo.auth_error = Some(StandardErrorBody {
-					kind: ErrorKind::forbidden(),
-					message: "Invalid username or password.".to_owned(),
-				});
-
-				return Ok((false, uiaainfo));
-			}
-
-			// Password was correct! Let's add it to `completed`
-			uiaainfo.completed.push(AuthType::Password);
-		},
-		| AuthData::ReCaptcha(r) => {
-			let Some(ref private_site_key) = self.services.config.recaptcha_private_site_key
-			else {
-				return Err!(Request(Forbidden("ReCaptcha is not configured.")));
-			};
-			match recaptcha_verify::verify_v3(private_site_key, r.response.as_str(), None).await {
-				| Ok(()) => {
-					uiaainfo.completed.push(AuthType::ReCaptcha);
-				},
-				| Err(e) => {
-					error!("ReCaptcha verification failed: {e:?}");
-					uiaainfo.auth_error = Some(StandardErrorBody {
-						kind: ErrorKind::forbidden(),
-						message: "ReCaptcha verification failed.".to_owned(),
-					});
-					return Ok((false, uiaainfo));
-				},
-			}
-		},
-		| AuthData::RegistrationToken(t) => {
-			let token = t.token.trim().to_owned();
-
-			if let Some(valid_token) = self
-				.services
-				.registration_tokens
-				.validate_token(token)
-				.await
-			{
-				self.services
-					.registration_tokens
-					.mark_token_as_used(valid_token);
-
-				uiaainfo.completed.push(AuthType::RegistrationToken);
-			} else {
-				uiaainfo.auth_error = Some(StandardErrorBody {
-					kind: ErrorKind::forbidden(),
-					message: "Invalid registration token.".to_owned(),
-				});
-				return Ok((false, uiaainfo));
-			}
-		},
-		| AuthData::Dummy(_) => {
-			uiaainfo.completed.push(AuthType::Dummy);
-		},
-		| AuthData::FallbackAcknowledgement(_) => {
+		if let &AuthData::FallbackAcknowledgement(_) = auth {
 			// The client is checking if authentication has succeeded out-of-band. This is
 			// possible if the client is using "fallback auth" (see spec section
 			// 4.9.1.4), which we don't support (and probably never will, because it's a
@@ -232,109 +175,276 @@ pub async fn try_auth(
 
 			// Return early to tell the client that no, authentication did not succeed while
 			// it wasn't looking.
-			return Ok((false, uiaainfo));
-		},
-		| k => error!("type not supported: {:?}", k),
-	}
-
-	// Check if a flow now succeeds
-	let mut completed = false;
-	'flows: for flow in &mut uiaainfo.flows {
-		for stage in &flow.stages {
-			if !uiaainfo.completed.contains(stage) {
-				continue 'flows;
-			}
+			return Ok(UiaaStatus::Retry(session.get().info.clone()));
+		}
+
+		let completed = 'completed: {
+			let UiaaSession { info, identity } = session.get_mut();
+
+			let completed_stages: HashSet<_> = info
+				.completed
+				.iter()
+				.map(AuthType::as_str)
+				.map(ToOwned::to_owned)
+				.collect();
+
+			// If the provided stage has already been completed, return early
+			if completed_stages
+				.contains(auth.auth_type().expect("auth type should be set").as_str())
+			{
+				return Ok(UiaaStatus::Retry(session.get().info.clone()));
+			}
+
+			match self.check_stage(auth, identity.clone()).await {
+				| Ok((completed_stage, updated_identity)) => {
+					info.completed.push(completed_stage);
+					*identity = updated_identity;
+				},
+				| Err(error) => {
+					info.auth_error = Some(error);
+				},
+			}
+
+			// Check all flows to see if any of them succeeded
+
+			for flow in &info.flows {
+				let flow_stages = flow
+					.stages
+					.iter()
+					.map(AuthType::as_str)
+					.map(ToOwned::to_owned)
+					.collect();
+
+				if completed_stages.is_superset(&flow_stages) {
+					// All stages in this flow are completed
+					break 'completed true;
+				}
+			}
+
+			// No flows had all their stages completed
+			break 'completed false;
+		};
+
+		if completed {
+			// This session is complete, remove it and return success
+			let (_, UiaaSession { identity, .. }) = session.remove_entry();
+
+			Ok(UiaaStatus::Success(identity))
+		} else {
+			// The client needs to try again, return the updated session
+			Ok(UiaaStatus::Retry(session.get().info.clone()))
 		}
-		// We didn't break, so this flow succeeded!
-		completed = true;
 	}
 
-	if !completed {
-		self.update_uiaa_session(
-			user_id,
-			device_id,
-			uiaainfo.session.as_ref().expect("session is always set"),
-			Some(&uiaainfo),
-		);
+	/// Perform the full UIAA authentication sequence for a route given its
+	/// authentication data.
+	#[inline]
+	pub async fn authenticate(
+		&self,
+		auth: &Option<AuthData>,
+		flows: Vec<AuthFlow>,
+		params: Box<RawValue>,
+		identity: Option<Identity>,
+	) -> Result<Identity> {
+		match auth.as_ref() {
+			| None => {
+				let info = self.create_session(flows, params, identity).await;
 
-		return Ok((false, uiaainfo));
+				Err(Error::Uiaa(info))
+			},
+			| Some(auth) => match self.continue_session(auth).await? {
+				| UiaaStatus::Retry(info) => Err(Error::Uiaa(info)),
+				| UiaaStatus::Success(identity) => Ok(identity),
+			},
+		}
 	}
 
-	// UIAA was successful! Remove this session and return true
-	self.update_uiaa_session(
-		user_id,
-		device_id,
-		uiaainfo.session.as_ref().expect("session is always set"),
-		None,
-	);
-
-	Ok((true, uiaainfo))
-}
-
-#[implement(Service)]
-fn set_uiaa_request(
-	&self,
-	user_id: &UserId,
-	device_id: &DeviceId,
-	session: &str,
-	request: &CanonicalJsonValue,
-) {
-	let key = (user_id.to_owned(), device_id.to_owned(), session.to_owned());
-	self.userdevicesessionid_uiaarequest
-		.write()
-		.insert(key, request.to_owned());
-}
-
-#[implement(Service)]
-pub fn get_uiaa_request(
-	&self,
-	user_id: &UserId,
-	device_id: Option<&DeviceId>,
-	session: &str,
-) -> Option<CanonicalJsonValue> {
-	let key = (
-		user_id.to_owned(),
-		device_id.unwrap_or_else(|| EMPTY.into()).to_owned(),
-		session.to_owned(),
-	);
-
-	self.userdevicesessionid_uiaarequest
-		.read()
-		.get(&key)
-		.cloned()
-}
-
-#[implement(Service)]
-fn update_uiaa_session(
-	&self,
-	user_id: &UserId,
-	device_id: &DeviceId,
-	session: &str,
-	uiaainfo: Option<&UiaaInfo>,
-) {
-	let key = (user_id, device_id, session);
-
-	if let Some(uiaainfo) = uiaainfo {
-		self.db
-			.userdevicesessionid_uiaainfo
-			.put(key, Json(uiaainfo));
-	} else {
-		self.db.userdevicesessionid_uiaainfo.del(key);
-	}
-}
-
-#[implement(Service)]
-async fn get_uiaa_session(
-	&self,
-	user_id: &UserId,
-	device_id: &DeviceId,
-	session: &str,
-) -> Result<UiaaInfo> {
-	let key = (user_id, device_id, session);
-	self.db
-		.userdevicesessionid_uiaainfo
-		.qry(&key)
+	/// A helper to perform UIAA authentication with just a password stage.
+	#[inline]
+	pub async fn authenticate_password(
+		&self,
+		auth: &Option<AuthData>,
+		identity: Option<Identity>,
+	) -> Result<Identity> {
+		self.authenticate(
+			auth,
+			vec![AuthFlow::new(vec![AuthType::Password])],
+			Box::default(),
+			identity,
+		)
 		.await
-		.deserialized()
-		.map_err(|_| err!(Request(Forbidden("UIAA session does not exist."))))
+	}
+
+	/// Check if the provided authentication data is valid.
+	///
+	/// Returns the completed stage's type on success and error information on
+	/// failure.
+	async fn check_stage(
+		&self,
+		auth: &AuthData,
+		mut identity: Identity,
+	) -> Result<(AuthType, Identity), StandardErrorBody> {
+		// Note: This function takes ownership of `identity` because mutations to the
+		// identity must not be applied unless checking the stage succeeds. The
+		// updated identity is returned as part of the Ok value, and
+		// `continue_session` handles saving it to `uiaa_sessions`.
+		//
+		// This also means it's fine to mutate `identity` at any point in this function,
+		// because those mutations won't be saved unless the function returns Ok.
+
+		match auth {
+			| AuthData::Dummy(_) => Ok(AuthType::Dummy),
+			| AuthData::EmailIdentity(EmailIdentity {
+				thirdparty_id_creds: ThirdpartyIdCredentials { client_secret, sid, .. },
+				..
+			}) => {
+				match self
+					.services
+					.threepid
+					.consume_valid_session(sid.as_str(), client_secret.as_str())
+					.await
+				{
+					| Ok(email) => {
+						if let Some(localpart) =
+							self.services.threepid.get_localpart_for_email(&email).await
+						{
+							identity.try_set_localpart(localpart)?;
+						}
+
+						identity.try_set_email(email)?;
+
+						Ok(AuthType::EmailIdentity)
+					},
+					| Err(message) => Err(StandardErrorBody {
+						kind: ErrorKind::ThreepidAuthFailed,
+						message: message.into_owned(),
+					}),
+				}
+			},
+			#[allow(clippy::useless_let_if_seq)]
+			| AuthData::Password(Password { identifier, password, .. }) => {
+				let user_id_or_localpart = match identifier {
+					| Some(UserIdentifier::UserIdOrLocalpart(username)) => username.to_owned(),
+					| Some(UserIdentifier::Email { address }) => {
+						let Ok(email) = Address::try_from(address.to_owned()) else {
+							return Err(StandardErrorBody {
+								kind: ErrorKind::InvalidParam,
+								message: "Email is invalid".to_owned(),
+							});
+						};
+
+						if let Some(localpart) =
+							self.services.threepid.get_localpart_for_email(&email).await
+						{
+							identity.try_set_email(email)?;
+
+							localpart
+						} else {
+							return Err(StandardErrorBody {
+								kind: ErrorKind::forbidden(),
+								message: "Invalid identifier or password".to_owned(),
+							});
+						}
+					},
+					| _ =>
+						return Err(StandardErrorBody {
+							kind: ErrorKind::Unrecognized,
+							message: "Identifier type not recognized".to_owned(),
+						}),
+				};
+
+				let Ok(user_id) = UserId::parse_with_server_name(
+					user_id_or_localpart,
+					self.services.globals.server_name(),
+				) else {
+					return Err(StandardErrorBody {
+						kind: ErrorKind::InvalidParam,
+						message: "User ID is invalid".to_owned(),
+					});
+				};
+
+				// Check if password is correct
+				let mut password_verified = false;
+
+				// First try local password hash verification
+				if let Ok(hash) = self.services.users.password_hash(&user_id).await {
+					password_verified = hash::verify_password(password, &hash).is_ok();
+				}
+
+				// If local password verification failed, try LDAP authentication
+				#[cfg(feature = "ldap")]
+				if !password_verified && self.services.config.ldap.enable {
+					// Search for user in LDAP to get their DN
+					if let Ok(dns) = self.services.users.search_ldap(&user_id).await {
+						if let Some((user_dn, _is_admin)) = dns.first() {
+							// Try to authenticate with LDAP
+							password_verified = self
+								.services
+								.users
+								.auth_ldap(user_dn, password)
+								.await
+								.is_ok();
+						}
+					}
+				}
+
+				if password_verified {
+					identity.try_set_localpart(user_id.localpart().to_owned())?;
+
+					Ok(AuthType::Password)
+				} else {
+					Err(StandardErrorBody {
+						kind: ErrorKind::forbidden(),
+						message: "Invalid identifier or password".to_owned(),
+					})
+				}
+			},
+			| AuthData::ReCaptcha(ReCaptcha { response, .. }) => {
+				let Some(ref private_site_key) = self.services.config.recaptcha_private_site_key
+				else {
+					return Err(StandardErrorBody {
+						kind: ErrorKind::forbidden(),
+						message: "ReCaptcha is not configured".to_owned(),
+					});
+				};
+
+				match recaptcha_verify::verify_v3(private_site_key, response, None).await {
+					| Ok(()) => Ok(AuthType::ReCaptcha),
+					| Err(e) => {
+						error!("ReCaptcha verification failed: {e:?}");
+						Err(StandardErrorBody {
+							kind: ErrorKind::forbidden(),
+							message: "ReCaptcha verification failed".to_owned(),
+						})
+					},
+				}
+			},
+			| AuthData::RegistrationToken(RegistrationToken { token, .. }) => {
+				let token = token.trim().to_owned();
+
+				if let Some(valid_token) = self
+					.services
+					.registration_tokens
+					.validate_token(token)
+					.await
+				{
+					self.services
+						.registration_tokens
+						.mark_token_as_used(valid_token);
+
+					Ok(AuthType::RegistrationToken)
+				} else {
+					Err(StandardErrorBody {
+						kind: ErrorKind::forbidden(),
+						message: "Invalid registration token".to_owned(),
+					})
+				}
+			},
+			| _ => Err(StandardErrorBody {
+				kind: ErrorKind::Unrecognized,
+				message: "Unsupported stage type".into(),
+			}),
+		}
+		.map(|auth_type| (auth_type, identity))
+	}
 }