macltok

public

Star0

Watch0

72 commits 3 branches 0 tags 0 releases3.0 MiB

use std::{collections::HashMap, sync::Arc};

use futures_util::future::{join_all, try_join};
use matrix_sdk::{
    Room, deserialized_responses::TimelineEvent as SdkTimelineEvent, room::MessagesOptions,
};
use ruma::{OwnedEventId, OwnedUserId, UInt, push::Action};
use serde_json::Value;

use crate::{
    Error, ErrorKind, EventContext, EventContextCommand, EventContextPage, MediaId, Sender,
    TimelineEvent, TimelineItem,
    engine::MediaService,
    normalise::{
        MediaRegistration, apply_profiles, normalise_context, normalise_observed, resolve_replies,
    },
};

const MAX_CONTEXT_LIMIT: u16 = 50;

pub(crate) async fn execute(
    room: Room,
    current_user_id: Option<OwnedUserId>,
    media: Arc<MediaService>,
    command: EventContextCommand,
) -> Result<crate::CommandOutput, Error> {
    let (token, limit, backwards) = match command {
        EventContextCommand::Fetch {
            event_id,
            before_limit,
            after_limit,
            ..
        } => {
            let (context, registrations, _) =
                capture(room, current_user_id, event_id, before_limit, after_limit).await?;
            media.register(registrations).await;
            return Ok(crate::CommandOutput::EventContext(context));
        }
        EventContextCommand::Before { token, limit, .. } => (token, limit, true),
        EventContextCommand::After { token, limit, .. } => (token, limit, false),
    };
    if limit == 0 {
        return Err(Error::new(
            ErrorKind::Invalid,
            "context page limit must be at least 1",
        ));
    }
    validate_limit(limit)?;
    if token.is_empty() {
        return Err(Error::new(
            ErrorKind::Invalid,
            "context pagination token is empty",
        ));
    }
    let side = fetch_side(&room, Some(&token), limit, backwards).await?;
    let mut page = Assembly::new(
        &room,
        current_user_id.as_deref(),
        &side.values,
        &side.highlights,
        None,
    )
    .await;
    let events = page.take(0..side.values.len());
    media.register(page.media).await;
    Ok(crate::CommandOutput::EventContextPage(EventContextPage {
        events,
        next_token: side.token,
    }))
}

pub(crate) async fn capture(
    room: Room,
    current_user_id: Option<OwnedUserId>,
    event_id: OwnedEventId,
    before_limit: u16,
    after_limit: u16,
) -> Result<(EventContext, Vec<MediaRegistration>, Vec<Value>), Error> {
    validate_limit(before_limit)?;
    validate_limit(after_limit)?;
    let raw = room
        .event_with_context(&event_id, false, UInt::from(0_u16), None)
        .await
        .map_err(|source| {
            let error = Error::from_sdk("fetching event context failed", source);
            if error
                .server
                .as_ref()
                .is_some_and(|server| server.status == 404)
            {
                Error::new(ErrorKind::EventNotFound, "event was not found")
            } else {
                error
            }
        })?;
    let (target, target_highlight) = raw
        .event
        .as_ref()
        .and_then(raw_event)
        .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "event was not found"))?;
    let (before, after) = try_join(
        fetch_side(&room, raw.prev_batch_token.as_deref(), before_limit, true),
        fetch_side(&room, raw.next_batch_token.as_deref(), after_limit, false),
    )
    .await?;
    let target_index = before.values.len();
    let mut values = before.values;
    values.push(target);
    values.extend(after.values);
    let mut highlights = before.highlights;
    highlights.push(target_highlight);
    highlights.extend(after.highlights);
    let mut page = Assembly::new(
        &room,
        current_user_id.as_deref(),
        &values,
        &highlights,
        Some(target_index),
    )
    .await;
    let events_before = page.take(0..target_index);
    let event = page
        .take(target_index..target_index + 1)
        .pop()
        .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "event could not be read"))?;
    let events_after = page.take(target_index + 1..values.len());
    Ok((
        EventContext {
            events_before,
            event,
            events_after,
            start_token: before.token,
            end_token: after.token,
        },
        page.media,
        values,
    ))
}

struct Side {
    /// in timeline order
    values: Vec<Value>,
    /// beside `values`, from the push actions the sdk worked out as it fetched
    highlights: Vec<bool>,
    token: Option<String>,
}

