arminius_server_client/handlers/

user_keys.rs

// SPDX-FileCopyrightText: 2025 BWI GmbH
//
// SPDX-License-Identifier: Apache-2.0

//! Implements APIs for client key upload from the matrix spec

#![allow(
	clippy::allow_attributes,
	reason = "This source file contains enums which allow `clippy::missing_docs_in_private_items`. But an exception to this lint will mistakenly report that it's unfulfilled when used with expect, even though it is not."
)]

mod tests;

use std::{collections::BTreeMap, time::Duration};

use arminius_lib_authentication::SessionInfo;
use arminius_lib_component::{
	Component, ServerState,
	error::{
		BadJson, InputJsonError, InternalServerError, MatrixApiError, ResultExt as _,
		create_custom_matrix_error,
	},
	ruma_requests::ArminiusResponse,
};
use arminius_lib_db::{Database as _, TransactedDb};
use arminius_lib_db_keys::{
	FallbackKeysQueries as _, OneTimeKeysQueries as _, UserKeysQueries as _,
};
use arminius_lib_db_signatures::LocalSignatureQueries as _;
use arminius_lib_macros::{FromDbError, FromNatsError};
use arminius_lib_model::user_keys::DbKeyType;
use arminius_lib_user_keys::{
	FederatedKeyFailureReason,
	identity_keys::{FederatedIdentityKeysResponse, RawUserKeys},
	one_time_keys::FederatedOTKResponse,
	signatures::{SignatureVerificationError, verify_signatures},
	traits::SigningKey as _,
};
use arminius_server_federation_output::clients::FederationNatsClient;
use axum::{Router, extract::State, routing::post};
use reqwest::StatusCode;
use ruma::{
	DeviceId, OwnedServerName, UserId,
	api::{
		client::keys::{claim_keys, get_keys, upload_keys},
		error::{Error as MatrixError, ErrorKind},
	},
	encryption::{DeviceKeys, OneTimeKey},
	serde::{Base64, base64::Standard},
};
use serde_json::json;
use tokio_util::task::JoinMap;

use crate::client_request::AuthenticatedClientRequest;

/// Thrown when signature verification fails.
#[derive(Debug, thiserror::Error)]
#[error("Invalid signature: {0}")]
pub struct SignatureError(String);

// Manually implemented, because `ErrorKind` does not agree with custom error
// codes.
impl MatrixApiError for SignatureError {
	fn create_matrix_error(&self) -> MatrixError {
		create_custom_matrix_error(
			StatusCode::BAD_REQUEST,
			"M_INVALID_SIGNATURE",
			&self.to_string(),
		)
	}
}

/// Converts a [`SignatureVerificationError`] into a [`UserKeysError`] so that
/// it can be returned as Matrix error.
impl From<SignatureVerificationError> for UserKeysError {
	fn from(signature_error: SignatureVerificationError) -> Self {
		match signature_error {
			SignatureVerificationError::BadJson(bad_json) => {
				Self::InputJsonError(InputJsonError::BadJson(bad_json))
			}
			other_error => Self::SignatureVerification(SignatureError(other_error.to_string())),
		}
	}
}

