tuwunel_service/rooms/state_res/

resolve.rs

#[cfg(test)]
mod tests;

mod auth_difference;
mod conflicted_subgraph;
mod iterative_auth_check;
mod mainline_sort;
mod power_sort;
mod split_conflicted;

use std::{
	collections::{BTreeMap, HashSet},
	ops::Deref,
	vec::IntoIter,
};

use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt};
use ruma::{OwnedEventId, events::StateEventType, room_version_rules::RoomVersionRules};
use tuwunel_core::{
	Result, debug,
	itertools::Itertools,
	matrix::{TypeStateKey, event_id::RandomState},
	smallvec::SmallVec,
	trace,
	utils::{
		BoolExt,
		stream::{BroadbandExt, IterStream, ReadyExt},
	},
};

use self::{
	auth_difference::auth_difference, conflicted_subgraph::conflicted_subgraph_dfs,
	iterative_auth_check::iterative_auth_check, mainline_sort::mainline_sort,
	power_sort::power_sort, split_conflicted::split_conflicted_state,
};
use super::FetchEvent;
#[cfg(test)]
use super::test_utils;

/// A mapping of event type and state_key to some value `T`, usually an
/// `EventId`.
pub type StateMap<Id> = BTreeMap<TypeStateKey, Id>;

/// Full recursive auth chain for one candidate [`StateMap`].
///
/// Values are distinct and immutable after construction. Their order is
/// arbitrary, and consumers must not depend on it.
#[derive(Clone)]
pub struct AuthSet<Id>(Vec<Id>);

/// Conflicting event ids for each contested state key.
pub type ConflictMap<Id> = StateMap<ConflictVec<Id>>;

/// Event ids contesting one state key.
///
/// Two forks disputing a key is the modal conflict, so two ids stay inline.
type ConflictVec<Id> = SmallVec<[Id; 2]>;

/// The full conflicted set (arbitrary order).
type ConflictedSet = HashSet<OwnedEventId, RandomState>;

impl<Id> AuthSet<Id> {
	/// Creates an auth set from distinct identifiers.
	///
	/// The caller must ensure `ids` contains no duplicates. Duplicates are
	/// not checked, so hot paths avoid redundant work.
	#[inline]
	#[must_use]
	pub(crate) fn from_distinct(ids: Vec<Id>) -> Self { Self(ids) }
}

impl<Id> Default for AuthSet<Id> {
	fn default() -> Self { Self(Vec::new()) }
}

impl<Id: Ord> FromIterator<Id> for AuthSet<Id> {
	fn from_iter<I: IntoIterator<Item = Id>>(iter: I) -> Self {
		Self::from_distinct(
			iter.into_iter()
				.sorted_unstable()
				.dedup()
				.collect(),
		)
	}
}

impl<Id> IntoIterator for AuthSet<Id> {
	type IntoIter = IntoIter<Id>;
	type Item = Id;

	fn into_iter(self) -> Self::IntoIter { self.0.into_iter() }
}

