tuwunel_api/server/

send.rs

use std::{
	collections::BTreeMap,
	net::IpAddr,
	sync::atomic::{AtomicBool, Ordering},
	time::{Duration, Instant},
};

use axum::extract::State;
use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt};
use ruma::{
	CanonicalJsonObject, CanonicalJsonValue, MilliSecondsSinceUnixEpoch, OwnedEventId,
	OwnedRoomId, OwnedUserId, RoomId, ServerName, TransactionId, UserId,
	api::{
		error::ErrorKind,
		federation::transactions::{
			edu::{
				DeviceListUpdateContent, DirectDeviceContent, Edu, PresenceContent,
				PresenceUpdate, ReceiptContent, ReceiptData, ReceiptMap, SigningKeyUpdateContent,
				TypingContent,
			},
			send_transaction_message,
		},
	},
	events::receipt::{ReceiptEvent, ReceiptEventContent, ReceiptType},
	int,
	serde::Raw,
	to_device::DeviceIdOrAllDevices,
	uint,
};
use tuwunel_core::{
	Err, Error, Result, debug,
	debug::INFO_SPAN_LEVEL,
	debug_warn, defer, err, error,
	itertools::Itertools,
	result::LogErr,
	smallvec::SmallVec,
	trace,
	utils::{
		debug::str_truncated,
		future::TryExtExt,
		millis_since_unix_epoch,
		stream::{BroadbandExt, IterStream, ReadyExt, TryBroadbandExt, automatic_width},
	},
	warn,
};
use tuwunel_service::{
	Services,
	rooms::state_res::{is_topologically_sorted_in_place, topological_sort},
	sending::{EDU_LIMIT, PDU_LIMIT},
};

use crate::{ClientIp, Ruma};

type ResolvedMap = BTreeMap<OwnedEventId, Result>;
type RoomsPdus = SmallVec<[RoomPdus; 1]>;
type RoomPdus = (OwnedRoomId, TxnPdus);
type TxnPdus = SmallVec<[(usize, Pdu); 1]>;
type Pdu = (OwnedRoomId, OwnedEventId, CanonicalJsonObject);

/// # `PUT /_matrix/federation/v1/send/{txnId}`
///
/// Push EDUs and PDUs to this server.
#[tracing::instrument(
	name = "txn",
	level = INFO_SPAN_LEVEL,
	skip_all,
	fields(
		txn = str_truncated(body.transaction_id.as_str(), 20),
		origin = body.origin().as_str(),
		%client,
	),
)]
pub(crate) async fn send_transaction_message_route(
	State(services): State<crate::State>,
	ClientIp(client): ClientIp,
	body: Ruma<send_transaction_message::v1::Request>,
) -> Result<send_transaction_message::v1::Response> {
	if body.origin() != body.body.origin {
		return Err!(Request(Forbidden(
			"Not allowed to send transactions on behalf of other servers"
		)));
	}

	if body.pdus.len() > PDU_LIMIT {
		return Err!(Request(Forbidden(
			"Not allowed to send more than {PDU_LIMIT} PDUs in one transaction"
		)));
	}

	if body.edus.len() > EDU_LIMIT {
		return Err!(Request(Forbidden(
			"Not allowed to send more than {EDU_LIMIT} EDUs in one transaction"
		)));
	}

	let txn_start_time = Instant::now();
	trace!(
		pdus = body.pdus.len(),
		edus = body.edus.len(),
		elapsed = ?txn_start_time.elapsed(),
		"Starting txn",
	);

	let pdus = body
		.pdus
		.iter()
		.stream()
		.enumerate()
		.broad_filter_map(|(i, pdu)| {
			services
				.event_handler
				.parse_incoming_pdu(pdu)
				.inspect_err(move |e| debug_warn!("Could not parse PDU[{i}]: {e}"))
				.map_ok(move |pdu| (i, pdu))
				.ok()
		});

	let edus = body
		.edus
		.iter()
		.stream()
		.enumerate()
		.ready_filter_map(|(i, edu)| {
			serde_json::from_str(edu.json().get())
				.inspect_err(|e| debug_warn!("Could not parse EDU[{i}]: {e}"))
				.map(|edu| (i, edu))
				.ok()
		});

	let results = handle(
		&services,
		&client,
		body.origin(),
		&body.transaction_id,
		txn_start_time,
		pdus,
		edus,
	)
	.await?;

	debug!(
		pdus = body.pdus.len(),
		edus = body.edus.len(),
		elapsed = ?txn_start_time.elapsed(),
		"Finished txn",
	);

	for (id, result) in &results {
		if let Err(e) = result
			&& matches!(e, Error::BadRequest(ErrorKind::NotFound, _))
		{
			warn!("Incoming PDU failed {id}: {e:?}");
		}
	}

	Ok(send_transaction_message::v1::Response {
		pdus: results
			.into_iter()
			.map(|(e, r)| (e, r.map_err(error::sanitized_message)))
			.collect(),
	})
}