#[allow(
	missing_docs,
	clippy::missing_docs_in_private_items,
	reason = "TODO(#2003): This exception predates us enabling `clippy::allow_attributes_without_reason`, so we don't know why it was enabled"
)]
#[derive(Debug, thiserror::Error, MatrixApiError, FromDbError, FromNatsError)]
enum UserKeysError {
	#[error(transparent)]
	InputJsonError(#[from] InputJsonError),

	/// The requested action was impossible because we do not have the user's
	/// device keys.
	#[error("Missing device keys")]
	#[matrix_api_error(status_code = StatusCode::BAD_REQUEST, error_kind = ErrorKind::MissingParam)]
	MissingDeviceKeys,

	#[error(transparent)]
	SignatureVerification(#[from] SignatureError),

	#[error(transparent)]
	InternalServerError(#[from] InternalServerError),
}

/// Given the user and device IDs of the authenticated request made, performs
/// various checks on a [`DeviceKeys`].
///
/// # Errors
///
/// This function returns a [`BadJson`] error for any of the following cases:
/// - The device key is missing or has a wrongly formatted field.
/// - The user ID or device ID does not match that of the authenticated user/device making the
///   request.
fn check_device_key(
	expected_user_id: &UserId,
	expected_device_id: &DeviceId,
	keys: &DeviceKeys,
) -> Result<(), InputJsonError> {
	if keys.keys.keys().any(|key_id| key_id.key_name() != expected_device_id) {
		return Err(InputJsonError::BadJson(BadJson(
			"Device keys contains keys for another device".to_owned(),
		)));
	}

	if keys.user_id != expected_user_id {
		return Err(InputJsonError::BadJson(BadJson(
			"Field `user_id` does not match in device key".to_owned(),
		)));
	}

	if keys.device_id != expected_device_id {
		return Err(InputJsonError::BadJson(BadJson(
			"Field `device_id` does not match the requester".to_owned(),
		)));
	}

	Ok(())
}

/// Parameter used to tell `check_onetime_key()` what to check.
#[allow(
	clippy::missing_docs_in_private_items,
	reason = "TODO(#2003): This exception predates us enabling `clippy::allow_attributes_without_reason`, so we don't know why it was enabled"
)]
#[derive(Debug, PartialEq, Clone, Copy)]
enum ShouldBeFallback {
	Yes,
	No,
}

/// Performs various checks on a [`OneTimeKey`].
///
/// # Errors
///
/// This function returns a [`BadJson`] error for any of the following cases:
/// - The `fallback` field is missing or not equal to `true` when the key is supposed to be a
///   fallback key.
/// - The `fallback` field is present and set to `true` when the key is supposed to be a one-time
///   key.
fn check_onetime_key(
	should_be_fallback: ShouldBeFallback,
	otk: &ruma::encryption::SignedKey,
) -> Result<(), InputJsonError> {
	if (should_be_fallback == ShouldBeFallback::Yes && !otk.fallback)
		|| (should_be_fallback == ShouldBeFallback::No && otk.fallback)
	{
		return Err(InputJsonError::BadJson(BadJson(
			"One-time key's fallback field does not match".to_owned(),
		)));
	}

	Ok(())
}

/// Handles a new device-keys upload by:
/// - Checking its correctness using [`check_device_key`]
/// - Check its signature is correct
/// - Extracts its signature(s)
/// - Store both the keys and the signatures
///
/// This function returns the keymap extracted from the given [DeviceKeys].
async fn handle_new_device_key(
	transaction: &mut TransactedDb<'_>,
	session: &SessionInfo,
	new_keys: &mut DeviceKeys,
) -> Result<(), UserKeysError> {
	check_device_key(&session.user_id, &session.device_id, new_keys)?;

	let new_keys_id = new_keys.key_id().map_err(InputJsonError::BadJson)?;
	verify_signatures(
		vec![&new_keys.keymap().map_err(InputJsonError::BadJson)?],
		new_keys,
		// Device keys MUST be signed by themselves
		Some(&BTreeMap::from([(session.user_id.clone(), vec![new_keys_id.clone()])])),
	)?;

	// SAFETY: We just checked keys were at least signed by themselves.
	#[expect(
		clippy::expect_used,
		reason = "TODO(#2003): This exception predates us enabling `clippy::allow_attributes_without_reason`, so we don't know why it was enabled"
	)]
	let signatures = new_keys
		.signatures
		.remove(&session.user_id)
		.expect("We just checked keys were at least signed by themselves");

	let device_keys_nid =
		transaction.upsert_device_keys(&session.user_id, &session.device_id, new_keys).await?;

	for (signing_key_id, signature) in signatures {
		transaction
			.upsert_local_signature(
				&session.user_id,
				// Device keys can either be signed by themselves,
				if signing_key_id == new_keys_id {
					DbKeyType::Device
				} else {
					// or by Self-Signing keys
					DbKeyType::SelfSigning
				},
				&signing_key_id,
				device_keys_nid,
				Base64::<Standard>::parse(signature)
					.unrecoverable_error("Failed to parse base64")?
					.as_bytes(),
			)
			.await?;
	}

	Ok(())
}

