tuwunel_service/rooms/pdu_metadata/

bundling.rs

use std::collections::BTreeSet;

use futures::{Stream, StreamExt, TryFutureExt, pin_mut};
use ruma::{OwnedUserId, UserId, api::Direction, events::room::encrypted::Relation};
use tuwunel_core::{
	PduId,
	arrayvec::ArrayVec,
	implement,
	matrix::{Event, Pdu, PduCount, RawPduId},
	result::LogErr,
	utils::{
		BoolExt,
		stream::{ReadyExt, TryIgnore},
		u64_from_u8,
	},
};

use super::{
	ExtractRelatesTo, IgnoredThreadView,
	IgnoredThreadView::{Adjusted, Unchanged, WithoutSummary},
	Service,
	typed_relations::{CHILD_COUNT_OFFSET, KEY_LEN, Tag, prefix},
};

type Seek = ArrayVec<u8, KEY_LEN>;

/// Fold read-time bundled aggregations into a served event's `unsigned`,
/// per-requester. MSC3816: the stored `m.thread` bundle carries a shared
/// `current_user_participated`, recomputed here for `sender_user`. MSC3925:
/// when `bundle_edit_relations` is enabled, the newest `m.replace` edit is
/// folded in as the full replacement event, and the bundled thread
/// `latest_event` carries its own newest edit (MSC3856). MSC3267: when
/// `bundle_reference_relations` is enabled, the `m.reference` children are
/// folded in as a `{ chunk: [{ event_id }] }` summary. The thread presence gate
/// keeps the common no-bundle case to a substring scan; the edit and reference
/// folds are skipped unless enabled.
#[implement(Service)]
#[tracing::instrument(skip_all, level = "trace")]
pub async fn bundle_aggregations(&self, sender_user: &UserId, mut pdu: Pdu) -> Pdu {
	// MSC4025: an erased sender's event serves as the pruned clone, and a
	// pruned event carries no aggregations.
	if let Some(pruned) = self
		.services
		.state_accessor
		.erased_view(sender_user, &pdu)
		.await
	{
		return pruned;
	}

	let has_thread = pdu.has_thread_bundle().log_err().unwrap_or(true);

	if has_thread {
		if pdu
			.remove_thread_latest_transaction_id_unless_sender(sender_user)
			.log_err()
			.is_err()
		{
			drop_thread_bundle(&mut pdu);
		} else {
			let participated = self
				.services
				.threads
				.user_participated(pdu.event_id(), sender_user)
				.await;

			let participation = pdu
				.set_thread_participated(participated)
				.log_err();

			if thread_result_or_drop(&mut pdu, participation).is_some() {
				self.erase_thread_latest(sender_user, &mut pdu)
					.await;

				if self.services.server.config.bundle_edit_relations {
					self.bundle_thread_latest_edit(sender_user, &mut pdu)
						.await;
				}
			}
		}
	}

	let replacement = self
		.services
		.server
		.config
		.bundle_edit_relations
		.then_async(|| self.newest_replacement(&pdu))
		.await
		.flatten();

	if let Some(mut replacement) = replacement
		&& !self
			.services
			.state_accessor
			.erased_for(sender_user, &replacement)
			.await
		&& replacement
			.remove_transaction_id_unless_sender(Some(sender_user))
			.log_err()
			.is_ok()
	{
		pdu.set_replacement_bundle(&replacement.into_format())
			.log_err()
			.ok();
	}

	let references = self
		.services
		.server
		.config
		.bundle_reference_relations
		.then_async(|| self.references(&pdu))
		.await
		.unwrap_or_default();

	if !references.is_empty() {
		pdu.set_reference_bundle(&references)
			.log_err()
			.ok();
	}

	pdu
}

/// MSC4025: the stored thread bundle carries a full `latest_event` of any
/// sender; an erased hit swaps in the pruned form for this recipient. The
/// event load and membership check run only on the erased hit.
#[implement(Service)]
#[tracing::instrument(skip_all, level = "trace")]
async fn erase_thread_latest(&self, sender_user: &UserId, pdu: &mut Pdu) {
	let identity = pdu.thread_latest_event().log_err();

	let Some((event_id, sender)) = thread_result_or_drop(pdu, identity).flatten() else {
		return;
	};

	if !self.services.users.is_erased(&sender).await {
		return;
	}

	let latest = self.services.timeline.get_pdu(&event_id).await;
	let Some(latest) = thread_result_or_drop(pdu, latest) else {
		return;
	};

	let Some(pruned) = self
		.services
		.state_accessor
		.erased_view(sender_user, &latest)
		.await
	else {
		return;
	};

	if pdu
		.set_thread_latest_event(&pruned.into_format())
		.log_err()
		.is_err()
	{
		drop_thread_bundle(pdu);
	}
}