async fn handle(
	services: &Services,
	client: &IpAddr,
	origin: &ServerName,
	txn_id: &TransactionId,
	started: Instant,
	pdus: impl Stream<Item = (usize, Pdu)> + Send,
	edus: impl Stream<Item = (usize, Edu)> + Send,
) -> Result<ResolvedMap> {
	let results = handle_pdus(services, client, origin, txn_id, started, pdus).await?;

	handle_edus(services, client, origin, txn_id, edus).await?;

	Ok(results)
}

async fn handle_pdus(
	services: &Services,
	client: &IpAddr,
	origin: &ServerName,
	txn_id: &TransactionId,
	started: Instant,
	pdus: impl Stream<Item = (usize, Pdu)> + Send,
) -> Result<ResolvedMap> {
	pdus.collect()
		.map(Ok)
		.map_ok(|pdus: TxnPdus| {
			pdus.into_iter()
				.sorted_by(|(_, (room_a, ..)), (_, (room_b, ..))| room_a.cmp(room_b))
				.into_grouping_map_by(|(_, (room_id, ..))| room_id.clone())
				.collect()
				.into_iter()
				.try_stream()
		})
		.try_flatten_stream()
		.try_collect::<RoomsPdus>()
		.map_ok(IntoIterator::into_iter)
		.map_ok(IterStream::try_stream)
		.try_flatten_stream()
		.broad_and_then(async |(room_id, pdus)| {
			handle_room(services, client, origin, txn_id, started, room_id, pdus)
				.map_ok(ResolvedMap::into_iter)
				.map_ok(IterStream::try_stream)
				.await
		})
		.try_flatten()
		.try_collect()
		.await
}

#[tracing::instrument(
	name = "room",
	level = INFO_SPAN_LEVEL,
	skip_all,
	fields(%room_id)
)]
async fn handle_room(
	services: &Services,
	_client: &IpAddr,
	origin: &ServerName,
	txn_id: &TransactionId,
	txn_start_time: Instant,
	ref room_id: OwnedRoomId,
	pdus: TxnPdus,
) -> Result<ResolvedMap> {
	let pdus = sort_pdus(pdus).await;

	services
		.event_handler
		.mutex_federation
		.lock(room_id)
		.then(async |_lock| {
			pdus.into_iter()
				.enumerate()
				.try_stream()
				.and_then(async |pdu| {
					services.server.check_running().map(|()| pdu) // interruption point
				})
				.and_then(|(ri, (ti, (room_id, event_id, value)))| {
					let meta = (origin, txn_id, txn_start_time, ti);
					let pdu = (ri, (room_id, event_id, value));
					handle_pdu(services, meta, pdu).map(Ok)
				})
				.try_collect()
				.await
		})
		.await
}

/// Reorder a room's transaction PDUs so each event follows the in-batch events
/// it references. An already-ordered batch is returned unchanged; references to
/// events outside the batch are non-edges. The sort is an optimization, so a
/// failure falls back to the arrival order.
async fn sort_pdus(mut pdus: TxnPdus) -> TxnPdus {
	if already_sorted(&pdus) {
		return pdus;
	}

	let event_ids: BTreeMap<&str, &OwnedEventId> = pdus
		.iter()
		.map(|(_, (_, event_id, _))| (event_id.as_str(), event_id))
		.collect();

	let graph = pdus
		.iter()
		.map(|(_, (_, event_id, value))| {
			let references = prev_event_ids(value)
				.filter_map(|prev| event_ids.get(prev).copied())
				.map(ToOwned::to_owned)
				.collect();

			(event_id.clone(), references)
		})
		.collect();

	// Causal order alone matters here, so the tie-break inputs are constant.
	let query = async |_event_id: OwnedEventId| {
		Ok((int!(0).into(), MilliSecondsSinceUnixEpoch(uint!(0))))
	};

	let Ok(order) = topological_sort(graph, &query).await else {
		return pdus;
	};

	let position: BTreeMap<&str, usize> = order
		.iter()
		.enumerate()
		.map(|(i, event_id)| (event_id.as_str(), i))
		.collect();

	pdus.sort_by_key(|(_, (_, event_id, _))| position.get(event_id.as_str()).copied());
	pdus
}