/// Handler for `POST /_matrix/client/v3/keys/upload` (See
/// <https://spec.matrix.org/latest/client-server-api/#post_matrixclientv3keysupload>).
///
/// Endpoint used by users to upload keys.
async fn upload_keys_handler(
	mut state: State<ServerState>,
	AuthenticatedClientRequest(claims, request): AuthenticatedClientRequest<
		upload_keys::v3::Request,
	>,
) -> Result<ArminiusResponse<upload_keys::v3::Response>, UserKeysError> {
	let mut transaction = state.database.begin_transaction().await?;

	// Uploaded device keys appear in the `device_lists.changed` field of
	// `/sync` responses of other users, so we need to notify the sync endpoint
	// about them after committing. One-time/fallback keys are not synced.
	let uploaded_device_keys = request.device_keys.is_some();

	let device_keys = match request.device_keys {
		Some(new_device_keys) => {
			let mut new_device_keys =
				new_device_keys.deserialize().map_err(InputJsonError::from)?;
			handle_new_device_key(&mut transaction, &claims, &mut new_device_keys).await?;

			Some(new_device_keys)
		}
		None => transaction
			.fetch_device_keys_for_users(&BTreeMap::from([(
				claims.user_id.clone(),
				vec![claims.device_id.clone()],
			)]))
			.await?
			.remove(&claims.user_id)
			.and_then(|mut user_keys| user_keys.remove(&claims.device_id))
			.map(|db_key| db_key.key),
	};

	if !request.one_time_keys.is_empty() || !request.fallback_keys.is_empty() {
		// If we don't have any device key at this point we won't be able to
		// verify the signatures of the new one-time and/or fallback keys so
		// we have to error out.
		let Some(device_keys) = device_keys else {
			return Err(UserKeysError::MissingDeviceKeys);
		};

		let keymap = device_keys.keymap().unrecoverable_error("Failed to get keymap")?;

		// The user's device that doing the request must have signed the other
		// keys.
		let required_signers = BTreeMap::from([(
			claims.user_id.clone(),
			vec![device_keys.key_id().map_err(InputJsonError::BadJson)?],
		)]);

		for (key_id, raw_otk) in request.one_time_keys {
			let one_time_key = raw_otk.deserialize().map_err(InputJsonError::from)?;
			let OneTimeKey::SignedKey(otk) = &one_time_key else {
				return Err(InputJsonError::BadJson(BadJson(
					"Expected one-time keys to be signed keys".to_owned(),
				))
				.into());
			};

			check_onetime_key(ShouldBeFallback::No, otk)?;
			verify_signatures(vec![&keymap], otk, Some(&required_signers))?;

			transaction
				.try_insert_one_time_key(
					&key_id.algorithm(),
					key_id.key_name(),
					&claims.user_id,
					&claims.device_id,
					&one_time_key,
				)
				.await?;
		}

		for (key_id, raw_fallback_key) in request.fallback_keys {
			let fallback_key = raw_fallback_key.deserialize().map_err(InputJsonError::from)?;
			let OneTimeKey::SignedKey(otk) = &fallback_key else {
				return Err(InputJsonError::BadJson(BadJson(
					"Expected fallback keys to be signed keys".to_owned(),
				))
				.into());
			};

			check_onetime_key(ShouldBeFallback::Yes, otk)?;
			verify_signatures(vec![&keymap], otk, Some(&required_signers))?;

			// NOTE: This will overwrite any existing fallback key with the same
			// algorithm.
			transaction
				.upsert_fallback_key(
					&key_id.algorithm(),
					key_id.key_name(),
					&claims.user_id,
					&claims.device_id,
					&fallback_key,
				)
				.await?;
		}
	}

	let one_time_key_counts =
		transaction.fetch_one_time_keys_counts(&claims.user_id, &claims.device_id).await?;

	transaction.commit().await?;

	if uploaded_device_keys {
		// Notify possibly listening clients that a user updated their device
		// keys
		state.nats_connection.notify_new_sync_data().await?;
	}

	Ok(upload_keys::v3::Response::new(one_time_key_counts).into())
}