fn thread_result_or_drop<T, E>(pdu: &mut Pdu, result: Result<T, E>) -> Option<T> {
	result
		.inspect_err(|_| drop_thread_bundle(pdu))
		.ok()
}

fn drop_thread_bundle(pdu: &mut Pdu) { pdu.remove_thread_bundle().log_err().ok(); }

/// The thread module's aggregated `latest_event` (MSC3856): when the edit
/// fold is enabled, the bundled latest reply carries its own newest
/// `m.replace` edit, so thread previews track edits. Erased-sender bundles
/// stay in their pruned form.
#[implement(Service)]
#[tracing::instrument(skip_all, level = "trace")]
async fn bundle_thread_latest_edit(&self, sender_user: &UserId, pdu: &mut Pdu) {
	let identity = pdu.thread_latest_event().log_err();

	let Some((event_id, _)) = thread_result_or_drop(pdu, identity).flatten() else {
		return;
	};

	let latest_event = self
		.services
		.timeline
		.get_pdu(&event_id)
		.await
		.log_err();

	let Some(mut latest_event) = thread_result_or_drop(pdu, latest_event) else {
		return;
	};

	if self
		.services
		.state_accessor
		.erased_for(sender_user, &latest_event)
		.await
	{
		return;
	}

	let Some(mut replacement_event) = self.newest_replacement(&latest_event).await else {
		return;
	};

	if self
		.services
		.state_accessor
		.erased_for(sender_user, &replacement_event)
		.await
	{
		return;
	}

	let sanitized = replacement_event
		.remove_transaction_id_unless_sender(Some(sender_user))
		.and_then(|()| latest_event.remove_transaction_id_unless_sender(Some(sender_user)))
		.log_err();

	if thread_result_or_drop(pdu, sanitized).is_none() {
		return;
	}

	let replacement = latest_event
		.set_replacement_bundle(&replacement_event.into_format())
		.log_err();

	if thread_result_or_drop(pdu, replacement).is_none() {
		return;
	}

	let latest = pdu
		.set_thread_latest_event(&latest_event.into_format())
		.log_err();

	thread_result_or_drop(pdu, latest);
}

/// MSC3925: the newest `m.replace` edit of `parent` as a full event, or `None`
/// when `parent` is redacted or has no valid edit. An edit counts only when it
/// shares the parent's sender and type and is not itself redacted; newest is by
/// `origin_server_ts`, which the typed index sorts on.
#[implement(Service)]
#[tracing::instrument(skip_all, level = "trace")]
async fn newest_replacement(&self, parent: &Pdu) -> Option<Pdu> {
	if parent.is_redacted() {
		return None;
	}

	let parent_id: PduId = self
		.services
		.timeline
		.get_pdu_id(parent.event_id())
		.map_ok(Into::into)
		.await
		.ok()?;

	let replacements = self.replacement_children(parent, parent_id);

	pin_mut!(replacements);
	replacements.next().await
}

/// Stream `parent`'s valid `m.replace` children, newest `origin_server_ts`
/// first, from the typed index. A child counts only when it shares the parent's
/// sender and type and is not itself redacted.
#[implement(Service)]
fn replacement_children<'a>(
	&'a self,
	parent: &'a Pdu,
	parent_id: PduId,
) -> impl Stream<Item = Pdu> + Send + 'a {
	let shortroomid = parent_id.shortroomid;
	let prefix = prefix(shortroomid, parent_id.count, Tag::Replace);

	let mut seek = Seek::new();

	seek.extend(prefix.iter().copied());
	seek.extend([u8::MAX; size_of::<u64>() * 2]);

	self.db
		.relatesto_typed
		.rev_raw_keys_from(seek.as_slice())
		.ignore_err()
		.ready_take_while(move |key| key.starts_with(&prefix))
		.map(|key| u64_from_u8(&key[CHILD_COUNT_OFFSET..KEY_LEN]))
		.map(PduCount::from_unsigned)
		.map(move |count| (shortroomid, count))
		.filter_map(async |(shortroomid, count)| {
			let child_id: RawPduId = PduId { shortroomid, count }.into();
			self.services
				.timeline
				.get_pdu_from_id(&child_id)
				.await
				.ok()
				.filter(|child| !child.is_redacted())
				.filter(|child| child.sender() == parent.sender())
				.filter(|child| child.kind() == parent.kind())
		})
}