fn raw_event(event: &SdkTimelineEvent) -> Option<(Value, bool)> {
    let value = serde_json::from_str(event.raw().json().get()).ok()?;
    let highlight = event
        .push_actions()
        .is_some_and(|actions| actions.iter().any(Action::is_highlight));
    Some((value, highlight))
}

async fn fetch_side(
    room: &Room,
    token: Option<&str>,
    limit: u16,
    backwards: bool,
) -> Result<Side, Error> {
    let Some(token) = token else {
        return Ok(Side {
            values: Vec::new(),
            highlights: Vec::new(),
            token: None,
        });
    };
    if limit == 0 {
        return Ok(Side {
            values: Vec::new(),
            highlights: Vec::new(),
            token: Some(token.to_owned()),
        });
    }
    let mut options = if backwards {
        MessagesOptions::backward()
    } else {
        MessagesOptions::forward()
    };
    options = options.from(token);
    options.limit = UInt::from(limit);
    let messages = room
        .messages(options)
        .await
        .map_err(|source| Error::from_sdk("fetching event context events failed", source))?;
    let mut events: Vec<_> = messages.chunk.iter().filter_map(raw_event).collect();
    if backwards {
        events.reverse();
    }
    let (values, highlights) = events.into_iter().unzip();
    Ok(Side {
        values,
        highlights,
        token: messages.end,
    })
}

/// a window of raw events normalised together, so relations inside it land on their targets
struct Assembly {
    items: Vec<(usize, TimelineItem)>,
    media: Vec<MediaRegistration>,
}

impl Assembly {
    async fn new(
        room: &Room,
        current_user_id: Option<&ruma::UserId>,
        values: &[Value],
        highlights: &[bool],
        keep: Option<usize>,
    ) -> Self {
        let page = normalise_context(values, keep, current_user_id, room.room_id());
        let mut items = page.items;
        for (index, item) in &mut items {
            if let TimelineItem::Event(event) = item {
                event.highlight = highlights[*index];
            }
        }
        let mut media = page.media;
        let mut targets: HashMap<OwnedEventId, TimelineEvent> = HashMap::new();
        let mut missing = Vec::new();
        for (_, item) in &items {
            if let TimelineItem::Event(event) = item
                && let Some(event_id) = &event.event_id
            {
                targets.insert(event_id.clone(), (**event).clone());
            }
        }
        for event in events(&mut items) {
            if let Some(reply) = &event.reply_to
                && !targets.contains_key(&reply.event_id)
                && !missing.contains(&reply.event_id)
            {
                missing.push(reply.event_id.clone());
            }
        }
        let fetched = join_all(
            missing
                .into_iter()
                .map(|event_id| fetch_reply(room, current_user_id, event_id)),
        )
        .await;
        for (event, registrations) in fetched.into_iter().flatten() {
            media.extend(registrations);
            if let Some(event_id) = event.event_id.clone() {
                targets.insert(event_id, event);
            }
        }
        let profiles = profiles(
            room,
            items
                .iter()
                .filter_map(|(_, item)| match item {
                    TimelineItem::Event(event) => Some(&event.sender.user_id),
                    _ => None,
                })
                .chain(targets.values().map(|event| &event.sender.user_id)),
        )
        .await;
        apply_profiles(targets.values_mut(), &profiles);
        apply_profiles(events(&mut items), &profiles);
        resolve_replies(events(&mut items), &targets);
        Self { items, media }
    }

    /// the items that came from raw events in `range`
    fn take(&mut self, range: std::ops::Range<usize>) -> Vec<TimelineItem> {
        let (taken, kept) = std::mem::take(&mut self.items)
            .into_iter()
            .partition(|(index, _)| range.contains(index));
        self.items = kept;
        taken.into_iter().map(|(_, item)| item).collect()
    }
}

fn events(items: &mut [(usize, TimelineItem)]) -> impl Iterator<Item = &mut TimelineEvent> {
    items.iter_mut().filter_map(|(_, item)| match item {
        TimelineItem::Event(event) => Some(&mut **event),
        _ => None,
    })
}