/// Handler for `POST /_matrix/client/v3/keys/query` (See
/// <https://spec.matrix.org/latest/client-server-api/#post_matrixclientv3keysquery>).
///
/// Endpoint used to query device and user identity keys.
async fn get_keys_handler(
	mut state: State<ServerState>,
	AuthenticatedClientRequest(claims, request): AuthenticatedClientRequest<get_keys::v3::Request>,
) -> Result<ArminiusResponse<get_keys::v3::Response>, UserKeysError> {
	let (req_local_keys, req_foreign_keys) = request.device_keys.into_iter().fold(
		(BTreeMap::new(), BTreeMap::<OwnedServerName, BTreeMap<_, _>>::new()),
		|(mut local_keys, mut foreign_keys), (user_id, user_keys)| {
			let server_name = user_id.server_name();

			if server_name == state.config.server_name {
				local_keys.insert(user_id, user_keys);
			} else {
				foreign_keys.entry(server_name.to_owned()).or_default().insert(user_id, user_keys);
			}

			(local_keys, foreign_keys)
		},
	);

	let mut response = get_keys::v3::Response::new();

	// PERF: Querying local and foreign keys could be done in parallel.

	{
		// This endpoint expects us to create entries for every requested user,
		// even if we did not find their key, so we do this first.
		for user_id in req_local_keys.keys().cloned() {
			response.device_keys.entry(user_id).or_default();
		}

		let mut local_keys: RawUserKeys = arminius_lib_user_keys::identity_keys::get_local_keys(
			&mut state.database,
			Some(&claims.user_id),
			req_local_keys,
		)
		.await?
		.try_into()?;

		response.device_keys.append(&mut local_keys.device_keys);

		response.master_keys.append(&mut local_keys.master_keys);
		response.self_signing_keys.append(&mut local_keys.self_signing_keys);
		response.user_signing_keys.append(&mut local_keys.user_signing_keys);
	}

	{
		// Recommended default timeout is 10 seconds, see
		// <https://spec.matrix.org/latest/client-server-api/#post_matrixclientv3keysquery>.
		let timeout = request.timeout.unwrap_or(Duration::from_secs(10));

		let mut tasks = JoinMap::new();

		for (server_name, requested_keys) in req_foreign_keys {
			let client = FederationNatsClient::new(state.nats_connection.client.clone());
			tasks.spawn(
				server_name.clone(),
				tokio::time::timeout(timeout, async move {
					client.get_federated_user_keys(server_name, requested_keys, Some(timeout)).await
				}),
			);
		}

		while let Some((server_name, result)) = tasks.join_next().await {
			let result = result.unrecoverable_error("Failed to join spawned task")?;

			// Check if the request timed-out
			let Ok(federation_response) = result else {
				response.failures.insert(
					server_name.to_string(),
					json!({
						"reason": FederatedKeyFailureReason::TimedOut.to_string()
					}),
				);

				continue;
			};

			let federation_response =
				federation_response.unrecoverable_error("Fetching foreign keys failed")?;

			match federation_response {
				FederatedIdentityKeysResponse::Success(mut federated_keys) => {
					response.device_keys.append(&mut federated_keys.device_keys);
					response.master_keys.append(&mut federated_keys.master_keys);
					response.self_signing_keys.append(&mut federated_keys.self_signing_keys);
				}
				FederatedIdentityKeysResponse::Failure(failure) => {
					response.failures.insert(
						failure.server_name.to_string(),
						json!({
							"reason": failure.reason.to_string()
						}),
					);
				}
			}
		}
	}

	Ok(response.into())
}