/// Evaluate one thread root against the requester's ignore list.
///
/// A cheap participant intersection gates the reply walk. One walk then yields
/// the replacement `latest_event`, the ignored-aware `count`, and the
/// summary-omission verdict when every reply is ignored. A root whose replies
/// are not indexed adjusts nothing beyond its own redacted form.
#[implement(Service)]
#[tracing::instrument(skip_all, level = "trace")]
pub async fn ignored_thread_view(
	&self,
	sender_user: &UserId,
	ignored: &BTreeSet<OwnedUserId>,
	root: &Pdu,
) -> IgnoredThreadView {
	let Ok(root_id) = self
		.services
		.timeline
		.get_pdu_id(root.event_id())
		.await
	else {
		return Unchanged;
	};

	let participants = self
		.services
		.threads
		.get_participants(&root_id)
		.await
		.unwrap_or_default();

	if !participants
		.iter()
		.any(|user| ignored.contains(user))
	{
		return Unchanged;
	}

	let root_pid: PduId = root_id.into();
	let replies = self
		.get_relations(
			root_pid.shortroomid,
			root_pid.count,
			None,
			Direction::Backward,
			Some(sender_user),
		)
		.ready_filter_map(|(_, pdu)| {
			pdu.get_content()
				.is_ok_and(|content: ExtractRelatesTo| {
					matches!(content.relates_to, Relation::Thread(_))
				})
				.then_some(pdu)
		});

	let fold = |(total, unignored, latest): (usize, usize, Option<Pdu>), pdu: Pdu| match ignored
		.contains(pdu.sender())
	{
		| true => (total.saturating_add(1), unignored, latest),
		| false => (total.saturating_add(1), unignored.saturating_add(1), latest.or(Some(pdu))),
	};

	let (total, unignored, latest) = replies.ready_fold((0, 0, None), fold).await;

	if total == 0 {
		return match self.redacted_root(ignored, root).await {
			| None => Unchanged,
			| root => Adjusted { root, count: None, latest: None },
		};
	}

	if unignored == 0 {
		return WithoutSummary {
			root: self.redacted_root(ignored, root).await,
		};
	}

	let swap = root
		.thread_latest_event()
		.ok()
		.flatten()
		.is_some_and(|(_, sender)| ignored.contains(&sender));

	let latest = match swap.then_some(latest).flatten() {
		| None => None,
		| Some(reply) => {
			// MSC4025: the swapped-in reply must not reopen the erased-sender
			// seam the bundle pass gates on the stored latest.
			let reply = self
				.services
				.state_accessor
				.erased_view(sender_user, &reply)
				.await
				.unwrap_or(reply);

			Some(reply.into_format())
		},
	};

	let count = unignored.ne(&total).then_some(unignored);

	let root = self.redacted_root(ignored, root).await;

	if root.is_none() && count.is_none() && latest.is_none() {
		return Unchanged;
	}

	Adjusted { root, count, latest }
}

/// The spec'd redacted form of an ignored sender's thread root, content side
/// only; `None` when the sender is not ignored, or on a redaction failure
/// (serving unredacted then matches the reference implementation).
#[implement(Service)]
#[tracing::instrument(skip_all, level = "trace")]
async fn redacted_root(&self, ignored: &BTreeSet<OwnedUserId>, root: &Pdu) -> Option<Box<Pdu>> {
	ignored
		.contains(root.sender())
		.then_async(async || {
			self.services
				.state
				.get_room_version_rules(root.room_id())
				.await
				.log_err()
				.ok()
				.and_then(|rules| root.redacted(&rules.redaction).log_err().ok())
				.map(Box::new)
		})
		.await
		.flatten()
}

#[cfg(test)]
mod tests {
	use serde_json::json;

	use super::*;

	#[test]
	fn thread_latest_load_error_drops_bundle() {
		let mut pdu: Pdu = serde_json::from_value(json!({
			"type": "m.room.member",
			"content": { "membership": "join" },
			"event_id": "$member:example.com",
			"room_id": "!room:example.com",
			"sender": "@alice:example.com",
			"state_key": "@alice:example.com",
			"prev_events": ["$prev:example.com"],
			"auth_events": ["$auth:example.com"],
			"origin_server_ts": 1_838_188_000,
			"depth": 12,
			"hashes": { "sha256": "thishashcoversallfieldsincasethisisredacted" },
			"unsigned": {
				"age": 4612,
				"m.relations": {
					"m.replace": { "event_id": "$edit:example.com" },
					"m.thread": { "count": 3 },
				},
			},
		}))
		.expect("test fixture should deserialize as a valid PDU");

		let latest: Option<()> = thread_result_or_drop(&mut pdu, Err("missing latest event"));

		assert!(latest.is_none(), "load failure returned a latest event");

		let unsigned: serde_json::Value = serde_json::from_str(
			pdu.unsigned
				.as_ref()
				.expect("sibling unsigned data should remain")
				.json()
				.get(),
		)
		.expect("remaining unsigned data should be valid JSON");

		assert!(
			unsigned["m.relations"].get("m.thread").is_none(),
			"failed load retained the thread bundle",
		);

		assert_eq!(
			unsigned["m.relations"]["m.replace"],
			json!({ "event_id": "$edit:example.com" }),
			"failed load removed a sibling relation",
		);
		assert_eq!(unsigned["age"], 4612, "failed load removed outer unsigned data");
	}
}