Skip to main content

tuwunel_service/rooms/pdu_metadata/
relations.rs

1use futures::{Stream, StreamExt, TryFutureExt, future::Either};
2use ruma::{
3	EventId, OwnedUserId, UserId,
4	api::Direction,
5	events::{reaction::ReactionEventContent, relation::RelationType},
6};
7use tuwunel_core::{
8	PduId,
9	arrayvec::ArrayVec,
10	implement, is_equal_to,
11	matrix::{Event, Pdu, PduCount, RawPduId, event::RelationTypeEqual},
12	result::LogErr,
13	utils::{
14		stream::{ReadyExt, TryIgnore, WidebandExt},
15		u64_from_u8,
16	},
17};
18
19use super::Service;
20use crate::rooms::short::ShortRoomId;
21
22type StartKey = ArrayVec<u8, 16>;
23
24#[implement(Service)]
25#[tracing::instrument(skip(self, from, to), level = "debug")]
26pub fn add_relation(&self, from: PduCount, to: PduCount) {
27	const BUFSIZE: usize = size_of::<u64>() * 2;
28
29	match (from, to) {
30		| (PduCount::Normal(from), PduCount::Normal(to)) => {
31			let key: &[u64] = &[to, from];
32
33			self.db
34				.tofrom_relation
35				.aput_raw::<BUFSIZE, _, _>(key, []);
36		},
37		| _ => {}, // TODO: Relations with backfilled pdus
38	}
39}
40
41/// Query relations of an event to determine if matching any of the trailing
42/// arguments. When all criteria are None the mere presence of a relation causes
43/// this function to return true.
44#[implement(Service)]
45pub async fn event_has_relation(
46	&self,
47	event_id: &EventId,
48	user_id: Option<&UserId>,
49	rel_type: Option<&RelationType>,
50	key: Option<&str>,
51) -> bool {
52	let Ok(pdu_id) = self.services.timeline.get_pdu_id(event_id).await else {
53		return false;
54	};
55
56	self.has_relation(pdu_id.into(), user_id, rel_type, key)
57		.await
58}
59
60/// Query relations of an event by PduId to determine if matching any of the
61/// trailing arguments. When all criteria are None the mere presence of a
62/// relation causes this function to return true.
63#[implement(Service)]
64pub async fn has_relation(
65	&self,
66	target: PduId,
67	user_id: Option<&UserId>,
68	rel_type: Option<&RelationType>,
69	key: Option<&str>,
70) -> bool {
71	self.get_relations(target.shortroomid, target.count, None, Direction::Forward, None)
72		.ready_filter(|(_, pdu)| user_id.is_none_or(is_equal_to!(pdu.sender())))
73		.ready_filter(|(_, pdu)| {
74			debug_assert!(
75				key.is_none() || rel_type.is_none_or(is_equal_to!(&RelationType::Annotation)),
76				"key argument only applies to Annotation type relations."
77			);
78
79			// When key is supplied we don't need to double-parse the content here and
80			// below.
81			key.is_some() || rel_type.is_none_or(|rel_type| rel_type.relation_type_equal(&pdu))
82		})
83		.ready_filter(|(_, pdu)| {
84			key.is_none_or(|key| {
85				pdu.get_content()
86					.map(|content: ReactionEventContent| content.relates_to.key == key)
87					.unwrap_or(false)
88			})
89		})
90		.ready_any(|_| true)
91		.await
92}
93
94/// MSC3440 `related_by_*`: whether any event relates to `target` with a
95/// `rel_type` in `rel_types` and a `sender` in `senders`. An empty list is
96/// unconstrained on that axis; a single relating event must satisfy both.
97#[implement(Service)]
98pub async fn has_incoming_relation(
99	&self,
100	target: PduId,
101	senders: &[OwnedUserId],
102	rel_types: &[RelationType],
103) -> bool {
104	self.get_relations(target.shortroomid, target.count, None, Direction::Forward, None)
105		.ready_any(|(_, pdu)| {
106			let sender_matches =
107				senders.is_empty() || senders.iter().any(is_equal_to!(pdu.sender()));
108
109			let rel_type_matches = rel_types.is_empty()
110				|| rel_types
111					.iter()
112					.any(|rel_type| rel_type.relation_type_equal(&pdu));
113
114			sender_matches && rel_type_matches
115		})
116		.await
117}
118
119#[implement(Service)]
120pub fn get_relations<'a>(
121	&'a self,
122	shortroomid: ShortRoomId,
123	target: PduCount,
124	from: Option<PduCount>,
125	dir: Direction,
126	user_id: Option<&'a UserId>,
127) -> impl Stream<Item = (PduCount, Pdu)> + Send + '_ {
128	let target = target.to_be_bytes();
129	let from = from
130		.map(|from| from.saturating_inc(dir))
131		.unwrap_or_else(|| match dir {
132			| Direction::Backward => PduCount::max(),
133			| Direction::Forward => PduCount::default(),
134		})
135		.to_be_bytes();
136
137	let mut buf = StartKey::new();
138	let start = {
139		buf.extend(target);
140		buf.extend(from);
141		buf.as_slice()
142	};
143
144	match dir {
145		| Direction::Backward => Either::Left(self.db.tofrom_relation.rev_raw_keys_from(start)),
146		| Direction::Forward => Either::Right(self.db.tofrom_relation.raw_keys_from(start)),
147	}
148	.ignore_err()
149	.ready_take_while(move |key| key.starts_with(&target))
150	.map(|to_from| u64_from_u8(&to_from[8..16]))
151	.map(PduCount::from_unsigned)
152	.map(move |count| (user_id, shortroomid, count))
153	.wide_filter_map(async |(user_id, shortroomid, count)| {
154		let pdu_id: RawPduId = PduId { shortroomid, count }.into();
155
156		self.services
157			.timeline
158			.get_pdu_from_id(&pdu_id)
159			.map_ok(move |mut pdu| {
160				if user_id.is_none_or(|user_id| pdu.sender() != user_id) {
161					pdu.as_mut_pdu()
162						.remove_transaction_id()
163						.log_err()
164						.ok();
165				}
166
167				(count, pdu)
168			})
169			.await
170			.ok()
171	})
172}