async fn fetch_reply(
    room: &Room,
    current_user_id: Option<&ruma::UserId>,
    event_id: OwnedEventId,
) -> Option<(TimelineEvent, Vec<MediaRegistration>)> {
    let event = room
        .load_or_fetch_event(&event_id, None)
        .await
        .inspect_err(|error| tracing::debug!(%event_id, %error, "fetching a reply target failed"))
        .ok()?;
    let value: Value = serde_json::from_str(event.raw().json().get()).ok()?;
    let identity = format!("{}\n{event_id}", room.room_id());
    match normalise_observed(&value, current_user_id, &identity)? {
        (TimelineItem::Event(event), media) => Some((*event, media)),
        _ => None,
    }
}

/// under lazy-loaded members a sender may be unknown locally, and `get_member` fetches the
/// member list once for the room when it is
async fn profiles<'a>(
    room: &Room,
    user_ids: impl Iterator<Item = &'a OwnedUserId>,
) -> HashMap<OwnedUserId, Sender> {
    let mut unique: Vec<&OwnedUserId> = user_ids.collect();
    unique.sort();
    unique.dedup();
    join_all(unique.into_iter().map(|user_id| async move {
        let member = room.get_member(user_id).await.ok().flatten()?;
        Some((
            user_id.clone(),
            Sender {
                user_id: user_id.clone(),
                display_name: member.display_name().map(str::to_owned),
                avatar: member.avatar_url().map(MediaId::plain),
                name_ambiguous: member.name_ambiguous(),
            },
        ))
    }))
    .await
    .into_iter()
    .flatten()
    .collect()
}

fn validate_limit(limit: u16) -> Result<(), Error> {
    if limit > MAX_CONTEXT_LIMIT {
        Err(Error::new(ErrorKind::Invalid, "context limit exceeds 50"))
    } else {
        Ok(())
    }
}