/// Whether the batch is already in causal order, in which case the sort can be
/// skipped.
fn already_sorted(pdus: &[(usize, Pdu)]) -> bool {
	is_topologically_sorted_in_place(
		pdus,
		|(_, (_, id, _))| id.as_str(),
		|(_, (_, _, value))| prev_event_ids(value),
	)
}

/// The `prev_events` of a PDU held as canonical JSON.
fn prev_event_ids(value: &CanonicalJsonObject) -> impl Iterator<Item = &str> + '_ {
	value
		.get("prev_events")
		.and_then(CanonicalJsonValue::as_array)
		.into_iter()
		.flatten()
		.filter_map(CanonicalJsonValue::as_str)
}

#[tracing::instrument(
	name = "pdu",
	level = INFO_SPAN_LEVEL,
	skip_all,
	fields(%event_id, %ti, %ri)
)]
async fn handle_pdu(
	services: &Services,
	(origin, txn_id, txn_start_time, ti): (&ServerName, &TransactionId, Instant, usize),
	(ri, (ref room_id, event_id, value)): (usize, Pdu),
) -> (OwnedEventId, Result) {
	let pdu_start_time = Instant::now();
	let completed: AtomicBool = Default::default();
	defer! {{
		if completed.load(Ordering::Acquire) {
			return;
		}

		if pdu_start_time.elapsed() >= Duration::from_secs(services.config.client_request_timeout) {
			error!(
				%origin, %txn_id, %room_id, %event_id, %ri, %ti,
				elapsed = ?pdu_start_time.elapsed(),
				"Incoming transaction processing timed out.",
			);
		} else {
			debug_warn!(
				%origin, %txn_id, %room_id, %event_id, %ri, %ti,
				elapsed = ?pdu_start_time.elapsed(),
				"Incoming transaction processing interrupted.",
			);
		}
	}}

	let result = services
		.event_handler
		.handle_incoming_pdu(origin, room_id, &event_id, value, true)
		.map_ok(|_| ())
		.boxed()
		.await;

	completed.store(true, Ordering::Release);
	debug!(
		%event_id, ri, ti,
		pdu_elapsed = ?pdu_start_time.elapsed(),
		txn_elapsed = ?txn_start_time.elapsed(),
		"Finished PDU",
	);

	(event_id.clone(), result)
}

#[tracing::instrument(name = "edus", level = "debug", skip_all)]
async fn handle_edus(
	services: &Services,
	client: &IpAddr,
	origin: &ServerName,
	txn_id: &TransactionId,
	edus: impl Stream<Item = (usize, Edu)> + Send,
) -> Result {
	edus.for_each_concurrent(automatic_width(), |(i, edu)| {
		handle_edu(services, client, origin, txn_id, i, edu)
	})
	.await;

	Ok(())
}

#[tracing::instrument(
	name = "edu",
	level = "debug",
	skip_all,
	fields(%i),
)]
async fn handle_edu(
	services: &Services,
	client: &IpAddr,
	origin: &ServerName,
	_txn_id: &TransactionId,
	i: usize,
	edu: Edu,
) {
	match edu {
		| Edu::Presence(presence) if services.server.config.allow_incoming_presence =>
			handle_edu_presence(services, client, origin, presence).await,

		| Edu::Receipt(receipt)
			if services
				.server
				.config
				.allow_incoming_read_receipts =>
			handle_edu_receipt(services, client, origin, receipt).await,

		| Edu::Typing(typing) if services.server.config.allow_incoming_typing =>
			handle_edu_typing(services, client, origin, typing).await,

		| Edu::DeviceListUpdate(content) =>
			handle_edu_device_list_update(services, client, origin, content).await,

		| Edu::DirectToDevice(content) =>
			handle_edu_direct_to_device(services, client, origin, content).await,

		| Edu::SigningKeyUpdate(content) =>
			handle_edu_signing_key_update(services, client, origin, content).await,

		| Edu::_Custom(ref _custom) => debug_warn!(?i, ?edu, "received custom/unknown EDU"),

		| _ => trace!(?i, ?edu, "skipped"),
	}
}

