tuwunel_service/rooms/state_res/resolve/

power_sort.rs

use std::{collections::HashMap, iter::once};

use futures::{StreamExt, TryFutureExt, TryStreamExt, stream::FuturesUnordered};
use ruma::{
	EventId, OwnedEventId,
	events::{TimelineEventType, room::power_levels::UserPowerLevel},
	room_version_rules::RoomVersionRules,
};
use serde::Deserialize;
use tuwunel_core::{
	Result, err,
	matrix::{Event, PduEvent},
	result::NotFound,
	utils::stream::{IterStream, TryBroadbandExt, TryReadyExt},
};

use super::{
	super::{
		FetchEvent,
		events::{
			PowerEvent, RoomCreateEvent, RoomPowerLevelsEvent, RoomPowerLevelsIntField,
			power_levels::RoomPowerLevelsEventOptionExt,
		},
		fetch_event::AuthRefs,
		topological_sort,
		topological_sort::ReferencedIds,
	},
	ConflictedSet,
};

/// Enlarge the given list of conflicted power events by adding the events in
/// their auth chain that are in the full conflicted set, and sort it using
/// reverse topological power ordering.
///
/// ## Arguments
///
/// * `full_conflicted_set` - The full conflicted set.
///
/// * `rules` - The authorization rules for the current room version.
///
/// * `fetch` - Handle for reading events in the room by event ID.
///
/// ## Returns
///
/// Returns the ordered list of event IDs from earliest to latest.
#[tracing::instrument(
	level = "debug",
	skip_all,
	fields(
		conflicted = full_conflicted_set.len(),
	)
)]
pub(super) async fn power_sort(
	rules: &RoomVersionRules,
	full_conflicted_set: &ConflictedSet,
	fetch: impl FetchEvent,
) -> Result<Vec<OwnedEventId>> {
	// A representation of the DAG, a map of event ID to its list of auth events
	// that are in the full conflicted set. Fill the graph.
	let graph = full_conflicted_set
		.iter()
		.try_stream()
		.broad_and_then(async |id| {
			is_power_event_id(id, fetch)
				.map_ok(|is_power| is_power.then(|| id.clone()))
				.await
		})
		.ready_try_filter_map(Result::Ok)
		.enumerate()
		.map(|(i, event_id)| event_id.map(|event_id| (i, event_id)))
		.try_fold(HashMap::new(), |graph, (i, event_id)| {
			add_event_auth_chain(full_conflicted_set, graph, event_id, fetch, i)
		})
		.await?;

	// The map of event ID to the power level of the sender of the event.
	// Get the power level of the sender of each event in the graph.
	let event_to_power_level: HashMap<_, _> = graph
		.keys()
		.try_stream()
		.map_ok(AsRef::as_ref)
		.broad_and_then(|event_id| {
			power_level_for_sender(event_id, rules, fetch)
				.map_ok(move |sender_power| (event_id.to_owned(), sender_power))
		})
		.try_collect()
		.await?;

	let query = async |event_id: OwnedEventId| {
		let power_level = *event_to_power_level
			.get(&event_id)
			.ok_or_else(|| err!(Request(NotFound("Missing PL event: {event_id}"))))?;

		let event: PduEvent = fetch.get(&event_id).await?;

		Ok((power_level, event.origin_server_ts()))
	};

	topological_sort(graph, &query).await
}