1	use std::{collections::HashMap, sync::Arc};
2
3	use futures_util::future::{join_all, try_join};
4	use matrix_sdk::{
5	Room, deserialized_responses::TimelineEvent as SdkTimelineEvent, room::MessagesOptions,
6	};
7	use ruma::{OwnedEventId, OwnedUserId, UInt, push::Action};
8	use serde_json::Value;
9
10	use crate::{
11	Error, ErrorKind, EventContext, EventContextCommand, EventContextPage, MediaId, Sender,
12	TimelineEvent, TimelineItem,
13	engine::MediaService,
14	normalise::{
15	MediaRegistration, apply_profiles, normalise_context, normalise_observed, resolve_replies,
16	},
17	};
18
19	const MAX_CONTEXT_LIMIT: u16 = 50;
20
21	pub(crate) async fn execute(
22	room: Room,
23	current_user_id: Option<OwnedUserId>,
24	media: Arc<MediaService>,
25	command: EventContextCommand,
26	) -> Result<crate::CommandOutput, Error> {
27	let (token, limit, backwards) = match command {
28	EventContextCommand::Fetch {
29	event_id,
30	before_limit,
31	after_limit,
32	..
33	} => {
34	let (context, registrations, _) =
35	capture(room, current_user_id, event_id, before_limit, after_limit).await?;
36	media.register(registrations).await;
37	return Ok(crate::CommandOutput::EventContext(context));
38	}
39	EventContextCommand::Before { token, limit, .. } => (token, limit, true),
40	EventContextCommand::After { token, limit, .. } => (token, limit, false),
41	};
42	if limit == 0 {
43	return Err(Error::new(
44	ErrorKind::Invalid,
45	"context page limit must be at least 1",
46	));
47	}
48	validate_limit(limit)?;
49	if token.is_empty() {
50	return Err(Error::new(
51	ErrorKind::Invalid,
52	"context pagination token is empty",
53	));
54	}
55	let side = fetch_side(&room, Some(&token), limit, backwards).await?;
56	let mut page = Assembly::new(
57	&room,
58	current_user_id.as_deref(),
59	&side.values,
60	&side.highlights,
61	None,
62	)
63	.await;
64	let events = page.take(0..side.values.len());
65	media.register(page.media).await;
66	Ok(crate::CommandOutput::EventContextPage(EventContextPage {
67	events,
68	next_token: side.token,
69	}))
70	}
71
72	pub(crate) async fn capture(
73	room: Room,
74	current_user_id: Option<OwnedUserId>,
75	event_id: OwnedEventId,
76	before_limit: u16,
77	after_limit: u16,
78	) -> Result<(EventContext, Vec<MediaRegistration>, Vec<Value>), Error> {
79	validate_limit(before_limit)?;
80	validate_limit(after_limit)?;
81	let raw = room
82	.event_with_context(&event_id, false, UInt::from(0_u16), None)
83	.await
84	.map_err(\|source\| {
85	let error = Error::from_sdk("fetching event context failed", source);
86	if error
87	.server
88	.as_ref()
89	.is_some_and(\|server\| server.status == 404)
90	{
91	Error::new(ErrorKind::EventNotFound, "event was not found")
92	} else {
93	error
94	}
95	})?;
96	let (target, target_highlight) = raw
97	.event
98	.as_ref()
99	.and_then(raw_event)
100	.ok_or_else(\|\| Error::new(ErrorKind::EventNotFound, "event was not found"))?;
101	let (before, after) = try_join(
102	fetch_side(&room, raw.prev_batch_token.as_deref(), before_limit, true),
103	fetch_side(&room, raw.next_batch_token.as_deref(), after_limit, false),
104	)
105	.await?;
106	let target_index = before.values.len();
107	let mut values = before.values;
108	values.push(target);
109	values.extend(after.values);
110	let mut highlights = before.highlights;
111	highlights.push(target_highlight);
112	highlights.extend(after.highlights);
113	let mut page = Assembly::new(
114	&room,
115	current_user_id.as_deref(),
116	&values,
117	&highlights,
118	Some(target_index),
119	)
120	.await;
121	let events_before = page.take(0..target_index);
122	let event = page
123	.take(target_index..target_index + 1)
124	.pop()
125	.ok_or_else(\|\| Error::new(ErrorKind::EventNotFound, "event could not be read"))?;
126	let events_after = page.take(target_index + 1..values.len());
127	Ok((
128	EventContext {
129	events_before,
130	event,
131	events_after,
132	start_token: before.token,
133	end_token: after.token,
134	},
135	page.media,
136	values,
137	))
138	}
139
140	struct Side {
141	/// in timeline order
142	values: Vec<Value>,
143	/// beside `values`, from the push actions the sdk worked out as it fetched
144	highlights: Vec<bool>,
145	token: Option<String>,
146	}
147
148	fn raw_event(event: &SdkTimelineEvent) -> Option<(Value, bool)> {
149	let value = serde_json::from_str(event.raw().json().get()).ok()?;
150	let highlight = event
151	.push_actions()
152	.is_some_and(\|actions\| actions.iter().any(Action::is_highlight));
153	Some((value, highlight))
154	}
155
156	async fn fetch_side(
157	room: &Room,
158	token: Option<&str>,
159	limit: u16,
160	backwards: bool,
161	) -> Result<Side, Error> {
162	let Some(token) = token else {
163	return Ok(Side {
164	values: Vec::new(),
165	highlights: Vec::new(),
166	token: None,
167	});
168	};
169	if limit == 0 {
170	return Ok(Side {
171	values: Vec::new(),
172	highlights: Vec::new(),
173	token: Some(token.to_owned()),
174	});
175	}
176	let mut options = if backwards {
177	MessagesOptions::backward()
178	} else {
179	MessagesOptions::forward()
180	};
181	options = options.from(token);
182	options.limit = UInt::from(limit);
183	let messages = room
184	.messages(options)
185	.await
186	.map_err(\|source\| Error::from_sdk("fetching event context events failed", source))?;
187	let mut events: Vec<_> = messages.chunk.iter().filter_map(raw_event).collect();
188	if backwards {
189	events.reverse();
190	}
191	let (values, highlights) = events.into_iter().unzip();
192	Ok(Side {
193	values,
194	highlights,
195	token: messages.end,
196	})
197	}
198
199	/// a window of raw events normalised together, so relations inside it land on their targets
200	struct Assembly {
201	items: Vec<(usize, TimelineItem)>,
202	media: Vec<MediaRegistration>,
203	}
204
205	impl Assembly {
206	async fn new(
207	room: &Room,
208	current_user_id: Option<&ruma::UserId>,
209	values: &[Value],
210	highlights: &[bool],
211	keep: Option<usize>,
212	) -> Self {
213	let page = normalise_context(values, keep, current_user_id, room.room_id());
214	let mut items = page.items;
215	for (index, item) in &mut items {
216	if let TimelineItem::Event(event) = item {
217	event.highlight = highlights[*index];
218	}
219	}
220	let mut media = page.media;
221	let mut targets: HashMap<OwnedEventId, TimelineEvent> = HashMap::new();
222	let mut missing = Vec::new();
223	for (_, item) in &items {
224	if let TimelineItem::Event(event) = item
225	&& let Some(event_id) = &event.event_id
226	{
227	targets.insert(event_id.clone(), (**event).clone());
228	}
229	}
230	for event in events(&mut items) {
231	if let Some(reply) = &event.reply_to
232	&& !targets.contains_key(&reply.event_id)
233	&& !missing.contains(&reply.event_id)
234	{
235	missing.push(reply.event_id.clone());
236	}
237	}
238	let fetched = join_all(
239	missing
240	.into_iter()
241	.map(\|event_id\| fetch_reply(room, current_user_id, event_id)),
242	)
243	.await;
244	for (event, registrations) in fetched.into_iter().flatten() {
245	media.extend(registrations);
246	if let Some(event_id) = event.event_id.clone() {
247	targets.insert(event_id, event);
248	}
249	}
250	let profiles = profiles(
251	room,
252	items
253	.iter()
254	.filter_map(\|(_, item)\| match item {
255	TimelineItem::Event(event) => Some(&event.sender.user_id),
256	_ => None,
257	})
258	.chain(targets.values().map(\|event\| &event.sender.user_id)),
259	)
260	.await;
261	apply_profiles(targets.values_mut(), &profiles);
262	apply_profiles(events(&mut items), &profiles);
263	resolve_replies(events(&mut items), &targets);
264	Self { items, media }
265	}
266
267	/// the items that came from raw events in `range`
268	fn take(&mut self, range: std::ops::Range<usize>) -> Vec<TimelineItem> {
269	let (taken, kept) = std::mem::take(&mut self.items)
270	.into_iter()
271	.partition(\|(index, _)\| range.contains(index));
272	self.items = kept;
273	taken.into_iter().map(\|(_, item)\| item).collect()
274	}
275	}
276
277	fn events(items: &mut [(usize, TimelineItem)]) -> impl Iterator<Item = &mut TimelineEvent> {
278	items.iter_mut().filter_map(\|(_, item)\| match item {
279	TimelineItem::Event(event) => Some(&mut **event),
280	_ => None,
281	})
282	}
283
284	async fn fetch_reply(
285	room: &Room,
286	current_user_id: Option<&ruma::UserId>,
287	event_id: OwnedEventId,
288	) -> Option<(TimelineEvent, Vec<MediaRegistration>)> {
289	let event = room
290	.load_or_fetch_event(&event_id, None)
291	.await
292	.inspect_err(\|error\| tracing::debug!(%event_id, %error, "fetching a reply target failed"))
293	.ok()?;
294	let value: Value = serde_json::from_str(event.raw().json().get()).ok()?;
295	let identity = format!("{}\n{event_id}", room.room_id());
296	match normalise_observed(&value, current_user_id, &identity)? {
297	(TimelineItem::Event(event), media) => Some((*event, media)),
298	_ => None,
299	}
300	}
301
302	/// under lazy-loaded members a sender may be unknown locally, and `get_member` fetches the
303	/// member list once for the room when it is
304	async fn profiles<'a>(
305	room: &Room,
306	user_ids: impl Iterator<Item = &'a OwnedUserId>,
307	) -> HashMap<OwnedUserId, Sender> {
308	let mut unique: Vec<&OwnedUserId> = user_ids.collect();
309	unique.sort();
310	unique.dedup();
311	join_all(unique.into_iter().map(\|user_id\| async move {
312	let member = room.get_member(user_id).await.ok().flatten()?;
313	Some((
314	user_id.clone(),
315	Sender {
316	user_id: user_id.clone(),
317	display_name: member.display_name().map(str::to_owned),
318	avatar: member.avatar_url().map(MediaId::plain),
319	name_ambiguous: member.name_ambiguous(),
320	},
321	))
322	}))
323	.await
324	.into_iter()
325	.flatten()
326	.collect()
327	}
328
329	fn validate_limit(limit: u16) -> Result<(), Error> {
330	if limit > MAX_CONTEXT_LIMIT {
331	Err(Error::new(ErrorKind::Invalid, "context limit exceeds 50"))
332	} else {
333	Ok(())
334	}
335	}