async fn handle_edu_presence(
	services: &Services,
	_client: &IpAddr,
	origin: &ServerName,
	presence: PresenceContent,
) {
	presence
		.push
		.into_iter()
		.stream()
		.for_each_concurrent(automatic_width(), |update| {
			handle_edu_presence_update(services, origin, update)
		})
		.await;
}

async fn handle_edu_presence_update(
	services: &Services,
	origin: &ServerName,
	update: PresenceUpdate,
) {
	if update.user_id.server_name() != origin {
		debug_warn!(
			%update.user_id, %origin,
			"received presence EDU for user not belonging to origin"
		);
		return;
	}

	services
		.presence
		.set_presence_from_federation(
			&update.user_id,
			&update.presence,
			update.currently_active,
			update.last_active_ago,
			update.status_msg.clone(),
		)
		.await
		.log_err()
		.ok();
}

async fn handle_edu_receipt(
	services: &Services,
	_client: &IpAddr,
	origin: &ServerName,
	receipt: ReceiptContent,
) {
	receipt
		.receipts
		.into_iter()
		.stream()
		.for_each_concurrent(automatic_width(), |(room_id, room_updates)| {
			handle_edu_receipt_room(services, origin, room_id, room_updates)
		})
		.await;
}

async fn handle_edu_receipt_room(
	services: &Services,
	origin: &ServerName,
	room_id: OwnedRoomId,
	room_updates: ReceiptMap,
) {
	if services
		.event_handler
		.acl_check(origin, &room_id)
		.await
		.is_err()
	{
		debug_warn!(
			%origin, %room_id,
			"received read receipt EDU from ACL'd server"
		);
		return;
	}

	let room_id = &room_id;
	room_updates
		.read
		.into_iter()
		.stream()
		.for_each_concurrent(automatic_width(), async |(user_id, user_updates)| {
			handle_edu_receipt_room_user(services, origin, room_id, &user_id, user_updates).await;
		})
		.await;
}

async fn handle_edu_receipt_room_user(
	services: &Services,
	origin: &ServerName,
	room_id: &RoomId,
	user_id: &UserId,
	user_updates: ReceiptData,
) {
	if user_id.server_name() != origin {
		debug_warn!(
			%user_id, %origin,
			"received read receipt EDU for user not belonging to origin"
		);
		return;
	}

	if !services
		.state_cache
		.server_in_room(origin, room_id)
		.await
	{
		debug_warn!(
			%user_id, %room_id, %origin,
			"received read receipt EDU from server who does not have a member in the room",
		);
		return;
	}

	let data = &user_updates.data;
	user_updates
		.event_ids
		.into_iter()
		.stream()
		.for_each_concurrent(automatic_width(), async |event_id| {
			let user_data = [(user_id.to_owned(), data.clone())];
			let receipts = [(ReceiptType::Read, BTreeMap::from(user_data))];
			let content = [(event_id.clone(), BTreeMap::from(receipts))];
			services
				.read_receipt
				.readreceipt_update(user_id, room_id, &ReceiptEvent {
					content: ReceiptEventContent(content.into()),
					room_id: room_id.to_owned(),
				})
				.await;
		})
		.await;
}