/// Add the event with the given event ID and all the events in its auth chain
/// that are in the full conflicted set to the graph.
///
/// Missing events are skipped. Other fetch errors preserve their original
/// kind.
#[tracing::instrument(
	name = "auth_chain",
	level = "trace",
	skip_all,
	fields(
		graph = graph.len(),
		?event_id,
		%i,
	)
)]
pub(super) async fn add_event_auth_chain(
	full_conflicted_set: &ConflictedSet,
	mut graph: HashMap<OwnedEventId, ReferencedIds>,
	event_id: OwnedEventId,
	fetch: impl FetchEvent,
	i: usize,
) -> Result<HashMap<OwnedEventId, ReferencedIds>> {
	let mut todo: FuturesUnordered<_> = once(fetch_optional_event(event_id, fetch)).collect();

	while let Some(event) = todo.next().await {
		let (event_id, event): (_, Option<AuthRefs>) = event?;
		let Some(event) = event else {
			continue;
		};

		graph.entry(event_id.clone()).or_default();

		for auth_event_id in event
			.auth_events
			.into_iter()
			.filter(|auth_event_id| full_conflicted_set.contains(auth_event_id))
		{
			if !graph.contains_key(&auth_event_id) {
				todo.push(fetch_optional_event(auth_event_id.clone(), fetch));
			}

			let references = graph
				.get_mut(&event_id)
				.expect("event_id present in graph");

			if !references.contains(&auth_event_id) {
				references.push(auth_event_id);
			}
		}
	}

	Ok(graph)
}

/// Finds the power level for an event sender.
///
/// This is only valid for topological sorting. Missing events retain the
/// default-power behavior; other dependency failures are returned unchanged.
#[tracing::instrument(
	name = "sender_power",
	level = "trace",
	skip_all,
	fields(
		?event_id,
	)
)]
pub(super) async fn power_level_for_sender(
	event_id: &EventId,
	rules: &RoomVersionRules,
	fetch: impl FetchEvent,
) -> Result<UserPowerLevel> {
	let event: Option<PduEvent> = fetch.get(event_id).await.optional()?;

	let hydra_room_id = rules
		.authorization
		.room_create_event_id_as_room_id;

	let mut create_event = None;
	let mut power_levels_event = None;
	if hydra_room_id && let Some(event) = event.as_ref() {
		let create_id = event.room_id().as_event_id()?;
		let fetched: PduEvent = fetch.get(&create_id).await?;

		_ = create_event.insert(RoomCreateEvent::new(fetched));
	}

	for auth_event_id in event
		.as_ref()
		.map(Event::auth_events)
		.into_iter()
		.flatten()
	{
		use TimelineEventType::{RoomCreate, RoomPowerLevels};

		let Some(auth_event) = fetch
			.get::<PduEvent>(auth_event_id)
			.await
			.optional()?
		else {
			continue;
		};

		if !hydra_room_id && auth_event.is_type_and_state_key(&RoomCreate, "") {
			_ = create_event.get_or_insert_with(|| RoomCreateEvent::new(auth_event));
		} else if auth_event.is_type_and_state_key(&RoomPowerLevels, "") {
			_ = power_levels_event.get_or_insert_with(|| RoomPowerLevelsEvent::new(auth_event));
		}

		if power_levels_event.is_some() && create_event.is_some() {
			break;
		}
	}

	let creators = create_event
		.as_ref()
		.map(|event| event.creators(&rules.authorization))
		.transpose()?;

	if let Some((event, creators)) = event.as_ref().zip(creators) {
		power_levels_event.user_power_level(event.sender(), creators, &rules.authorization)
	} else {
		power_levels_event
			.get_as_int_or_default(RoomPowerLevelsIntField::UsersDefault, &rules.authorization)
			.map(Into::into)
	}
}

/// Whether the given event ID belongs to a power event.
///
/// See the docs of `is_power_event()` for the definition of a power event.
#[tracing::instrument(
	name = "is_power_event",
	level = "trace",
	skip_all,
	fields(
		?event_id,
	)
)]
pub(super) async fn is_power_event_id(
	event_id: &EventId,
	fetch: impl FetchEvent,
) -> Result<bool> {
	Ok(fetch
		.get(event_id)
		.await
		.optional()?
		.is_some_and(|PowerEvent(is_power)| is_power))
}

async fn fetch_optional_event<T>(
	event_id: OwnedEventId,
	fetch: impl FetchEvent,
) -> Result<(OwnedEventId, Option<T>)>
where
	T: for<'de> Deserialize<'de> + Send,
{
	let event = fetch.get::<T>(&event_id).await.optional()?;

	Ok((event_id, event))
}