| 1 | use std::{ |
| 2 | collections::{HashMap, VecDeque}, |
| 3 | sync::Arc, |
| 4 | }; |
| 5 | |
| 6 | use futures_util::future::{join_all, try_join}; |
| 7 | use matrix_sdk::{ |
| 8 | Client, Room, RoomState, deserialized_responses::TimelineEvent as SdkTimelineEvent, |
| 9 | room::MessagesOptions, |
| 10 | }; |
| 11 | use ruma::{OwnedEventId, OwnedRoomId, OwnedUserId, UInt, push::Action}; |
| 12 | use serde_json::Value; |
| 13 | use tokio::sync::{Mutex, broadcast::error::RecvError}; |
| 14 | |
| 15 | use crate::{ |
| 16 | Error, ErrorKind, EventContext, EventContextCommand, EventContextPage, EventContextSession, |
| 17 | MediaId, Sender, TimelineEvent, TimelineItem, |
| 18 | engine::MediaService, |
| 19 | normalise::{ |
| 20 | MediaRegistration, MediaScope, apply_profiles, normalise_context, normalise_observed, |
| 21 | resolve_replies, |
| 22 | }, |
| 23 | }; |
| 24 | |
| 25 | const MAX_CONTEXT_LIMIT: u16 = 50; |
| 26 | /// a consumer that never closes its jump views still holds at most this many |
| 27 | const MAX_SESSIONS: usize = 32; |
| 28 | |
| 29 | /// open event contexts, oldest first, each owning the media its pages registered |
| 30 | #[derive(Default)] |
| 31 | pub(crate) struct Sessions { |
| 32 | open: Mutex<VecDeque<Session>>, |
| 33 | } |
| 34 | |
| 35 | struct Session { |
| 36 | id: String, |
| 37 | room_id: OwnedRoomId, |
| 38 | scope: MediaScope, |
| 39 | } |
| 40 | |
| 41 | impl Sessions { |
| 42 | async fn open( |
| 43 | &self, |
| 44 | media: &MediaService, |
| 45 | room_id: OwnedRoomId, |
| 46 | registrations: Vec<MediaRegistration>, |
| 47 | ) -> String { |
| 48 | let scope = media.context_scope(); |
| 49 | let id = format!("context-{}", scope.nonce()); |
| 50 | let mut open = self.open.lock().await; |
| 51 | media.register(&scope, registrations).await; |
| 52 | open.push_back(Session { |
| 53 | id: id.clone(), |
| 54 | room_id, |
| 55 | scope, |
| 56 | }); |
| 57 | while open.len() > MAX_SESSIONS { |
| 58 | if let Some(oldest) = open.pop_front() { |
| 59 | media.release(&oldest.scope).await; |
| 60 | } |
| 61 | } |
| 62 | id |
| 63 | } |
| 64 | |
| 65 | async fn room(&self, id: &str) -> Result<OwnedRoomId, Error> { |
| 66 | let open = self.open.lock().await; |
| 67 | let session = open.iter().find(|session| session.id == id); |
| 68 | session |
| 69 | .map(|session| session.room_id.clone()) |
| 70 | .ok_or_else(not_open) |
| 71 | } |
| 72 | |
| 73 | /// under the lock, so a page that lands after its close registers nothing |
| 74 | async fn register( |
| 75 | &self, |
| 76 | media: &MediaService, |
| 77 | id: &str, |
| 78 | registrations: Vec<MediaRegistration>, |
| 79 | ) -> Result<(), Error> { |
| 80 | let open = self.open.lock().await; |
| 81 | let session = open |
| 82 | .iter() |
| 83 | .find(|session| session.id == id) |
| 84 | .ok_or_else(not_open)?; |
| 85 | media.register(&session.scope, registrations).await; |
| 86 | Ok(()) |
| 87 | } |
| 88 | |
| 89 | async fn close(&self, media: &MediaService, id: &str) -> Result<(), Error> { |
| 90 | let mut open = self.open.lock().await; |
| 91 | let index = open |
| 92 | .iter() |
| 93 | .position(|session| session.id == id) |
| 94 | .ok_or_else(not_open)?; |
| 95 | if let Some(session) = open.remove(index) { |
| 96 | media.release(&session.scope).await; |
| 97 | } |
| 98 | Ok(()) |
| 99 | } |
| 100 | |
| 101 | async fn close_left(&self, media: &MediaService, client: &Client) { |
| 102 | let mut open = self.open.lock().await; |
| 103 | let mut kept = VecDeque::with_capacity(open.len()); |
| 104 | for session in open.drain(..) { |
| 105 | let left = client |
| 106 | .get_room(&session.room_id) |
| 107 | .is_none_or(|room| matches!(room.state(), RoomState::Left | RoomState::Banned)); |
| 108 | if left { |
| 109 | media.release(&session.scope).await; |
| 110 | } else { |
| 111 | kept.push_back(session); |
| 112 | } |
| 113 | } |
| 114 | *open = kept; |
| 115 | } |
| 116 | |
| 117 | /// the media service is cleared beside this, so nothing is released one by one |
| 118 | pub(crate) async fn clear(&self) { |
| 119 | self.open.lock().await.clear(); |
| 120 | } |
| 121 | } |
| 122 | |
| 123 | fn not_open() -> Error { |
| 124 | Error::new( |
| 125 | ErrorKind::EventContextNotOpen, |
| 126 | "event context is closed or was never opened", |
| 127 | ) |
| 128 | } |
| 129 | |
| 130 | /// leaving a room closes its event contexts |
| 131 | pub(super) async fn close_on_leave( |
| 132 | client: Client, |
| 133 | media: Arc<MediaService>, |
| 134 | sessions: Arc<Sessions>, |
| 135 | ) { |
| 136 | let mut info = client.room_info_notable_update_receiver(); |
| 137 | loop { |
| 138 | match info.recv().await { |
| 139 | Ok(_) | Err(RecvError::Lagged(_)) => sessions.close_left(&media, &client).await, |
| 140 | Err(RecvError::Closed) => return, |
| 141 | } |
| 142 | } |
| 143 | } |
| 144 | |
| 145 | pub(crate) async fn execute( |
| 146 | client: Client, |
| 147 | media: Arc<MediaService>, |
| 148 | sessions: Arc<Sessions>, |
| 149 | command: EventContextCommand, |
| 150 | ) -> Result<crate::CommandOutput, Error> { |
| 151 | let current_user_id = client.user_id().map(ToOwned::to_owned); |
| 152 | let room = |room_id: &OwnedRoomId| client.get_room(room_id).ok_or_else(Error::room_unknown); |
| 153 | let (context_id, token, limit, backwards) = match command { |
| 154 | EventContextCommand::Fetch { |
| 155 | room_id, |
| 156 | event_id, |
| 157 | before_limit, |
| 158 | after_limit, |
| 159 | } => { |
| 160 | let (context, registrations, _) = capture( |
| 161 | room(&room_id)?, |
| 162 | current_user_id, |
| 163 | event_id, |
| 164 | before_limit, |
| 165 | after_limit, |
| 166 | ) |
| 167 | .await?; |
| 168 | let context_id = sessions.open(&media, room_id, registrations).await; |
| 169 | return Ok(crate::CommandOutput::EventContext(EventContextSession { |
| 170 | context_id, |
| 171 | context, |
| 172 | })); |
| 173 | } |
| 174 | EventContextCommand::Close { context_id } => { |
| 175 | sessions.close(&media, &context_id).await?; |
| 176 | return Ok(crate::CommandOutput::Done); |
| 177 | } |
| 178 | EventContextCommand::Before { |
| 179 | context_id, |
| 180 | token, |
| 181 | limit, |
| 182 | } => (context_id, token, limit, true), |
| 183 | EventContextCommand::After { |
| 184 | context_id, |
| 185 | token, |
| 186 | limit, |
| 187 | } => (context_id, token, limit, false), |
| 188 | }; |
| 189 | if limit == 0 { |
| 190 | return Err(Error::new( |
| 191 | ErrorKind::Invalid, |
| 192 | "context page limit must be at least 1", |
| 193 | )); |
| 194 | } |
| 195 | validate_limit(limit)?; |
| 196 | if token.is_empty() { |
| 197 | return Err(Error::new( |
| 198 | ErrorKind::Invalid, |
| 199 | "context pagination token is empty", |
| 200 | )); |
| 201 | } |
| 202 | let room = room(&sessions.room(&context_id).await?)?; |
| 203 | let side = fetch_side(&room, Some(&token), limit, backwards).await?; |
| 204 | let mut page = Assembly::new( |
| 205 | &room, |
| 206 | current_user_id.as_deref(), |
| 207 | &side.values, |
| 208 | &side.highlights, |
| 209 | None, |
| 210 | ) |
| 211 | .await; |
| 212 | let events = page.take(0..side.values.len()); |
| 213 | sessions.register(&media, &context_id, page.media).await?; |
| 214 | Ok(crate::CommandOutput::EventContextPage(EventContextPage { |
| 215 | events, |
| 216 | next_token: side.token, |
| 217 | })) |
| 218 | } |
| 219 | |
| 220 | pub(crate) async fn capture( |
| 221 | room: Room, |
| 222 | current_user_id: Option<OwnedUserId>, |
| 223 | event_id: OwnedEventId, |
| 224 | before_limit: u16, |
| 225 | after_limit: u16, |
| 226 | ) -> Result<(EventContext, Vec<MediaRegistration>, Vec<Value>), Error> { |
| 227 | validate_limit(before_limit)?; |
| 228 | validate_limit(after_limit)?; |
| 229 | let raw = room |
| 230 | .event_with_context(&event_id, false, UInt::from(0_u16), None) |
| 231 | .await |
| 232 | .map_err(|source| { |
| 233 | let error = Error::from_sdk("fetching event context failed", source); |
| 234 | if error |
| 235 | .server |
| 236 | .as_ref() |
| 237 | .is_some_and(|server| server.status == 404) |
| 238 | { |
| 239 | Error::new(ErrorKind::EventNotFound, "event was not found") |
| 240 | } else { |
| 241 | error |
| 242 | } |
| 243 | })?; |
| 244 | let (target, target_highlight) = raw |
| 245 | .event |
| 246 | .as_ref() |
| 247 | .and_then(raw_event) |
| 248 | .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "event was not found"))?; |
| 249 | let (before, after) = try_join( |
| 250 | fetch_side(&room, raw.prev_batch_token.as_deref(), before_limit, true), |
| 251 | fetch_side(&room, raw.next_batch_token.as_deref(), after_limit, false), |
| 252 | ) |
| 253 | .await?; |
| 254 | let target_index = before.values.len(); |
| 255 | let mut values = before.values; |
| 256 | values.push(target); |
| 257 | values.extend(after.values); |
| 258 | let mut highlights = before.highlights; |
| 259 | highlights.push(target_highlight); |
| 260 | highlights.extend(after.highlights); |
| 261 | let mut page = Assembly::new( |
| 262 | &room, |
| 263 | current_user_id.as_deref(), |
| 264 | &values, |
| 265 | &highlights, |
| 266 | Some(target_index), |
| 267 | ) |
| 268 | .await; |
| 269 | let events_before = page.take(0..target_index); |
| 270 | let event = page |
| 271 | .take(target_index..target_index + 1) |
| 272 | .pop() |
| 273 | .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "event could not be read"))?; |
| 274 | let events_after = page.take(target_index + 1..values.len()); |
| 275 | Ok(( |
| 276 | EventContext { |
| 277 | events_before, |
| 278 | event, |
| 279 | events_after, |
| 280 | start_token: before.token, |
| 281 | end_token: after.token, |
| 282 | }, |
| 283 | page.media, |
| 284 | values, |
| 285 | )) |
| 286 | } |
| 287 | |
| 288 | struct Side { |
| 289 | /// in timeline order |
| 290 | values: Vec<Value>, |
| 291 | /// beside `values`, from the push actions the sdk worked out as it fetched |
| 292 | highlights: Vec<bool>, |
| 293 | token: Option<String>, |
| 294 | } |
| 295 | |
| 296 | fn raw_event(event: &SdkTimelineEvent) -> Option<(Value, bool)> { |
| 297 | let value = serde_json::from_str(event.raw().json().get()).ok()?; |
| 298 | let highlight = event |
| 299 | .push_actions() |
| 300 | .is_some_and(|actions| actions.iter().any(Action::is_highlight)); |
| 301 | Some((value, highlight)) |
| 302 | } |
| 303 | |
| 304 | async fn fetch_side( |
| 305 | room: &Room, |
| 306 | token: Option<&str>, |
| 307 | limit: u16, |
| 308 | backwards: bool, |
| 309 | ) -> Result<Side, Error> { |
| 310 | let Some(token) = token else { |
| 311 | return Ok(Side { |
| 312 | values: Vec::new(), |
| 313 | highlights: Vec::new(), |
| 314 | token: None, |
| 315 | }); |
| 316 | }; |
| 317 | if limit == 0 { |
| 318 | return Ok(Side { |
| 319 | values: Vec::new(), |
| 320 | highlights: Vec::new(), |
| 321 | token: Some(token.to_owned()), |
| 322 | }); |
| 323 | } |
| 324 | let mut options = if backwards { |
| 325 | MessagesOptions::backward() |
| 326 | } else { |
| 327 | MessagesOptions::forward() |
| 328 | }; |
| 329 | options = options.from(token); |
| 330 | options.limit = UInt::from(limit); |
| 331 | let messages = room |
| 332 | .messages(options) |
| 333 | .await |
| 334 | .map_err(|source| Error::from_sdk("fetching event context events failed", source))?; |
| 335 | let mut events: Vec<_> = messages.chunk.iter().filter_map(raw_event).collect(); |
| 336 | if backwards { |
| 337 | events.reverse(); |
| 338 | } |
| 339 | let (values, highlights) = events.into_iter().unzip(); |
| 340 | Ok(Side { |
| 341 | values, |
| 342 | highlights, |
| 343 | token: messages.end, |
| 344 | }) |
| 345 | } |
| 346 | |
| 347 | /// a window of raw events normalised together, so relations inside it land on their targets |
| 348 | struct Assembly { |
| 349 | items: Vec<(usize, TimelineItem)>, |
| 350 | media: Vec<MediaRegistration>, |
| 351 | } |
| 352 | |
| 353 | impl Assembly { |
| 354 | async fn new( |
| 355 | room: &Room, |
| 356 | current_user_id: Option<&ruma::UserId>, |
| 357 | values: &[Value], |
| 358 | highlights: &[bool], |
| 359 | keep: Option<usize>, |
| 360 | ) -> Self { |
| 361 | let page = normalise_context(values, keep, current_user_id, room.room_id()); |
| 362 | let mut items = page.items; |
| 363 | for (index, item) in &mut items { |
| 364 | if let TimelineItem::Event(event) = item { |
| 365 | event.highlight = highlights[*index]; |
| 366 | } |
| 367 | } |
| 368 | let mut media = page.media; |
| 369 | let mut targets: HashMap<OwnedEventId, TimelineEvent> = HashMap::new(); |
| 370 | let mut missing = Vec::new(); |
| 371 | for (_, item) in &items { |
| 372 | if let TimelineItem::Event(event) = item |
| 373 | && let Some(event_id) = &event.event_id |
| 374 | { |
| 375 | targets.insert(event_id.clone(), (**event).clone()); |
| 376 | } |
| 377 | } |
| 378 | for event in events(&mut items) { |
| 379 | if let Some(reply) = &event.reply_to |
| 380 | && !targets.contains_key(&reply.event_id) |
| 381 | && !missing.contains(&reply.event_id) |
| 382 | { |
| 383 | missing.push(reply.event_id.clone()); |
| 384 | } |
| 385 | } |
| 386 | let fetched = join_all( |
| 387 | missing |
| 388 | .into_iter() |
| 389 | .map(|event_id| fetch_reply(room, current_user_id, event_id)), |
| 390 | ) |
| 391 | .await; |
| 392 | for (event, registrations) in fetched.into_iter().flatten() { |
| 393 | media.extend(registrations); |
| 394 | if let Some(event_id) = event.event_id.clone() { |
| 395 | targets.insert(event_id, event); |
| 396 | } |
| 397 | } |
| 398 | let profiles = profiles( |
| 399 | room, |
| 400 | items |
| 401 | .iter() |
| 402 | .filter_map(|(_, item)| match item { |
| 403 | TimelineItem::Event(event) => Some(&event.sender.user_id), |
| 404 | _ => None, |
| 405 | }) |
| 406 | .chain(targets.values().map(|event| &event.sender.user_id)), |
| 407 | ) |
| 408 | .await; |
| 409 | apply_profiles(targets.values_mut(), &profiles); |
| 410 | apply_profiles(events(&mut items), &profiles); |
| 411 | resolve_replies(events(&mut items), &targets); |
| 412 | Self { items, media } |
| 413 | } |
| 414 | |
| 415 | /// the items that came from raw events in `range` |
| 416 | fn take(&mut self, range: std::ops::Range<usize>) -> Vec<TimelineItem> { |
| 417 | let (taken, kept) = std::mem::take(&mut self.items) |
| 418 | .into_iter() |
| 419 | .partition(|(index, _)| range.contains(index)); |
| 420 | self.items = kept; |
| 421 | taken.into_iter().map(|(_, item)| item).collect() |
| 422 | } |
| 423 | } |
| 424 | |
| 425 | fn events(items: &mut [(usize, TimelineItem)]) -> impl Iterator<Item = &mut TimelineEvent> { |
| 426 | items.iter_mut().filter_map(|(_, item)| match item { |
| 427 | TimelineItem::Event(event) => Some(&mut **event), |
| 428 | _ => None, |
| 429 | }) |
| 430 | } |
| 431 | |
| 432 | async fn fetch_reply( |
| 433 | room: &Room, |
| 434 | current_user_id: Option<&ruma::UserId>, |
| 435 | event_id: OwnedEventId, |
| 436 | ) -> Option<(TimelineEvent, Vec<MediaRegistration>)> { |
| 437 | let event = room |
| 438 | .load_or_fetch_event(&event_id, None) |
| 439 | .await |
| 440 | .inspect_err(|error| tracing::debug!(%event_id, %error, "fetching a reply target failed")) |
| 441 | .ok()?; |
| 442 | let value: Value = serde_json::from_str(event.raw().json().get()).ok()?; |
| 443 | let identity = format!("{}\n{event_id}", room.room_id()); |
| 444 | match normalise_observed(&value, current_user_id, &identity)? { |
| 445 | (TimelineItem::Event(event), media) => Some((*event, media)), |
| 446 | _ => None, |
| 447 | } |
| 448 | } |
| 449 | |
| 450 | /// under lazy-loaded members a sender may be unknown locally, and `get_member` fetches the |
| 451 | /// member list once for the room when it is |
| 452 | async fn profiles<'a>( |
| 453 | room: &Room, |
| 454 | user_ids: impl Iterator<Item = &'a OwnedUserId>, |
| 455 | ) -> HashMap<OwnedUserId, Sender> { |
| 456 | let mut unique: Vec<&OwnedUserId> = user_ids.collect(); |
| 457 | unique.sort(); |
| 458 | unique.dedup(); |
| 459 | join_all(unique.into_iter().map(|user_id| async move { |
| 460 | let member = room.get_member(user_id).await.ok().flatten()?; |
| 461 | Some(( |
| 462 | user_id.clone(), |
| 463 | Sender { |
| 464 | user_id: user_id.clone(), |
| 465 | display_name: member.display_name().map(str::to_owned), |
| 466 | avatar: member.avatar_url().map(MediaId::plain), |
| 467 | name_ambiguous: member.name_ambiguous(), |
| 468 | }, |
| 469 | )) |
| 470 | })) |
| 471 | .await |
| 472 | .into_iter() |
| 473 | .flatten() |
| 474 | .collect() |
| 475 | } |
| 476 | |
| 477 | fn validate_limit(limit: u16) -> Result<(), Error> { |
| 478 | if limit > MAX_CONTEXT_LIMIT { |
| 479 | Err(Error::new(ErrorKind::Invalid, "context limit exceeds 50")) |
| 480 | } else { |
| 481 | Ok(()) |
| 482 | } |
| 483 | } |