continuwuity/src/service/uiaa/mod.rs

use std::{
	borrow::Cow,
	collections::{HashMap, HashSet, hash_map::Entry},
	sync::Arc,
};

use conduwuit::{Err, Error, Result, utils};
use futures::StreamExt;
use lettre::Address;
use ruma::{
	DeviceId, UserId,
	api::{
		client::uiaa::{
			AuthData, AuthFlow, AuthType, EmailIdentity, EmailUserIdentifier,
			MatrixUserIdentifier, Password, ReCaptcha, RegistrationToken,
			ThirdpartyIdCredentials, UiaaInfo, UserIdentifier,
		},
		error::{ErrorKind, StandardErrorBody},
	},
	assign,
};
use serde_json::{
	json,
	value::{RawValue, to_raw_value},
};
use tokio::sync::Mutex;

use crate::{
	Dep, client, config, firstrun, globals,
	oauth::{self, OAuthTicket},
	registration_tokens, threepid, users,
};

pub struct Service {
	services: Services,
	uiaa_sessions: Mutex<HashMap<String, UiaaSession>>,
}

struct Services {
	config: Dep<config::Service>,
	client: Dep<client::Service>,
	firstrun: Dep<firstrun::Service>,
	globals: Dep<globals::Service>,
	oauth: Dep<oauth::Service>,
	registration_tokens: Dep<registration_tokens::Service>,
	threepid: Dep<threepid::Service>,
	users: Dep<users::Service>,
}

#[derive(Debug)]
pub enum TrustedFlowStatus {
	Unavailable,
	Available,
}

#[derive(Debug)]
pub enum UntrustedFlowStatus {
	Unavailable,
	Available {
		require_email: bool,
	},
}

pub struct FlowStatus {
	pub trusted: TrustedFlowStatus,
	pub untrusted: UntrustedFlowStatus,
}

impl FlowStatus {
	#[must_use]
	pub fn any_available(&self) -> bool {
		matches!(self.trusted, TrustedFlowStatus::Available)
			|| matches!(self.untrusted, UntrustedFlowStatus::Available { .. })
	}
}

impl crate::Service for Service {
	fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {
		Ok(Arc::new(Self {
			services: Services {
				config: args.depend::<config::Service>("config"),
				client: args.depend::<client::Service>("client"),
				firstrun: args.depend::<firstrun::Service>("firstrun"),
				globals: args.depend::<globals::Service>("globals"),
				oauth: args.depend::<oauth::Service>("oauth"),
				registration_tokens: args
					.depend::<registration_tokens::Service>("registration_tokens"),
				threepid: args.depend::<threepid::Service>("threepid"),
				users: args.depend::<users::Service>("users"),
			},
			uiaa_sessions: Mutex::new(HashMap::new()),
		}))
	}

	fn name(&self) -> &str { crate::service::make_name(std::module_path!()) }
}

struct UiaaSession {
	session_metadata: UiaaSessionMetadata,
	info: UiaaInfo,
}

#[derive(Clone)]
enum UiaaSessionMetadata {
	Legacy {
		identity: Identity,
	},
	OAuth {
		localpart: String,
		ticket: OAuthTicket,
	},
}

impl UiaaSessionMetadata {
	fn into_identity(self) -> Identity {
		match self {
			| Self::Legacy { identity } => identity,
			| Self::OAuth { localpart, .. } =>
				assign!(Identity::default(), { localpart: Some(localpart) }),
		}
	}
}

/// Information about the user which is initiating this UIAA session.
pub struct UiaaInitiator<'a> {
	user_id: &'a UserId,
	device_id: Option<&'a DeviceId>,
	oauth_ticket: Option<OAuthTicket>,
}

impl<'a> UiaaInitiator<'a> {
	#[must_use]
	pub fn new(user_id: &'a UserId, device_id: Option<&'a DeviceId>) -> Self {
		Self { user_id, device_id, oauth_ticket: None }
	}