async fn handle_edu_typing(
	services: &Services,
	_client: &IpAddr,
	origin: &ServerName,
	typing: TypingContent,
) {
	if typing.user_id.server_name() != origin {
		debug_warn!(
			%typing.user_id, %origin,
			"received typing EDU for user not belonging to origin"
		);
		return;
	}

	if services
		.event_handler
		.acl_check(typing.user_id.server_name(), &typing.room_id)
		.await
		.is_err()
	{
		debug_warn!(
			%typing.user_id, %typing.room_id, %origin,
			"received typing EDU for ACL'd user's server"
		);
		return;
	}

	if !services
		.state_cache
		.is_joined(&typing.user_id, &typing.room_id)
		.await
	{
		debug_warn!(
			%typing.user_id, %typing.room_id, %origin,
			"received typing EDU for user not in room"
		);
		return;
	}

	if typing.typing {
		let secs = services.server.config.typing_federation_timeout_s;
		let timeout = millis_since_unix_epoch().saturating_add(secs.saturating_mul(1000));

		services
			.typing
			.typing_add(&typing.user_id, &typing.room_id, timeout)
			.await
			.log_err()
			.ok();
	} else {
		services
			.typing
			.typing_remove(&typing.user_id, &typing.room_id)
			.await
			.log_err()
			.ok();
	}
}

async fn handle_edu_device_list_update(
	services: &Services,
	_client: &IpAddr,
	origin: &ServerName,
	content: DeviceListUpdateContent,
) {
	let DeviceListUpdateContent { user_id, .. } = content;

	if user_id.server_name() != origin {
		debug_warn!(
			%user_id, %origin,
			"received device list update EDU for user not belonging to origin"
		);
		return;
	}

	services
		.users
		.mark_device_key_update(&user_id)
		.await;
}

async fn handle_edu_direct_to_device(
	services: &Services,
	_client: &IpAddr,
	origin: &ServerName,
	content: DirectDeviceContent,
) {
	let DirectDeviceContent {
		ref sender,
		ref ev_type,
		ref message_id,
		messages,
	} = content;

	if sender.server_name() != origin {
		debug_warn!(
			%sender, %origin,
			"received direct to device EDU for user not belonging to origin"
		);
		return;
	}

	// Check if this is a new transaction id
	if services
		.transaction_ids
		.existing_txnid(sender, None, message_id)
		.await
		.is_ok()
	{
		return;
	}

	// process messages concurrently for different users
	let ev_type = ev_type.to_string();
	messages
		.into_iter()
		.stream()
		.for_each_concurrent(automatic_width(), |(target_user_id, map)| {
			handle_edu_direct_to_device_user(services, target_user_id, sender, &ev_type, map)
		})
		.await;

	// Save transaction id with empty data
	services
		.transaction_ids
		.add_txnid(sender, None, message_id, &[]);
}

async fn handle_edu_direct_to_device_user<Event: Send + Sync>(
	services: &Services,
	target_user_id: OwnedUserId,
	sender: &UserId,
	ev_type: &str,
	map: BTreeMap<DeviceIdOrAllDevices, Raw<Event>>,
) {
	map.into_iter()
		.stream()
		.ready_filter_map(|(tid, raw)| {
			raw.deserialize_as()
				.map_err(|e| {
					err!(Request(InvalidParam(error!("To-Device event is invalid: {e}"))))
				})
				.ok()
				.map(|ev| (tid, ev))
		})
		.for_each_concurrent(automatic_width(), |(tid, ev)| {
			handle_edu_direct_to_device_event(services, &target_user_id, sender, tid, ev_type, ev)
		})
		.await;
}

async fn handle_edu_direct_to_device_event(
	services: &Services,
	target_user_id: &UserId,
	sender: &UserId,
	target_device_id_maybe: DeviceIdOrAllDevices,
	ev_type: &str,
	event: serde_json::Value,
) {
	match target_device_id_maybe {
		| DeviceIdOrAllDevices::DeviceId(ref target_device_id) => {
			services.users.add_to_device_event(
				sender,
				target_user_id,
				target_device_id,
				ev_type,
				&event,
			);
		},

		| DeviceIdOrAllDevices::AllDevices => {
			services
				.users
				.all_device_ids(target_user_id)
				.ready_for_each(|target_device_id| {
					services.users.add_to_device_event(
						sender,
						target_user_id,
						target_device_id,
						ev_type,
						&event,
					);
				})
				.await;
		},
	}
}

async fn handle_edu_signing_key_update(
	services: &Services,
	_client: &IpAddr,
	origin: &ServerName,
	content: SigningKeyUpdateContent,
) {
	let SigningKeyUpdateContent { user_id, master_key, self_signing_key } = content;

	if user_id.server_name() != origin {
		debug_warn!(
			%user_id, %origin,
			"received signing key update EDU from server that does not belong to user's server"
		);
		return;
	}

	services
		.users
		.add_cross_signing_keys(&user_id, &master_key, &self_signing_key, &None, true)
		.await
		.log_err()
		.ok();
}