/// Handler for `POST /_matrix/client/v3/keys/claim` (See
/// <https://spec.matrix.org/latest/client-server-api/#post_matrixclientv3keysclaim>).
///
/// Endpoint that allows users to claim one-time keys.
async fn claim_keys_handler(
	mut state: State<ServerState>,
	AuthenticatedClientRequest(_, request): AuthenticatedClientRequest<claim_keys::v3::Request>,
) -> Result<ArminiusResponse<claim_keys::v3::Response>, UserKeysError> {
	let (req_local_keys, req_foreign_keys): (Vec<_>, Vec<_>) = request
		.one_time_keys
		.into_iter()
		.partition(|(user_id, _)| user_id.server_name() == state.config.server_name);

	let req_local_keys = req_local_keys.into_iter().collect();
	let req_foreign_keys = req_foreign_keys.into_iter().fold(
		BTreeMap::new(),
		|mut keys: BTreeMap<OwnedServerName, BTreeMap<_, _>>, (user_id, user_keys)| {
			keys.entry(user_id.server_name().to_owned()).or_default().insert(user_id, user_keys);
			keys
		},
	);

	let mut response = claim_keys::v3::Response::new(BTreeMap::new());

	// PERF: Querying local and foreign keys could be done in parallel.

	{
		response.one_time_keys.append(
			&mut arminius_lib_user_keys::one_time_keys::claim_local_keys(
				&mut state.database,
				req_local_keys,
			)
			.await?,
		);
	}

	{
		// Recommended default timeout is 10 seconds, see
		// <https://spec.matrix.org/latest/client-server-api/#post_matrixclientv3keysclaim>.
		let timeout = request.timeout.unwrap_or(Duration::from_secs(10));

		let mut join_map = JoinMap::new();
		for (server_name, requested_keys) in req_foreign_keys {
			let client = FederationNatsClient::new(state.nats_connection.client.clone());
			join_map.spawn(
				server_name.clone(),
				tokio::time::timeout(timeout, async move {
					client
						.get_federated_otks(server_name.clone(), requested_keys, Some(timeout))
						.await
				}),
			);
		}

		while let Some((server_name, result)) = join_map.join_next().await {
			let result = result.unrecoverable_error("Failed to join spawned task")?;

			// Check if the request timed-out
			let Ok(federation_response) = result else {
				response.failures.insert(
					server_name.to_string(),
					json!({
						"reason": format!("Server {server_name} failed to respond before the timeout")
					}),
				);

				continue;
			};

			let federation_response =
				federation_response.unrecoverable_error("Fetching foreign keys failed")?;

			match federation_response {
				FederatedOTKResponse::Success(mut federated_otks) => {
					response.one_time_keys.append(&mut federated_otks.one_time_keys);
				}
				FederatedOTKResponse::Failure(failure) => {
					response.failures.insert(
						failure.server_name.to_string(),
						json!({
							"reason": failure.reason
						}),
					);
				}
			}
		}
	}

	Ok(response.into())
}

/// Component implementing endpoints for user keys
#[derive(Debug)]
pub struct UserKeysComponent;

impl Component for UserKeysComponent {
	fn router() -> Router<ServerState> {
		Router::new()
			.route("/_matrix/client/v3/keys/upload", post(upload_keys_handler))
			.route("/_matrix/client/v3/keys/query", post(get_keys_handler))
			.route("/_matrix/client/v3/keys/claim", post(claim_keys_handler))
	}
}