	#[must_use]
	pub fn with_oauth_ticket(
		user_id: &'a UserId,
		device_id: Option<&'a DeviceId>,
		oauth_ticket: OAuthTicket,
	) -> Self {
		Self {
			user_id,
			device_id,
			oauth_ticket: Some(oauth_ticket),
		}
	}
}

/// 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::new(ErrorKind::InvalidParam, $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]
	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;

	/// Perform the full UIAA authentication sequence for a route given its
	/// authentication data.
	pub async fn authenticate(
		&self,
		auth: &Option<AuthData>,
		flows: Vec<AuthFlow>,
		params: Box<RawValue>,
		initiator: Option<UiaaInitiator<'_>>,
	) -> Result<Identity> {
		match auth.as_ref() {
			| None => {
				let info = self.create_session(flows, params, initiator).await?;

				Err(Error::Uiaa(info))
			},
			| Some(auth) => {
				let session: Cow<'_, str> = match auth.session() {
					| Some(session) => session.into(),
					| None => {
						// Clients are allowed to send UIAA requests with an
						// auth dict and no session if they want to start
						// the UIAA exchange with existing authentication
						// data. If that happens, we create a new session
						// here.
						self.create_session(flows, params, initiator)
							.await?
							.session
							.unwrap()
							.into()
					},
				};

				match self.continue_session(auth, &session).await? {
					| Ok(identity) => Ok(identity),
					| Err(info) => Err(Error::Uiaa(info)),
				}
			},
		}
	}

	/// A helper to perform UIAA authentication with just a password stage.
	#[inline]
	pub async fn authenticate_password(
		&self,
		auth: &Option<AuthData>,
		user_id: &UserId,
		device_id: Option<&DeviceId>,
		oauth_ticket: Option<OAuthTicket>,
	) -> Result<Identity> {
		self.authenticate(
			auth,
			vec![AuthFlow::new(vec![AuthType::Password])],
			Box::default(),
			Some(UiaaInitiator { user_id, device_id, oauth_ticket }),
		)
		.await
	}

	/// 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 `initiator` parameter. Authentication will fail if
	/// flows provide different values for known identity information.
	///
	/// Returns the info of the newly created session.
	async fn create_session(
		&self,
		flows: Vec<AuthFlow>,
		params: Box<RawValue>,
		initiator: Option<UiaaInitiator<'_>>,
	) -> Result<UiaaInfo> {
		let mut uiaa_sessions = self.uiaa_sessions.lock().await;

		let session_id = utils::random_string(Self::SESSION_ID_LENGTH);

		let mut info = assign!(UiaaInfo::new(flows), { params: Some(params), session: Some(session_id.clone()) });

		let session_metadata = if let Some(initiator) = initiator {
			let is_oauth = if let Some(device_id) = initiator.device_id {
				self.services
					.oauth
					.get_session_info_for_device(initiator.user_id, device_id)
					.await
					.is_some()
			} else {
				// Appservices never have oauth sessions
				false
			};

			if is_oauth {
				if let Some(oauth_ticket) = initiator.oauth_ticket {
					let ticket_url = self
						.services
						.config
						.get_client_domain()
						.join(&format!(
							"{}{}",
							conduwuit_core::ROUTE_PREFIX,
							oauth_ticket.ticket_issue_path()
						))
						.unwrap();

					info.flows = vec![AuthFlow::new(vec![AuthType::OAuth])];
					info.params = Some(
						to_raw_value(&json!({
							AuthType::OAuth.as_str(): {
								"url": ticket_url,
							},
							// TODO(compat): This is necessary for older versions of matrix-rust-sdk
							"org.matrix.cross_signing_reset": {
								"url": ticket_url,
							}
						}))
						.unwrap(),
					);

					UiaaSessionMetadata::OAuth {
						localpart: initiator.user_id.localpart().to_owned(),
						ticket: oauth_ticket,
					}
				} else {
					return Err!(Request(Forbidden(
						"Clients authorized with OAuth cannot use this route."
					)));
				}
			} else {
				UiaaSessionMetadata::Legacy {
					identity: Identity::from_user_id(initiator.user_id),
				}
			}
		} else {
			UiaaSessionMetadata::Legacy { identity: Identity::default() }
		};

		// Legacy sessions aren't available if OAuth is required
		if matches!(&session_metadata, UiaaSessionMetadata::Legacy { .. })
			&& !self
				.services
				.config
				.oauth
				.compatibility_mode()
				.uiaa_available()
		{
			return Err!(Request(Unrecognized(
				"User-interactive authentication is unavailable on this server"
			)));
		}

		uiaa_sessions.insert(session_id, UiaaSession { session_metadata, info: info.clone() });

		Ok(info)
	}

	/// Proceed with UIAA authentication given a client's authorization data.
	async fn continue_session(
		&self,
		auth: &AuthData,
		session: &str,
	) -> Result<Result<Identity, UiaaInfo>> {
		// 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 Entry::Occupied(mut session) = uiaa_sessions.entry(session.to_owned()) else {
			return Err!(Request(InvalidParam("Invalid session")));
		};

		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 disgusting hack).

			// Return early to tell the client that no, authentication did not
			// succeed while it wasn't looking.
			return Ok(Err(session.get().info.clone()));
		}

		let completed = {
			let UiaaSession { session_metadata, info } = session.get_mut();

			let auth_type = auth.auth_type().expect("auth type should be set");

			let flow_stages: Vec<HashSet<_>> = info
				.flows
				.iter()
				.map(|flow| {
					flow.stages
						.iter()
						.map(AuthType::as_str)
						.map(ToOwned::to_owned)
						.collect()
				})
				.collect();

			let mut completed_stages: HashSet<_> = info
				.completed
				.iter()
				.map(AuthType::as_str)
				.map(ToOwned::to_owned)
				.collect();

			// Don't allow stages which aren't in any flows
			if !flow_stages
				.iter()
				.any(|stages| stages.contains(auth_type.as_str()))
			{
				return Err!(Request(InvalidParam("No flows include the supplied stage")));
			}

			// If the provided stage hasn't already been completed, check it for
			// completion
			if !completed_stages.contains(auth_type.as_str()) {
				match self.check_stage(auth, session_metadata.clone()).await {
					| Ok((completed_stage, updated_metadata)) => {
						info.auth_error = None;
						completed_stages.insert(completed_stage.to_string());
						info.completed.push(completed_stage);
						*session_metadata = updated_metadata;
					},
					| Err(error) => {
						info.auth_error = Some(Box::new(error));
					},
				}
			}

			// UIAA is completed if all stages in any flow are completed
			flow_stages
				.iter()
				.any(|stages| completed_stages.is_superset(stages))
		};

		if completed {
			// This session is complete, remove it and return success
			let (_, UiaaSession { session_metadata, .. }) = session.remove_entry();

			Ok(Ok(session_metadata.into_identity()))
		} else {
			// The client needs to try again, return the updated session
			Ok(Err(session.get().info.clone()))
		}
	}

	/// 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 session_metadata: UiaaSessionMetadata,
	) -> Result<(AuthType, UiaaSessionMetadata), StandardErrorBody> {
		// Note: This function takes ownership of `session_metadata` because
		// mutations to the identity (if it's a legacy session) 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.

		let completed_auth_type = match &mut session_metadata {
			| UiaaSessionMetadata::OAuth { localpart, ticket } => {
				// m.oauth is the only valid stage for oauth sessions
				assert!(
					matches!(auth, AuthData::OAuth(_)),
					"got non-oauth auth data for oauth session"
				);

				if self.services.oauth.try_consume_ticket(localpart, *ticket) {
					Ok(AuthType::OAuth)
				} else {
					Err(StandardErrorBody::new(
						ErrorKind::Forbidden,
						"No OAuth ticket available".to_owned(),
					))
				}
			},
			| UiaaSessionMetadata::Legacy { identity } => match auth {
				| AuthData::Dummy(_) => Ok(AuthType::Dummy),
				| AuthData::EmailIdentity(EmailIdentity {
					thirdparty_id_creds: ThirdpartyIdCredentials { client_secret, sid, .. },
					..
				}) => match self.services.threepid.get_valid_session(sid, client_secret) {
					| Ok(session) => {
						let email = session.consume();

						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::new(
						ErrorKind::ThreepidAuthFailed,
						message.into_owned(),
					)),
				},
				#[allow(clippy::useless_let_if_seq)]
				| AuthData::Password(Password { identifier, password, .. }) => {
					let user_id_or_localpart = match identifier {
						| UserIdentifier::Matrix(MatrixUserIdentifier { user, .. }) =>
							user.to_owned(),
						| UserIdentifier::Email(EmailUserIdentifier { address, .. }) => {
							let Ok(email) = Address::try_from(address.to_owned()) else {
								return Err(StandardErrorBody::new(
									ErrorKind::InvalidParam,
									"Email is malformed".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::new(
									ErrorKind::Forbidden,
									"Invalid identifier or password".to_owned(),
								));
							}
						},
						| _ =>
							return Err(StandardErrorBody::new(
								ErrorKind::Unrecognized,
								"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::new(
							ErrorKind::InvalidParam,
							"User ID is malformed".to_owned(),
						));
					};

					if self
						.services
						.users
						.check_password(&user_id, password)
						.await
						.is_ok()
					{
						identity.try_set_localpart(user_id.localpart().to_owned())?;

						Ok(AuthType::Password)
					} else {
						Err(StandardErrorBody::new(
							ErrorKind::Forbidden,
							"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::new(
							ErrorKind::Forbidden,
							"ReCaptcha is not configured".to_owned(),
						));
					};

					if self
						.services
						.client
						.recaptcha_verify(private_site_key, response)
						.await
					{
						Ok(AuthType::ReCaptcha)
					} else {
						Err(StandardErrorBody::new(
							ErrorKind::CaptchaInvalid,
							"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::new(
							ErrorKind::Forbidden,
							"Invalid registration token".to_owned(),
						))
					}
				},
				| AuthData::Terms(_) => Ok(AuthType::Terms),
				| unknown => {
					// We already checked that the stage type is one that exists
					// in the flow, so we can only get here if we ourselves
					// served a flow with a stage that we don't understand.
					panic!("tried to check an unsupported stage type: {unknown:?}");
				},
			},
		}?;

		Ok((completed_auth_type, session_metadata))
	}

	pub async fn registration_flow_status(&self) -> FlowStatus {
		// Allow registration if it's enabled in the config file or if this is
		// the first run (so the first user account can be created)
		let allow_registration =
			self.services.config.allow_registration || self.services.firstrun.is_first_run();

		// Trusted flow is only available if any registration tokens exist
		let trusted = {
			if !allow_registration {
				TrustedFlowStatus::Unavailable
			} else if self
				.services
				.registration_tokens
				.iterate_tokens()
				.next()
				.await
				.is_some()
			{
				TrustedFlowStatus::Available
			} else {
				TrustedFlowStatus::Unavailable
			}
		};

		// Untrusted flow is available if email is required for registration,
		// or reCAPTCHA is configured, or open registration is enabled
		let untrusted = {
			let require_email = self
				.services
				.config
				.smtp
				.as_ref()
				.is_some_and(|smtp| smtp.require_email_for_registration);

			if !allow_registration || self.services.firstrun.is_first_run() {
				UntrustedFlowStatus::Unavailable
			} else if self.services.config.recaptcha_private_site_key.is_some() || require_email {
				UntrustedFlowStatus::Available { require_email }
			} else if self.services.config.enable_challengeless_registration() {
				UntrustedFlowStatus::Available { require_email: false }
			} else {
				UntrustedFlowStatus::Unavailable
			}
		};

		FlowStatus { trusted, untrusted }
	}
}