#[cfg(test)]
mod tests {
	use ruma::{CanonicalJsonObject, OwnedEventId, event_id, room_id};
	use serde_json::json;

	use super::{Pdu, TxnPdus, already_sorted, prev_event_ids, sort_pdus};

	fn pdu(index: usize, id: &OwnedEventId, prev: &[&OwnedEventId]) -> (usize, Pdu) {
		let prev_events: Vec<&str> = prev.iter().map(|e| e.as_str()).collect();
		let value: CanonicalJsonObject =
			serde_json::from_value(json!({ "prev_events": prev_events }))
				.expect("valid canonical json");

		(index, (room_id!("!r:example.com").to_owned(), id.clone(), value))
	}

	fn ids() -> (OwnedEventId, OwnedEventId, OwnedEventId) {
		(
			event_id!("$a:example.com").to_owned(),
			event_id!("$b:example.com").to_owned(),
			event_id!("$c:example.com").to_owned(),
		)
	}

	fn order(pdus: &[(usize, Pdu)]) -> Vec<&str> {
		pdus.iter()
			.map(|(_, (_, id, _))| id.as_str())
			.collect()
	}

	#[test]
	fn sorted_when_parents_lead() {
		let (a, b, c) = ids();
		let pdus = [pdu(0, &a, &[]), pdu(1, &b, &[&a]), pdu(2, &c, &[&b])];

		assert!(already_sorted(&pdus));
	}

	#[test]
	fn unsorted_when_child_leads() {
		let (a, b, _c) = ids();
		let pdus = [pdu(0, &b, &[&a]), pdu(1, &a, &[])];

		assert!(!already_sorted(&pdus));
	}

	#[test]
	fn sorted_ignores_out_of_batch_references() {
		let (a, b, c) = ids();
		let pdus = [pdu(0, &b, &[&c]), pdu(1, &a, &[&c])];

		assert!(already_sorted(&pdus));
	}

	#[tokio::test]
	async fn sort_orders_parents_before_children() {
		let (a, b, c) = ids();
		let pdus: TxnPdus = [pdu(0, &c, &[&b]), pdu(1, &b, &[&a]), pdu(2, &a, &[])]
			.into_iter()
			.collect();

		let sorted = sort_pdus(pdus).await;

		assert_eq!(order(&sorted), ["$a:example.com", "$b:example.com", "$c:example.com"]);
	}

	#[tokio::test]
	async fn sort_is_noop_when_already_ordered() {
		let (a, b, c) = ids();
		let pdus: TxnPdus = [pdu(0, &a, &[]), pdu(1, &b, &[&a]), pdu(2, &c, &[&b])]
			.into_iter()
			.collect();

		let sorted = sort_pdus(pdus.clone()).await;

		assert_eq!(order(&sorted), order(&pdus));
	}

	#[tokio::test]
	async fn sort_preserves_duplicates() {
		let (a, b, _c) = ids();
		let pdus: TxnPdus = [pdu(0, &b, &[&a]), pdu(1, &a, &[]), pdu(2, &b, &[&a])]
			.into_iter()
			.collect();

		let sorted = sort_pdus(pdus).await;

		assert_eq!(sorted.len(), 3);
	}

	#[tokio::test]
	async fn sort_preserves_a_cycle() {
		let (a, b, _c) = ids();
		let pdus: TxnPdus = [pdu(0, &a, &[&b]), pdu(1, &b, &[&a])]
			.into_iter()
			.collect();

		let sorted = sort_pdus(pdus).await;

		assert_eq!(sorted.len(), 2);
	}

	#[test]
	fn prev_event_ids_reads_the_array() {
		let (a, b, _c) = ids();
		let (_, (_, _, value)) = pdu(0, &a, &[&b]);

		let prev: Vec<&str> = prev_event_ids(&value).collect();

		assert_eq!(prev, ["$b:example.com"]);
	}

	#[test]
	fn prev_event_ids_empty_when_absent() {
		let value = CanonicalJsonObject::new();

		assert_eq!(prev_event_ids(&value).count(), 0);
	}
}