/// Apply the [state resolution] algorithm introduced in room version 2 to
/// resolve the state of a room.
///
/// ## Arguments
///
/// * `rules` - The rules to apply for the version of the current room.
///
/// * `state_maps` - The incoming states to resolve. Each `StateMap` represents
///   a possible fork in the state of a room.
///
/// * `auth_sets` - The list of full recursive sets of `auth_events` for each
///   event in the `state_maps`. Inputs must not contain duplicates.
///
/// * `fetch_event` - Function to fetch an event in the room given its event ID.
///
/// ## Invariants
///
/// The caller of `resolve` must ensure that all the events are from the same
/// room.
///
/// ## Returns
///
/// The resolved room state.
///
/// [state resolution]: https://spec.matrix.org/latest/rooms/v2/#state-resolution
#[tracing::instrument(level = "debug", skip_all)]
pub async fn resolve<States, AuthSets, Fetch>(
	rules: &RoomVersionRules,
	state_maps: States,
	auth_sets: AuthSets,
	fetch: Fetch,
	hydra_backports: bool,
) -> Result<StateMap<OwnedEventId>>
where
	States: Stream<Item = StateMap<OwnedEventId>> + Send,
	AuthSets: Stream<Item = AuthSet<OwnedEventId>> + Send,
	Fetch: FetchEvent,
{
	// Split the unconflicted state map and the conflicted state set.
	let (unconflicted_state, conflicted_states) = split_conflicted_state(state_maps).await;

	debug!(
		unconflicted = unconflicted_state.len(),
		conflicted_states = conflicted_states.len(),
		conflicted_events = conflicted_states
			.values()
			.fold(0_usize, |a, s| a.saturating_add(s.len())),
		"unresolved states"
	);

	trace!(
		?unconflicted_state,
		?conflicted_states,
		unconflicted = unconflicted_state.len(),
		conflicted_states = conflicted_states.len(),
		"unresolved states"
	);

	if conflicted_states.is_empty() {
		return Ok(unconflicted_state.into_iter().collect());
	}

	// 0. The full conflicted set is the union of the conflicted state set and the
	//    auth difference. Don't honor events that don't exist.
	let full_conflicted_set =
		full_conflicted_set(rules, conflicted_states, auth_sets, fetch, hydra_backports).await?;

	// 1. Select the set X of all power events that appear in the full conflicted
	//    set. For each such power event P, enlarge X by adding the events in the
	//    auth chain of P which also belong to the full conflicted set. Sort X into
	//    a list using the reverse topological power ordering.
	let sorted_power_set: Vec<_> = power_sort(rules, &full_conflicted_set, fetch)
		.inspect_ok(|list| debug!(count = list.len(), "sorted power events"))
		.inspect_ok(|list| trace!(?list, "sorted power events"))
		.boxed() // size firewall
		.await?;

	let power_set_event_ids: Vec<_> = sorted_power_set
		.iter()
		.sorted_unstable()
		.collect();

	let sorted_power_set = sorted_power_set
		.iter()
		.stream()
		.map(AsRef::as_ref);

	let begin_with_empty_state_map = rules
		.state_res
		.v2_rules()
		.is_some_and(|r| r.begin_iterative_auth_checks_with_empty_state_map)
		|| hydra_backports;

	let initial_state = begin_with_empty_state_map
		.is_false()
		.then(|| unconflicted_state.clone())
		.unwrap_or_default();

	// 2. Apply the iterative auth checks algorithm, starting from the unconflicted
	//    state map, to the list of events from the previous step to get a partially
	//    resolved state.
	// query-depth firewall
	let partially_resolved_state =
		iterative_auth_check(rules, sorted_power_set, initial_state, fetch)
			.boxed()
			.inspect_ok(|map| debug!(count = map.len(), "partially resolved power state"))
			.inspect_ok(|map| trace!(?map, "partially resolved power state"))
			.await?;

	// This "epochs" power level event
	let power_ty_sk = (StateEventType::RoomPowerLevels, "".into());
	let power_event = partially_resolved_state.get(&power_ty_sk);
	debug!(event_id = ?power_event, "epoch power event");

	let remaining_events: Vec<_> = full_conflicted_set
		.into_iter()
		.filter(|id| power_set_event_ids.binary_search(&id).is_err())
		.collect();

	debug!(count = remaining_events.len(), "remaining events");
	trace!(list = ?remaining_events, "remaining events");

	let have_remaining_events = !remaining_events.is_empty();
	let remaining_events = remaining_events
		.iter()
		.stream()
		.map(AsRef::as_ref);

	// 3. Take all remaining events that weren’t picked in step 1 and order them by
	//    the mainline ordering based on the power level in the partially resolved
	//    state obtained in step 2.
	let sorted_remaining_events = have_remaining_events
		.then_async(move || mainline_sort(power_event.cloned(), remaining_events, fetch))
		.boxed(); // size firewall

	let sorted_remaining_events = sorted_remaining_events
		.await
		.unwrap_or(Ok(Vec::new()))?;

	debug!(count = sorted_remaining_events.len(), "sorted remaining events");
	trace!(list = ?sorted_remaining_events, "sorted remaining events");

	let sorted_remaining_events = sorted_remaining_events
		.iter()
		.stream()
		.map(AsRef::as_ref);

	// 4. Apply the iterative auth checks algorithm on the partial resolved state
	//    and the list of events from the previous step.
	// query-depth firewall
	let mut resolved_state =
		iterative_auth_check(rules, sorted_remaining_events, partially_resolved_state, fetch)
			.boxed()
			.await?;

	// 5. Update the result by replacing any event with the event with the same key
	//    from the unconflicted state map, if such an event exists, to get the final
	//    resolved state.
	resolved_state.extend(unconflicted_state);

	debug!(resolved_state = resolved_state.len(), "resolved state");
	trace!(?resolved_state, "resolved state");

	Ok(resolved_state)
}

#[tracing::instrument(
	name = "conflicted",
	level = "debug",
	skip_all,
	fields(
		states = conflicted_states.len(),
		events = conflicted_states.values().flatten().count()
	),
)]
async fn full_conflicted_set<AuthSets>(
	rules: &RoomVersionRules,
	conflicted_states: ConflictMap<OwnedEventId>,
	auth_sets: AuthSets,
	fetch: impl FetchEvent,
	hydra_backports: bool,
) -> Result<ConflictedSet>
where
	AuthSets: Stream<Item = AuthSet<OwnedEventId>> + Send,
{
	let consider_conflicted_subgraph = rules
		.state_res
		.v2_rules()
		.is_some_and(|rules| rules.consider_conflicted_state_subgraph)
		|| hydra_backports;

	let conflicted_state_set: Vec<_> = conflicted_states
		.values()
		.flatten()
		.sorted_unstable()
		.dedup()
		.collect();

	// Since `org.matrix.hydra.11`, fetch the conflicted state subgraph.
	let conflicted_subgraph = consider_conflicted_subgraph
		.then_async(async || conflicted_subgraph_dfs(&conflicted_state_set, fetch))
		.map(Option::into_iter)
		.map(IterStream::stream)
		.flatten_stream()
		.flatten()
		.boxed(); // erase region

	let conflicted_state_ids = conflicted_state_set
		.iter()
		.map(Deref::deref)
		.cloned()
		.stream();

	auth_difference(auth_sets)
		.chain(conflicted_state_ids)
		.broad_then(async |id| {
			let exists = fetch.exists(&id).await?;

			Ok(exists.then_some(id))
		})
		.ready_filter_map(Result::transpose)
		.chain(conflicted_subgraph)
		.try_collect::<ConflictedSet>()
		.inspect_ok(|set| debug!(count = set.len(), "full conflicted set"))
		.inspect_ok(|set| trace!(?set, "full conflicted set"))
		.await
}