atagen

macltok

public
Star0
Watch0
913832a327a82f6b019568117ba636d069143124
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git
src/mentions/service.rsRust1582 lines61.3 KB
1
use std::{
2
    collections::{BTreeMap, BTreeSet},
3
    sync::{
4
        Arc,
5
        atomic::{AtomicBool, Ordering},
6
    },
7
    time::Duration,
8
};
9
10
use futures_util::{StreamExt, stream::FuturesUnordered};
11
use matrix_sdk::{
12
    Client, Room, RoomState,
13
    event_cache::RedecryptorReport,
14
    sliding_sync::{ResponseCommitBatch, ResponseCommitHook},
15
};
16
use ruma::{
17
    OwnedRoomId, OwnedUserId, RoomId,
18
    events::{
19
        AnySyncTimelineEvent, ignored_user_list::IgnoredUserListEventContent,
20
        room::encrypted::OriginalSyncRoomEncryptedEvent,
21
    },
22
    serde::Raw,
23
};
24
use serde_json::Value;
25
use tokio::sync::{mpsc, oneshot, watch};
26
27
use super::state::{Candidate, Checkpoint, Data, SavedEvent, SyncState, Window, id, invalid, now};
28
use crate::{
29
    CommandOutput, EngineConfig, Error, ErrorKind, Event, MentionCaptureState, MentionEntry,
30
    MentionGapReason, MentionKey, MentionReason, MentionsCommand, MentionsPage, MentionsUpdate,
31
    TimelineContent, TimelineItem,
32
    engine::{MediaService, MentionsStore, context, executor::EngineBoxFuture},
33
    normalise::{MediaRegistration, normalise_observed},
34
    transport::EventSender,
35
};
36
37
#[cfg(all(test, feature = "sqlite"))]
38
#[path = "service_tests.rs"]
39
mod tests;
40
41
const HOOK_TIMEOUT: Duration = Duration::from_secs(10);
42
const FLUSH_RETRY_MS: u64 = 60_000;
43
44
#[derive(Clone)]
45
pub(crate) struct Handle {
46
    requests: mpsc::Sender<Request>,
47
    timed_out: Arc<AtomicBool>,
48
}
49
50
enum Request {
51
    Batch(ResponseCommitBatch, oneshot::Sender<()>),
52
    Command(
53
        MentionsCommand,
54
        oneshot::Sender<Result<CommandOutput, Error>>,
55
    ),
56
    Stop(oneshot::Sender<()>),
57
}
58
59
pub(crate) struct Service {
60
    client: Client,
61
    account: OwnedUserId,
62
    media: Arc<MediaService>,
63
    store: Option<MentionsStore>,
64
    data: Data,
65
    sync: SyncState,
66
    dirty: bool,
67
    flush_after: u64,
68
    ignored: BTreeSet<OwnedUserId>,
69
    leaving: BTreeSet<OwnedRoomId>,
70
    timed_out: Arc<AtomicBool>,
71
    config: EngineConfig,
72
    requests: mpsc::Receiver<Request>,
73
    updates: watch::Sender<MentionsPage>,
74
    recovery_budget: usize,
75
    latest_scan: bool,
76
    recovery_retry: u64,
77
}
78
79
enum Work {
80
    Capture(
81
        String,
82
        u64,
83
        Result<(crate::EventContext, Vec<MediaRegistration>, Vec<Value>), Error>,
84
    ),
85
    Candidate(String, Result<Value, Error>),
86
    History(
87
        bool,
88
        Option<String>,
89
        Result<ruma::api::client::push::get_notifications::v3::Response, Error>,
90
    ),
91
}
92
93
impl Handle {
94
    /// never fails, a mentions failure becomes a coverage gap rather than a stalled sync
95
    pub(crate) fn hook(&self) -> ResponseCommitHook {
96
        let handle = self.clone();
97
        ResponseCommitHook::new(move |batch| {
98
            let handle = handle.clone();
99
            async move {
100
                let committed = async {
101
                    let (reply, response) = oneshot::channel();
102
                    handle
103
                        .requests
104
                        .send(Request::Batch(batch, reply))
105
                        .await
106
                        .ok()?;
107
                    response.await.ok()
108
                };
109
                match matrix_sdk::timeout::timeout(committed, HOOK_TIMEOUT).await {
110
                    Ok(Some(())) => {}
111
                    Ok(None) => tracing::warn!("mentions service stopped, sync carries on"),
112
                    Err(_) => {
113
                        handle.timed_out.store(true, Ordering::Release);
114
                        tracing::warn!("mentions commit timed out, sync carries on");
115
                    }
116
                }
117
                Ok(())
118
            }
119
        })
120
    }
121
122
    pub(crate) async fn execute(&self, command: MentionsCommand) -> Result<CommandOutput, Error> {
123
        let (reply, response) = oneshot::channel();
124
        self.requests
125
            .send(Request::Command(command, reply))
126
            .await
127
            .map_err(|_| Error::interrupted())?;
128
        response.await.map_err(|_| Error::interrupted())?
129
    }
130
131
    pub(crate) async fn stop(&self) {
132
        let (reply, response) = oneshot::channel();
133
        if self.requests.send(Request::Stop(reply)).await.is_ok() {
134
            let _ = response.await;
135
        }
136
    }
137
}
138
139
#[cfg(any(feature = "sqlite", feature = "indexeddb"))]
140
async fn open_store(
141
    store: Option<Result<MentionsStore, Error>>,
142
) -> (
143
    Option<MentionsStore>,
144
    Data,
145
    SyncState,
146
    Vec<MentionGapReason>,
147
) {
148
    let mut store = match store {
149
        None => return (None, Data::default(), SyncState::default(), Vec::new()),
150
        Some(Err(error)) => {
151
            tracing::error!(%error, "mentions store unavailable, the inbox runs in memory");
152
            return (
153
                None,
154
                Data::default(),
155
                SyncState::default(),
156
                vec![MentionGapReason::StoreFailed],
157
            );
158
        }
159
        Some(Ok(store)) => store,
160
    };
161
    let loaded = (|| {
162
        let data: Data = store.read("inbox")?.unwrap_or_default();
163
        let sync: SyncState = store.read("sync")?.unwrap_or_default();
164
        if sync.inbox_revision != data.revision {
165
            return Err(Error::new(
166
                ErrorKind::Store,
167
                "mentions checkpoint does not match its inbox",
168
            ));
169
        }
170
        Ok((data, sync))
171
    })();
172
    let (data, sync) = match loaded {
173
        Ok(loaded) => loaded,
174
        Err(error) => {
175
            tracing::warn!(%error, "resetting an unreadable mentions inbox");
176
            if let Err(error) = store.reset().await {
177
                tracing::error!(%error, "mentions store unavailable, the inbox runs in memory");
178
                return (
179
                    None,
180
                    Data::default(),
181
                    SyncState::default(),
182
                    vec![MentionGapReason::StoreFailed],
183
                );
184
            }
185
            (Data::default(), SyncState::default())
186
        }
187
    };
188
    let gaps = if store.was_reset() {
189
        vec![MentionGapReason::StoreReset]
190
    } else {
191
        Vec::new()
192
    };
193
    (Some(store), data, sync, gaps)
194
}
195
196
#[cfg(not(any(feature = "sqlite", feature = "indexeddb")))]
197
async fn open_store(
198
    store: Option<Result<MentionsStore, Error>>,
199
) -> (
200
    Option<MentionsStore>,
201
    Data,
202
    SyncState,
203
    Vec<MentionGapReason>,
204
) {
205
    let gaps = match store {
206
        None => Vec::new(),
207
        Some(Err(_)) => vec![MentionGapReason::StoreFailed],
208
        Some(Ok(store)) => match store {},
209
    };
210
    (None, Data::default(), SyncState::default(), gaps)
211
}
212
213
impl Service {
214
    pub(crate) async fn new(
215
        client: Client,
216
        media: Arc<MediaService>,
217
        store: Option<Result<MentionsStore, Error>>,
218
        config: EngineConfig,
219
    ) -> Result<(Handle, Self, watch::Receiver<MentionsPage>), Error> {
220
        let account = client
221
            .user_id()
222
            .ok_or_else(|| {
223
                Error::new(
224
                    ErrorKind::Unauthenticated,
225
                    "mentions require an authenticated account",
226
                )
227
            })?
228
            .to_owned();
229
        let (store, mut data, sync, gaps) = open_store(store).await;
230
        for reason in &gaps {
231
            data.gap(None, *reason, None);
232
        }
233
        let page = data.page(&account, None, 50)?;
234
        media
235
            .register(
236
                page.entries
237
                    .iter()
238
                    .filter_map(|entry| data.events.get(&id(&entry.key)))
239
                    .flat_map(|event| event.media.clone())
240
                    .collect(),
241
            )
242
            .await;
243
        let (updates, receiver) = watch::channel(page);
244
        let (sender, requests) = mpsc::channel(config.command_capacity.max(1));
245
        let timed_out = Arc::new(AtomicBool::new(false));
246
        let mut service = Self {
247
            client,
248
            account,
249
            media,
250
            store,
251
            data,
252
            sync,
253
            dirty: !gaps.is_empty(),
254
            flush_after: 0,
255
            ignored: BTreeSet::new(),
256
            leaving: BTreeSet::new(),
257
            timed_out: timed_out.clone(),
258
            config,
259
            requests,
260
            updates,
261
            recovery_budget: 20,
262
            latest_scan: true,
263
            recovery_retry: 0,
264
        };
265
        service.ignored = service.ignored_users().await;
266
        Ok((
267
            Handle {
268
                requests: sender,
269
                timed_out,
270
            },
271
            service,
272
            receiver,
273
        ))
274
    }
275
276
    async fn ignored_users(&self) -> BTreeSet<OwnedUserId> {
277
        self.client
278
            .account()
279
            .account_data::<IgnoredUserListEventContent>()
280
            .await
281
            .ok()
282
            .flatten()
283
            .and_then(|raw| raw.deserialize().ok())
284
            .map(|content| content.ignored_users.into_keys().collect())
285
            .unwrap_or_default()
286
    }
287
288
    async fn publish(&mut self) {
289
        let Ok(page) = self.data.page(&self.account, None, 50) else {
290
            return;
291
        };
292
        self.media
293
            .register(
294
                page.entries
295
                    .iter()
296
                    .filter_map(|entry| self.data.events.get(&id(&entry.key)))
297
                    .flat_map(|event| event.media.clone())
298
                    .collect(),
299
            )
300
            .await;
301
        self.updates.send_replace(page);
302
    }
303
304
    /// writes the sync record, and the inbox with it when `dirty`, in one commit so a
305
    /// checkpoint never lands without the inbox changes it covers
306
    async fn flush(&mut self) -> Result<(), Error> {
307
        if self.dirty {
308
            self.advance()?;
309
        }
310
        #[cfg(any(feature = "sqlite", feature = "indexeddb"))]
311
        if let Some(store) = &mut self.store {
312
            use super::state::json_len;
313
            use super::store::Write;
314
315
            let encode = |error| {
316
                Error::with_source(
317
                    ErrorKind::Store,
318
                    "encoding the mentions inbox failed",
319
                    error,
320
                )
321
            };
322
            let sync = serde_json::to_value(&self.sync).map_err(encode)?;
323
            let mut writes = Vec::new();
324
            if self.dirty {
325
                let available = store
326
                    .available(&["inbox", "sync"])
327
                    .saturating_sub(store.sealed_bound("sync", json_len(&sync)));
328
                let mut size = store.sealed_bound("inbox", json_len(&self.data));
329
                while size > available {
330
                    let Some(freed) = self.data.evict_one() else {
331
                        break;
332
                    };
333
                    size = size.saturating_sub(freed * 3);
334
                }
335
                writes.push(Write::Set(
336
                    "inbox".to_owned(),
337
                    serde_json::to_value(&self.data).map_err(encode)?,
338
                ));
339
            }
340
            writes.push(Write::Set("sync".to_owned(), sync));
341
            store.commit(store.revision(), writes).await?;
342
        }
343
        if self.dirty {
344
            self.dirty = false;
345
            self.publish().await;
346
        }
347
        Ok(())
348
    }
349
350
    fn advance(&mut self) -> Result<(), Error> {
351
        self.data.revision = self
352
            .data
353
            .revision
354
            .checked_add(1)
355
            .ok_or_else(|| invalid("mentions revision exhausted"))?;
356
        self.data.bound(
357
            self.config.max_mentions_entries,
358
            self.config.max_mention_context_events,
359
            self.config.max_mention_candidates,
360
        );
361
        self.sync.inbox_revision = self.data.revision;
362
        Ok(())
363
    }
364
365
    async fn apply(&mut self, data: Data) -> Result<(), Error> {
366
        let dirty = std::mem::replace(&mut self.dirty, true);
367
        let previous = std::mem::replace(&mut self.data, data);
368
        let result = self.flush().await;
369
        if result.is_err() {
370
            self.data = previous;
371
            self.dirty = dirty;
372
        }
373
        result
374
    }
375
376
    async fn store_failed(&mut self, error: Error) {
377
        tracing::error!(%error, "saving the mentions inbox failed");
378
        self.data.gap(None, MentionGapReason::StoreFailed, None);
379
        self.dirty = true;
380
        self.flush_after = now().saturating_add(FLUSH_RETRY_MS);
381
        self.publish().await;
382
    }
383
384
    async fn command(&mut self, command: MentionsCommand) -> Result<CommandOutput, Error> {
385
        match command {
386
            MentionsCommand::Page { cursor, limit } => {
387
                let page = self.data.page(&self.account, cursor.as_deref(), limit)?;
388
                self.media
389
                    .register(
390
                        page.entries
391
                            .iter()
392
                            .filter_map(|entry| self.data.events.get(&id(&entry.key)))
393
                            .flat_map(|event| event.media.clone())
394
                            .collect(),
395
                    )
396
                    .await;
397
                Ok(CommandOutput::MentionsPage(page))
398
            }
399
            MentionsCommand::Context(key) => {
400
                let (context, media) = self.data.context(&key)?;
401
                self.media.register(media).await;
402
                Ok(CommandOutput::MentionContext(context))
403
            }
404
            MentionsCommand::RetryContext(key) => {
405
                let room = self
406
                    .client
407
                    .get_room(&key.room_id)
408
                    .filter(|room| room.state() == RoomState::Joined)
409
                    .ok_or_else(|| {
410
                        Error::new(
411
                            ErrorKind::Forbidden,
412
                            "saved context cannot be refetched after leaving the room",
413
                        )
414
                    })?;
415
                drop(room);
416
                let mut next = self.data.clone();
417
                let entry = next
418
                    .entries
419
                    .get_mut(&id(&key))
420
                    .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "mention was not found"))?;
421
                entry.capture = MentionCaptureState::Pending;
422
                entry.revision = self.data.revision + 1;
423
                next.windows.entry(id(&key)).or_insert(Window {
424
                    before: Vec::new(),
425
                    target: Some(id(&key)),
426
                    after: Vec::new(),
427
                    start: None,
428
                    end: None,
429
                });
430
                self.apply(next).await?;
431
                let (context, media) = self.data.context(&key)?;
432
                self.media.register(media).await;
433
                Ok(CommandOutput::MentionContext(context))
434
            }
435
            MentionsCommand::SetRead { keys, read } => {
436
                let mut next = self.data.clone();
437
                for key in keys {
438
                    if let Some(entry) = next.entries.get_mut(&id(&key)) {
439
                        entry.read = read;
440
                        entry.revision = self.data.revision + 1;
441
                    }
442
                }
443
                self.apply(next).await?;
444
                Ok(CommandOutput::Done)
445
            }
446
            MentionsCommand::ReadThrough(boundary) => {
447
                if boundary.account != self.account || boundary.revision > self.data.revision {
448
                    return Err(invalid("invalid mentions read boundary"));
449
                }
450
                let mut next = self.data.clone();
451
                for entry in next
452
                    .entries
453
                    .values_mut()
454
                    .filter(|entry| entry.discovery_revision <= boundary.revision)
455
                {
456
                    entry.read = true;
457
                    entry.revision = self.data.revision + 1;
458
                }
459
                self.apply(next).await?;
460
                Ok(CommandOutput::Done)
461
            }
462
            MentionsCommand::Recover => {
463
                self.recovery_budget = 20;
464
                self.latest_scan = true;
465
                self.recovery_retry = 0;
466
                let mut next = self.data.clone();
467
                next.coverage.recovering = true;
468
                next.coverage.historical_supported = None;
469
                self.apply(next).await?;
470
                Ok(CommandOutput::Done)
471
            }
472
        }
473
    }
474
475
    fn observable(&self, room_id: &RoomId) -> Option<Room> {
476
        self.client
477
            .get_room(room_id)
478
            .filter(|room| room.state() == RoomState::Joined || self.leaving.contains(room_id))
479
    }
480
481
    async fn batch(&mut self, batch: ResponseCommitBatch) {
482
        let checkpoint = self.sync.checkpoints.get(&batch.connection_id).cloned();
483
        if checkpoint
484
            .as_ref()
485
            .is_some_and(|checkpoint| checkpoint.pos == batch.next_pos)
486
        {
487
            return;
488
        }
489
        self.ignored = self.ignored_users().await;
490
        self.leaving = batch.response.rooms.left.keys().cloned().collect();
491
        let mut next = self.data.clone();
492
        let mut sync = self.sync.clone();
493
        if self.timed_out.swap(false, Ordering::AcqRel) {
494
            next.gap(None, MentionGapReason::HookTimedOut, None);
495
        }
496
        // a crash after the commit but before the sdk saved `pos` replays from the
497
        // previous position, and that response is a superset of the one committed
498
        let continuous = match &checkpoint {
499
            None => batch.previous_pos.is_none(),
500
            Some(checkpoint) => {
501
                batch.previous_pos.as_ref() == Some(&checkpoint.pos)
502
                    || batch.previous_pos == checkpoint.previous
503
            }
504
        };
505
        if !continuous {
506
            next.gap(None, MentionGapReason::SyncDiscontinuity, None);
507
        }
508
        for (room_id, raw) in &batch.raw_rooms {
509
            if raw.limited && sync.observed_rooms.contains(room_id) {
510
                next.gap(
511
                    Some(room_id.clone()),
512
                    MentionGapReason::LimitedTimeline,
513
                    raw.prev_batch.clone(),
514
                );
515
            }
516
        }
517
        let room_ids: BTreeSet<_> = batch
518
            .raw_rooms
519
            .keys()
520
            .chain(batch.response.rooms.joined.keys())
521
            .chain(batch.response.rooms.left.keys())
522
            .cloned()
523
            .collect();
524
        for room_id in &room_ids {
525
            if self.observable(room_id).is_none() {
526
                continue;
527
            }
528
            sync.observed_rooms.insert(room_id.clone());
529
            let mut observed = BTreeMap::new();
530
            if let Some(raw) = batch.raw_rooms.get(room_id) {
531
                for event in &raw.timeline {
532
                    let Ok(value) = serde_json::from_str::<Value>(event.json().get()) else {
533
                        continue;
534
                    };
535
                    if let Some(event_id) = value["event_id"].as_str() {
536
                        observed.insert(event_id.to_owned(), value);
537
                    }
538
                }
539
            }
540
            let processed = batch
541
                .response
542
                .rooms
543
                .joined
544
                .get(room_id)
545
                .map(|update| &update.timeline.events)
546
                .into_iter()
547
                .chain(
548
                    batch
549
                        .response
550
                        .rooms
551
                        .left
552
                        .get(room_id)
553
                        .map(|update| &update.timeline.events),
554
                )
555
                .flatten();
556
            for event in processed {
557
                let Ok(value) = serde_json::from_str::<Value>(event.raw().json().get()) else {
558
                    continue;
559
                };
560
                if let Some(event_id) = value["event_id"].as_str() {
561
                    observed.insert(event_id.to_owned(), value);
562
                }
563
            }
564
            let mut observed: Vec<_> = observed.into_values().collect();
565
            observed.sort_by_key(|value| value["origin_server_ts"].as_u64().unwrap_or(0));
566
            for value in observed {
567
                self.observe(&mut next, room_id, value).await;
568
            }
569
        }
570
        for room_id in batch.response.rooms.left.keys() {
571
            next.gap(Some(room_id.clone()), MentionGapReason::RoomLeft, None);
572
            next.candidates
573
                .retain(|_, candidate| &candidate.room_id != room_id);
574
            let revision = next.revision + 1;
575
            for entry in next.entries.values_mut().filter(|entry| {
576
                &entry.key.room_id == room_id && entry.capture != MentionCaptureState::Unavailable
577
            }) {
578
                entry.capture = MentionCaptureState::Unavailable;
579
                entry.revision = revision;
580
            }
581
        }
582
        let reevaluate: Vec<_> = next
583
            .entries
584
            .values()
585
            .filter(|entry| room_ids.contains(&entry.key.room_id))
586
            .filter_map(|entry| {
587
                next.events
588
                    .get(&id(&entry.key))
589
                    .and_then(|event| event.raw.clone())
590
                    .map(|raw| (entry.key.room_id.clone(), raw))
591
            })
592
            .collect();
593
        for (room_id, raw) in reevaluate {
594
            self.observe(&mut next, &room_id, raw).await;
595
        }
596
        self.filter_ignored(&mut next);
597
        self.leaving.clear();
598
        sync.checkpoints.insert(
599
            batch.connection_id,
600
            Checkpoint {
601
                pos: batch.next_pos,
602
                previous: batch.previous_pos,
603
            },
604
        );
605
        let changed = next != self.data;
606
        if changed {
607
            self.data = next;
608
            self.dirty = true;
609
        }
610
        self.sync = sync;
611
        if now() < self.flush_after {
612
            if changed && self.advance().is_ok() {
613
                self.publish().await;
614
            }
615
        } else if let Err(error) = self.flush().await {
616
            self.store_failed(error).await;
617
        }
618
    }
619
620
    fn filter_ignored(&self, next: &mut Data) -> bool {
621
        let removed: Vec<_> = next
622
            .entries
623
            .iter()
624
            .filter(|(_, entry)| {
625
                matches!(&entry.event, TimelineItem::Event(event) if self.ignored.contains(&event.sender.user_id))
626
            })
627
            .map(|(key, _)| key.clone())
628
            .collect();
629
        for key in &removed {
630
            next.entries.remove(key);
631
            next.windows.remove(key);
632
        }
633
        next.collect_events();
634
        !removed.is_empty()
635
    }
636
637
    async fn observe(&self, data: &mut Data, room_id: &OwnedRoomId, mut value: Value) {
638
        let Some(event_id) = value["event_id"]
639
            .as_str()
640
            .and_then(|value| value.parse().ok())
641
        else {
642
            return;
643
        };
644
        let key = MentionKey {
645
            room_id: room_id.clone(),
646
            event_id,
647
        };
648
        let identity = id(&key);
649
        let Some(timestamp) = value["origin_server_ts"].as_u64() else {
650
            return;
651
        };
652
        data.coverage.earliest_timestamp_ms = Some(
653
            data.coverage
654
                .earliest_timestamp_ms
655
                .map_or(timestamp, |old| old.min(timestamp)),
656
        );
657
        if value["type"] == "m.room.redaction" {
658
            if let Some(target) = value["redacts"]
659
                .as_str()
660
                .or_else(|| value["content"]["redacts"].as_str())
661
                .and_then(|value| value.parse().ok())
662
            {
663
                let target = id(&MentionKey {
664
                    room_id: room_id.clone(),
665
                    event_id: target,
666
                });
667
                data.redacted.insert(target.clone(), timestamp);
668
                data.entries.remove(&target);
669
                data.windows.remove(&target);
670
                data.candidates.remove(&target);
671
                data.edits.remove(&target);
672
                data.candidates.retain(|_, candidate| {
673
                    candidate.room_id != *room_id
674
                        || candidate.raw["content"]["m.relates_to"]["event_id"].as_str()
675
                            != target.split_once('\n').map(|(_, event)| event)
676
                });
677
                if let Some(saved) = data.events.get_mut(&target) {
678
                    redact(saved);
679
                }
680
                data.touch_contexts(&target);
681
            }
682
            return;
683
        }
684
        if data.redacted.contains_key(&identity)
685
            || value["unsigned"].get("redacted_because").is_some()
686
        {
687
            data.redacted.insert(identity.clone(), timestamp);
688
            data.entries.remove(&identity);
689
            data.windows.remove(&identity);
690
            data.candidates.remove(&identity);
691
            if let Some(saved) = data.events.get_mut(&identity) {
692
                redact(saved);
693
            }
694
            data.touch_contexts(&identity);
695
            return;
696
        }
697
        if value["content"]["m.relates_to"]["rel_type"] == "m.replace" {
698
            if let Some(target) = value["content"]["m.relates_to"]["event_id"]
699
                .as_str()
700
                .and_then(|value| value.parse().ok())
701
            {
702
                let target = id(&MentionKey {
703
                    room_id: room_id.clone(),
704
                    event_id: target,
705
                });
706
                if data.redacted.contains_key(&target) {
707
                    return;
708
                }
709
                let newer = data.edits.get(&target).is_none_or(|old| {
710
                    (old["origin_server_ts"].as_u64(), old["event_id"].as_str())
711
                        < (Some(timestamp), value["event_id"].as_str())
712
                });
713
                if newer {
714
                    if let Some(original) =
715
                        data.events.get(&target).and_then(|saved| saved.raw.clone())
716
                    {
717
                        if original["sender"] != value["sender"] {
718
                            return;
719
                        }
720
                        data.edits.insert(target.clone(), value.clone());
721
                        value = edited(original, &value);
722
                        return Box::pin(self.observe(data, room_id, value)).await;
723
                    }
724
                    data.candidates
725
                        .entry(identity)
726
                        .and_modify(|candidate| candidate.raw = value.clone())
727
                        .or_insert_with(|| Candidate::new(room_id.clone(), value, 0));
728
                }
729
            }
730
            return;
731
        }
732
        if let Some(edit) = value["unsigned"]["m.relations"].get("m.replace").cloned() {
733
            if edit["sender"] == value["sender"]
734
                && edit["content"]["m.new_content"].is_object()
735
                && edit["content"]["m.relates_to"]["event_id"]
736
                    .as_str()
737
                    .is_none_or(|target| target == key.event_id.as_str())
738
                && edit["origin_server_ts"].as_u64().is_some()
739
                && data.edits.get(&identity).is_none_or(|old| {
740
                    (old["origin_server_ts"].as_u64(), old["event_id"].as_str())
741
                        < (edit["origin_server_ts"].as_u64(), edit["event_id"].as_str())
742
                })
743
            {
744
                data.edits.insert(identity.clone(), edit);
745
            }
746
        }
747
        if let Some(edit) = data.edits.get(&identity) {
748
            if value["sender"] == edit["sender"] {
749
                value = edited(value, edit);
750
            }
751
        }
752
        if let Some(edit) = value["unsigned"].get("macltok_edit").cloned() {
753
            data.edits.insert(identity.clone(), edit);
754
        }
755
        if value["type"] == "m.room.encrypted" {
756
            data.candidates
757
                .entry(identity)
758
                .or_insert_with(|| Candidate::new(room_id.clone(), value, 0));
759
            return;
760
        }
761
        let Some((item, media)) = normalise_observed(&value, Some(&self.account), &identity) else {
762
            return;
763
        };
764
        self.supplement(data, &key, &item, &media, &value);
765
        let (reasons, suppressed, unresolved) = self.classify(data, room_id, &value).await;
766
        if unresolved {
767
            data.candidates
768
                .entry(identity.clone())
769
                .or_insert_with(|| Candidate::new(room_id.clone(), value.clone(), now() + 30_000));
770
        } else {
771
            data.candidates.remove(&identity);
772
        }
773
        if reasons.is_empty() {
774
            data.entries.remove(&identity);
775
            data.windows.remove(&identity);
776
            if data.events.contains_key(&identity) || data.edits.contains_key(&identity) {
777
                data.events.insert(
778
                    identity,
779
                    SavedEvent {
780
                        item,
781
                        media,
782
                        raw: Some(value),
783
                    },
784
                );
785
            }
786
            return;
787
        }
788
        let previous = data.entries.get(&identity);
789
        let changed = previous.is_none_or(|entry| {
790
            entry.event != item
791
                || entry.reasons != reasons
792
                || entry.notification_suppressed != suppressed
793
        });
794
        let entry = MentionEntry {
795
            key,
796
            timestamp_ms: timestamp,
797
            discovered_at_ms: previous.map_or_else(now, |entry| entry.discovered_at_ms),
798
            discovery_revision: previous
799
                .map_or(data.revision + 1, |entry| entry.discovery_revision),
800
            revision: if changed {
801
                data.revision + 1
802
            } else {
803
                previous.unwrap().revision
804
            },
805
            event: item.clone(),
806
            reasons,
807
            read: previous.is_some_and(|entry| entry.read),
808
            capture: previous.map_or(MentionCaptureState::Pending, |entry| {
809
                if changed {
810
                    MentionCaptureState::Pending
811
                } else {
812
                    entry.capture.clone()
813
                }
814
            }),
815
            room_name: self.client.get_room(room_id).and_then(|room| room.name()),
816
            notification_suppressed: suppressed,
817
        };
818
        data.entries.insert(identity.clone(), entry);
819
        data.windows.entry(identity.clone()).or_insert(Window {
820
            before: Vec::new(),
821
            target: Some(identity.clone()),
822
            after: Vec::new(),
823
            start: None,
824
            end: None,
825
        });
826
        data.events.insert(
827
            identity,
828
            SavedEvent {
829
                item,
830
                media,
831
                raw: Some(value),
832
            },
833
        );
834
    }
835
836
    fn supplement(
837
        &self,
838
        data: &mut Data,
839
        key: &MentionKey,
840
        item: &TimelineItem,
841
        media: &[MediaRegistration],
842
        raw: &Value,
843
    ) {
844
        let TimelineItem::Event(event) = item else {
845
            return;
846
        };
847
        let identity = id(key);
848
        if data
849
            .events
850
            .get(&identity)
851
            .is_some_and(|saved| saved.item != *item)
852
        {
853
            data.touch_contexts(&identity);
854
        }
855
        let mut retained = data.events.contains_key(&identity);
856
        let saved_events = &data.events;
857
        for (target, window) in &mut data.windows {
858
            let Some(entry) = data.entries.get_mut(target) else {
859
                continue;
860
            };
861
            if entry.key.room_id != key.room_id
862
                || (entry.timestamp_ms, target.as_str()) >= (event.timestamp_ms, identity.as_str())
863
                || window.after.contains(&identity)
864
            {
865
                continue;
866
            }
867
            let previous = window.after.clone();
868
            window.after.push(identity.clone());
869
            window.after.sort_by_key(|key| {
870
                if key == &identity {
871
                    (event.timestamp_ms, key.clone())
872
                } else {
873
                    (
874
                        saved_events
875
                            .get(key)
876
                            .and_then(|saved| match &saved.item {
877
                                TimelineItem::Event(event) => Some(event.timestamp_ms),
878
                                _ => None,
879
                            })
880
                            .unwrap_or(0),
881
                        key.clone(),
882
                    )
883
                }
884
            });
885
            window.after.truncate(10);
886
            if window.after == previous {
887
                continue;
888
            }
889
            window.end = None;
890
            retained |= window.after.contains(&identity);
891
            if entry.capture == MentionCaptureState::Partial
892
                && window.after.len() == 10
893
                && (window.before.len() == 10 || window.start.is_none())
894
            {
895
                entry.capture = MentionCaptureState::Ready;
896
            }
897
            entry.revision = data.revision + 1;
898
        }
899
        if retained {
900
            data.events.insert(
901
                identity,
902
                SavedEvent {
903
                    item: item.clone(),
904
                    media: media.to_vec(),
905
                    raw: Some(raw.clone()),
906
                },
907
            );
908
        }
909
    }
910
911
    async fn classify(
912
        &self,
913
        data: &Data,
914
        room_id: &OwnedRoomId,
915
        value: &Value,
916
    ) -> (Vec<MentionReason>, bool, bool) {
917
        let Some(sender) = value["sender"]
918
            .as_str()
919
            .and_then(|sender| sender.parse::<OwnedUserId>().ok())
920
        else {
921
            return (vec![], false, false);
922
        };
923
        if sender == self.account
924
            || self.ignored.contains(&sender)
925
            || value["type"] != "m.room.message"
926
        {
927
            return (vec![], false, false);
928
        }
929
        let Some(room) = self.observable(room_id) else {
930
            return (vec![], false, false);
931
        };
932
        let raw: Raw<AnySyncTimelineEvent> =
933
            Raw::from_json_string(value.to_string()).expect("serialized event");
934
        let actions = room.event_push_actions(&raw).await.ok().flatten();
935
        let notify = actions.as_ref().is_some_and(|actions| {
936
            actions
937
                .iter()
938
                .any(|action| matches!(action, ruma::push::Action::Notify))
939
        });
940
        let highlight = actions.as_ref().is_some_and(|actions| actions.iter().any(|action| matches!(action, ruma::push::Action::SetTweak(ruma::push::Tweak::Highlight(value)) if value != &ruma::push::HighlightTweakValue::No)));
941
        let direct = value["content"]["m.mentions"]["user_ids"]
942
            .as_array()
943
            .is_some_and(|users| {
944
                users
945
                    .iter()
946
                    .any(|user| user.as_str() == Some(self.account.as_str()))
947
            });
948
        let mut reasons = Vec::new();
949
        if direct {
950
            reasons.push(MentionReason::Direct);
951
        }
952
        let mut unresolved = false;
953
        if let Some(reply) = value["content"]["m.relates_to"]["m.in_reply_to"]["event_id"]
954
            .as_str()
955
            .and_then(|value| value.parse().ok())
956
        {
957
            let target = data.events.get(&id(&MentionKey {
958
                room_id: room_id.clone(),
959
                event_id: reply,
960
            }));
961
            let reply_to_self = value["unsigned"]["macltok_reply_sender"].as_str() == Some(self.account.as_str()) || target.is_some_and(|target| matches!(&target.item, TimelineItem::Event(event) if event.sender.user_id == self.account));
962
            if reply_to_self && (direct || (notify && highlight)) {
963
                reasons.push(MentionReason::Reply);
964
            } else if target.is_none()
965
                && value["unsigned"]["macltok_reply_sender"].as_str().is_none()
966
                && (direct || (notify && highlight))
967
            {
968
                unresolved = true;
969
            }
970
        }
971
        if value["content"]["m.mentions"]["room"] == true {
972
            match room.power_levels().await {
973
                Ok(levels) if levels.user_can_trigger_room_notification(&sender) => {
974
                    if notify && highlight {
975
                        reasons.push(MentionReason::Room);
976
                    } else if actions.is_none() {
977
                        unresolved = true;
978
                    }
979
                }
980
                Err(_) => unresolved = true,
981
                _ => {}
982
            }
983
        }
984
        (reasons, !notify, unresolved)
985
    }
986
987
    pub(crate) async fn run(mut self) {
988
        let reports_client = self.client.clone();
989
        let reports = reports_client
990
            .event_cache()
991
            .subscribe_to_decryption_reports();
992
        futures_util::pin_mut!(reports);
993
        let mut ignores = self.client.subscribe_to_ignore_user_list_changes();
994
        let mut work: FuturesUnordered<EngineBoxFuture<'static, Work>> = FuturesUnordered::new();
995
        let mut captures = BTreeSet::new();
996
        let mut retries = BTreeSet::new();
997
        let mut history = false;
998
        let mut next_scan = now() + 60_000;
999
        let mut next_expiry = 0;
1000
        let mut tick = Box::pin(matrix_sdk::sleep::sleep(Duration::from_secs(1)));
1001
        loop {
1002
            self.schedule(&mut work, &mut captures, &mut retries, &mut history);
1003
            tokio::select! {
1004
                biased;
1005
                request = self.requests.recv() => match request {
1006
                    Some(Request::Stop(reply)) => {
1007
                        if self.dirty && let Err(error) = self.flush().await {
1008
                            tracing::error!(%error, "saving the mentions inbox at shutdown failed");
1009
                        }
1010
                        let _ = reply.send(());
1011
                        return;
1012
                    }
1013
                    Some(Request::Batch(batch, reply)) => { self.batch(batch).await; let _ = reply.send(()); }
1014
                    Some(Request::Command(command, reply)) => { let result = self.command(command).await; let _ = reply.send(result); }
1015
                    None => return,
1016
                },
1017
                Some(result) = work.next(), if !work.is_empty() => {
1018
                    let mut next = std::mem::take(&mut self.data);
1019
                    // a candidate that failed again only moved its retry time, which needs no write
1020
                    let changed = self.complete(&mut next, result, &mut captures, &mut retries, &mut history).await;
1021
                    self.data = next;
1022
                    self.dirty |= changed;
1023
                }
1024
                Some(report) = reports.next() => {
1025
                    let room = match report { Ok(RedecryptorReport::ResolvedUtds { room_id, .. }) => Some(room_id), _ => None };
1026
                    for candidate in self.data.candidates.values_mut().filter(|candidate| candidate.encrypted() && room.as_ref().is_none_or(|room| *room == candidate.room_id)) { candidate.retry_at_ms = 0; }
1027
                }
1028
                Some(_) = ignores.next() => {
1029
                    self.ignored = self.ignored_users().await;
1030
                    let mut next = std::mem::take(&mut self.data);
1031
                    self.dirty |= self.filter_ignored(&mut next);
1032
                    self.data = next;
1033
                }
1034
                () = &mut tick => {
1035
                    tick = Box::pin(matrix_sdk::sleep::sleep(Duration::from_secs(1)));
1036
                    if now() >= next_scan && !history { self.recovery_budget = 20; self.latest_scan = true; next_scan = now() + 60_000; }
1037
                    if now() >= next_expiry { self.dirty |= self.data.expire_candidates(now()); next_expiry = now() + 60_000; }
1038
                    if self.dirty && now() >= self.flush_after && let Err(error) = self.flush().await { self.store_failed(error).await; }
1039
                }
1040
            }
1041
        }
1042
    }
1043
1044
    async fn complete(
1045
        &mut self,
1046
        next: &mut Data,
1047
        result: Work,
1048
        captures: &mut BTreeSet<String>,
1049
        retries: &mut BTreeSet<String>,
1050
        history: &mut bool,
1051
    ) -> bool {
1052
        match result {
1053
            Work::Capture(key, revision, result) => {
1054
                captures.remove(&key);
1055
                let observed = result
1056
                    .as_ref()
1057
                    .ok()
1058
                    .map(|(_, _, observed)| observed.clone())
1059
                    .unwrap_or_default();
1060
                let room_id = next
1061
                    .entries
1062
                    .get(&key)
1063
                    .filter(|entry| entry.revision == revision)
1064
                    .map(|entry| entry.key.room_id.clone())
1065
                    .filter(|room_id| {
1066
                        self.client
1067
                            .get_room(room_id)
1068
                            .is_some_and(|room| room.state() == RoomState::Joined)
1069
                    });
1070
                self.captured(next, key, revision, result);
1071
                if let Some(room_id) = room_id {
1072
                    for value in observed {
1073
                        let identity = value["event_id"]
1074
                            .as_str()
1075
                            .map(|event_id| format!("{room_id}\n{event_id}"));
1076
                        if identity.as_ref().is_some_and(|key| {
1077
                            next.events
1078
                                .get(key)
1079
                                .is_some_and(|event| event.raw.is_some())
1080
                        }) {
1081
                            continue;
1082
                        }
1083
                        self.observe(next, &room_id, value).await;
1084
                    }
1085
                }
1086
            }
1087
            Work::Candidate(key, result) => {
1088
                retries.remove(&key);
1089
                if let Some(mut candidate) = next.candidates.remove(&key) {
1090
                    match result {
1091
                        Ok(value) => {
1092
                            self.observe(next, &candidate.room_id, value).await;
1093
                            if let Some(unresolved) = next.candidates.get_mut(&key) {
1094
                                unresolved.attempts = candidate.attempts.saturating_add(1);
1095
                                unresolved.first_seen_ms = candidate.first_seen_ms;
1096
                                unresolved.retry_at_ms =
1097
                                    now() + (1000u64 << unresolved.attempts.min(8));
1098
                            }
1099
                        }
1100
                        Err(error) => {
1101
                            if !candidate.encrypted()
1102
                                && (matches!(
1103
                                    error.kind,
1104
                                    ErrorKind::EventNotFound | ErrorKind::Forbidden
1105
                                ) || error
1106
                                    .server
1107
                                    .as_ref()
1108
                                    .is_some_and(|server| matches!(server.status, 403 | 404)))
1109
                            {
1110
                                next.gap(
1111
                                    Some(candidate.room_id),
1112
                                    MentionGapReason::ContextUnavailable,
1113
                                    None,
1114
                                );
1115
                            } else {
1116
                                candidate.attempts = candidate.attempts.saturating_add(1);
1117
                                candidate.retry_at_ms =
1118
                                    retry_at(&error, 1000u64 << candidate.attempts.min(8));
1119
                                next.candidates.insert(key, candidate);
1120
                                return false;
1121
                            }
1122
                        }
1123
                    }
1124
                }
1125
            }
1126
            Work::History(latest, from, result) => {
1127
                *history = false;
1128
                self.recovery_budget = self.recovery_budget.saturating_sub(1);
1129
                match result {
1130
                    Ok(page) => {
1131
                        next.coverage.historical_supported = Some(true);
1132
                        let mut reached_anchor = false;
1133
                        for notification in page.notifications {
1134
                            if let Some(event_id) = notification
1135
                                .event
1136
                                .get_field::<String>("event_id")
1137
                                .ok()
1138
                                .flatten()
1139
                            {
1140
                                let identity = format!("{}\n{}", notification.room_id, event_id);
1141
                                if latest && from.is_none() && next.latest_head.is_none() {
1142
                                    next.latest_head = Some(identity.clone());
1143
                                }
1144
                                if latest && next.latest_anchor.as_ref() == Some(&identity) {
1145
                                    reached_anchor = true;
1146
                                    break;
1147
                                }
1148
                            }
1149
                            if self
1150
                                .client
1151
                                .get_room(&notification.room_id)
1152
                                .is_some_and(|room| room.state() == RoomState::Joined)
1153
                                && let Ok(value) =
1154
                                    serde_json::from_str(notification.event.json().get())
1155
                            {
1156
                                self.observe(next, &notification.room_id, value).await;
1157
                            }
1158
                        }
1159
                        let token = page
1160
                            .next_token
1161
                            .filter(|token| Some(token) != from.as_ref() && !reached_anchor);
1162
                        if latest {
1163
                            next.latest_cursor = token.clone();
1164
                            if token.is_none() {
1165
                                self.latest_scan = false;
1166
                                next.latest_anchor =
1167
                                    next.latest_head.take().or(next.latest_anchor.take());
1168
                            }
1169
                        } else {
1170
                            next.history_cursor = token.clone();
1171
                            next.history_exhausted = token.is_none();
1172
                        }
1173
                    }
1174
                    Err(error) => {
1175
                        if error
1176
                            .server
1177
                            .as_ref()
1178
                            .is_some_and(|server| matches!(server.status, 404 | 405 | 501))
1179
                        {
1180
                            next.coverage.historical_supported = Some(false);
1181
                            self.recovery_budget = 0;
1182
                        }
1183
                        if error
1184
                            .server
1185
                            .as_ref()
1186
                            .is_some_and(|server| server.status == 400)
1187
                        {
1188
                            if latest {
1189
                                next.latest_cursor = None;
1190
                                next.latest_head = None;
1191
                            } else {
1192
                                next.history_cursor = None;
1193
                            }
1194
                        }
1195
                        next.gap(None, MentionGapReason::NotificationHistoryUnavailable, from);
1196
                        self.recovery_retry = retry_at(&error, 60_000);
1197
                    }
1198
                }
1199
                next.coverage.recovering = self.recovery_budget > 0
1200
                    && (self.latest_scan || !next.history_exhausted)
1201
                    && next.coverage.historical_supported != Some(false);
1202
            }
1203
        }
1204
        true
1205
    }
1206
1207
    fn schedule(
1208
        &mut self,
1209
        work: &mut FuturesUnordered<EngineBoxFuture<'static, Work>>,
1210
        captures: &mut BTreeSet<String>,
1211
        retries: &mut BTreeSet<String>,
1212
        history: &mut bool,
1213
    ) {
1214
        for (key, entry) in self.data.entries.iter().filter(|(_, entry)| matches!(entry.capture, MentionCaptureState::Pending) || matches!(entry.capture, MentionCaptureState::Failed {retry_at_ms, ..} if retry_at_ms <= now())).take(128) {
1215
            if captures.len() >= 2 { break; }
1216
            if captures.contains(key) { continue; }
1217
            let Some(room) = self.client.get_room(&entry.key.room_id).filter(|room| room.state() == RoomState::Joined) else { continue; };
1218
            captures.insert(key.clone());
1219
            let key = key.clone(); let event_id = entry.key.event_id.clone(); let account = self.account.clone(); let revision = entry.revision;
1220
            work.push(Box::pin(async move { Work::Capture(key, revision, context::capture(room, Some(account), event_id, 10, 10).await) }));
1221
        }
1222
        for (key, candidate) in &self.data.candidates {
1223
            if retries.len() >= 2 {
1224
                break;
1225
            }
1226
            if candidate.retry_at_ms > now() || retries.contains(key) {
1227
                continue;
1228
            }
1229
            let Some(room) = self
1230
                .client
1231
                .get_room(&candidate.room_id)
1232
                .filter(|room| room.state() == RoomState::Joined)
1233
            else {
1234
                continue;
1235
            };
1236
            let key = key.clone();
1237
            let candidate = candidate.clone();
1238
            let key_for_work = key.clone();
1239
            retries.insert(key);
1240
            work.push(Box::pin(async move {
1241
                let result = async {
1242
                    if candidate.raw["type"] == "m.room.encrypted" {
1243
                        let raw: Raw<OriginalSyncRoomEncryptedEvent> =
1244
                            Raw::from_json_string(candidate.raw.to_string()).map_err(|error| {
1245
                                Error::with_source(
1246
                                    ErrorKind::Invalid,
1247
                                    "invalid encrypted mention candidate",
1248
                                    error,
1249
                                )
1250
                            })?;
1251
                        let event = room.decrypt_event(&raw, None).await.map_err(|error| {
1252
                            Error::from_sdk("decrypting a mention candidate failed", error)
1253
                        })?;
1254
                        let value: Value =
1255
                            serde_json::from_str(event.raw().json().get()).map_err(|error| {
1256
                                Error::with_source(
1257
                                    ErrorKind::Invalid,
1258
                                    "invalid decrypted mention",
1259
                                    error,
1260
                                )
1261
                            })?;
1262
                        if value["type"] == "m.room.encrypted" {
1263
                            return Err(Error::new(
1264
                                ErrorKind::Encryption,
1265
                                "mention candidate still awaits keys",
1266
                            ));
1267
                        }
1268
                        return Ok(value);
1269
                    }
1270
                    let target = candidate.raw["content"]["m.relates_to"]["event_id"]
1271
                        .as_str()
1272
                        .filter(|_| {
1273
                            candidate.raw["content"]["m.relates_to"]["rel_type"] == "m.replace"
1274
                        });
1275
                    if let Some(target) =
1276
                        target.and_then(|target| target.parse::<ruma::OwnedEventId>().ok())
1277
                    {
1278
                        let event = room.event(&target, None).await.map_err(|error| {
1279
                            Error::from_sdk("fetching an edited mention failed", error)
1280
                        })?;
1281
                        let original: Value = serde_json::from_str(event.raw().json().get())
1282
                            .map_err(|error| {
1283
                                Error::with_source(
1284
                                    ErrorKind::Invalid,
1285
                                    "invalid edited mention",
1286
                                    error,
1287
                                )
1288
                            })?;
1289
                        if original["sender"] == candidate.raw["sender"] {
1290
                            let mut original = edited(original, &candidate.raw);
1291
                            original["unsigned"]["macltok_edit"] = candidate.raw;
1292
                            return Ok(original);
1293
                        }
1294
                        return Err(Error::new(
1295
                            ErrorKind::Forbidden,
1296
                            "an edit belongs to a different sender",
1297
                        ));
1298
                    }
1299
                    let mut value = candidate.raw;
1300
                    if let Some(target) =
1301
                        value["content"]["m.relates_to"]["m.in_reply_to"]["event_id"]
1302
                            .as_str()
1303
                            .and_then(|target| target.parse::<ruma::OwnedEventId>().ok())
1304
                    {
1305
                        let event = room.event(&target, None).await.map_err(|error| {
1306
                            Error::from_sdk("fetching a mention reply target failed", error)
1307
                        })?;
1308
                        let target: Value = serde_json::from_str(event.raw().json().get())
1309
                            .map_err(|error| {
1310
                                Error::with_source(
1311
                                    ErrorKind::Invalid,
1312
                                    "invalid mention reply target",
1313
                                    error,
1314
                                )
1315
                            })?;
1316
                        value["unsigned"]["macltok_reply_sender"] = target["sender"].clone();
1317
                    }
1318
                    Ok(value)
1319
                }
1320
                .await;
1321
                Work::Candidate(key_for_work, result)
1322
            }));
1323
        }
1324
        if !*history
1325
            && self.recovery_budget > 0
1326
            && now() >= self.recovery_retry
1327
            && (self.latest_scan || !self.data.history_exhausted)
1328
            && self.data.coverage.historical_supported != Some(false)
1329
        {
1330
            *history = true;
1331
            let latest = self.latest_scan;
1332
            let from = if latest {
1333
                self.data.latest_cursor.clone()
1334
            } else {
1335
                self.data.history_cursor.clone()
1336
            };
1337
            let client = self.client.clone();
1338
            work.push(Box::pin(async move {
1339
                let mut request = ruma::api::client::push::get_notifications::v3::Request::new();
1340
                request.from = from.clone();
1341
                request.limit = Some(ruma::uint!(50));
1342
                let result = client.send(request).await.map_err(|error| {
1343
                    Error::from_sdk("fetching mention history failed", error.into())
1344
                });
1345
                Work::History(latest, from, result)
1346
            }));
1347
        }
1348
    }
1349
1350
    fn captured(
1351
        &self,
1352
        data: &mut Data,
1353
        key: String,
1354
        revision: u64,
1355
        result: Result<(crate::EventContext, Vec<MediaRegistration>, Vec<Value>), Error>,
1356
    ) {
1357
        let Some(entry) = data.entries.get(&key) else {
1358
            return;
1359
        };
1360
        if entry.revision != revision {
1361
            return;
1362
        }
1363
        let room_id = entry.key.room_id.clone();
1364
        if !self
1365
            .client
1366
            .get_room(&room_id)
1367
            .is_some_and(|room| room.state() == RoomState::Joined)
1368
        {
1369
            return;
1370
        }
1371
        match result {
1372
            Ok((context, media, _)) => {
1373
                data.capture_attempts.remove(&key);
1374
                let live_after = data
1375
                    .windows
1376
                    .get(&key)
1377
                    .map(|window| window.after.clone())
1378
                    .unwrap_or_default();
1379
                let mut insert = |item: TimelineItem| -> Option<String> {
1380
                    let TimelineItem::Event(event) = &item else {
1381
                        return None;
1382
                    };
1383
                    let event_id = event.event_id.clone()?;
1384
                    let event_key = id(&MentionKey {
1385
                        room_id: room_id.clone(),
1386
                        event_id,
1387
                    });
1388
                    let registrations = media_for_item(&item, &media);
1389
                    let mut saved = SavedEvent {
1390
                        item,
1391
                        media: registrations,
1392
                        raw: None,
1393
                    };
1394
                    if let Some(existing) = data.events.get(&event_key) {
1395
                        saved = existing.clone();
1396
                    }
1397
                    if data.redacted.contains_key(&event_key) {
1398
                        redact(&mut saved);
1399
                    }
1400
                    data.events.insert(event_key.clone(), saved);
1401
                    Some(event_key)
1402
                };
1403
                let before: Vec<_> = context
1404
                    .events_before
1405
                    .into_iter()
1406
                    .filter_map(&mut insert)
1407
                    .collect();
1408
                let target = insert(context.event);
1409
                let mut after: Vec<_> = context
1410
                    .events_after
1411
                    .into_iter()
1412
                    .filter_map(&mut insert)
1413
                    .collect();
1414
                let extended = live_after.iter().any(|key| !after.contains(key));
1415
                after.extend(live_after);
1416
                after.sort_by_key(|key| {
1417
                    (
1418
                        data.events
1419
                            .get(key)
1420
                            .and_then(|saved| match &saved.item {
1421
                                TimelineItem::Event(event) => Some(event.timestamp_ms),
1422
                                _ => None,
1423
                            })
1424
                            .unwrap_or(0),
1425
                        key.clone(),
1426
                    )
1427
                });
1428
                after.dedup();
1429
                after.truncate(10);
1430
                let entry = data.entries.get_mut(&key).unwrap();
1431
                entry.capture = if target.is_none() {
1432
                    MentionCaptureState::Unavailable
1433
                } else if after.len() < 10 || (before.len() < 10 && context.start_token.is_some()) {
1434
                    MentionCaptureState::Partial
1435
                } else {
1436
                    MentionCaptureState::Ready
1437
                };
1438
                entry.revision = data.revision + 1;
1439
                data.windows.insert(
1440
                    key,
1441
                    Window {
1442
                        before,
1443
                        target,
1444
                        after,
1445
                        start: context.start_token,
1446
                        end: if extended { None } else { context.end_token },
1447
                    },
1448
                );
1449
            }
1450
            Err(error) => {
1451
                let attempts = data.capture_attempts.entry(key.clone()).or_default();
1452
                *attempts = attempts.saturating_add(1);
1453
                let delay = (1000u64 << (*attempts).min(11)).min(3_600_000);
1454
                let entry = data.entries.get_mut(&key).unwrap();
1455
                entry.capture = if matches!(
1456
                    error.kind,
1457
                    ErrorKind::EventNotFound | ErrorKind::Forbidden | ErrorKind::RoomUnknown
1458
                ) || error
1459
                    .server
1460
                    .as_ref()
1461
                    .is_some_and(|server| matches!(server.status, 403 | 404))
1462
                {
1463
                    MentionCaptureState::Unavailable
1464
                } else {
1465
                    MentionCaptureState::Failed {
1466
                        message: error.to_string(),
1467
                        retry_at_ms: retry_at(&error, delay),
1468
                    }
1469
                };
1470
                entry.revision = data.revision + 1;
1471
            }
1472
        }
1473
    }
1474
}
1475
1476
fn edited(mut original: Value, edit: &Value) -> Value {
1477
    original["content"] = edit["content"]["m.new_content"].clone();
1478
    original["unsigned"]["macltok_edited"] = Value::Bool(true);
1479
    original
1480
}
1481
1482
fn retry_at(error: &Error, fallback_ms: u64) -> u64 {
1483
    let mut source = std::error::Error::source(error);
1484
    while let Some(current) = source {
1485
        if let Some(sdk) = current.downcast_ref::<matrix_sdk::Error>() {
1486
            if let Some(ruma::api::error::ErrorKind::LimitExceeded(data)) =
1487
                sdk.client_api_error_kind()
1488
            {
1489
                if let Some(delay) = data.retry_after {
1490
                    let delay = match delay {
1491
                        ruma::api::error::RetryAfter::Delay(delay) => delay,
1492
                        ruma::api::error::RetryAfter::DateTime(time) => time
1493
                            .duration_since(ruma::time::SystemTime::now())
1494
                            .unwrap_or_default(),
1495
                    };
1496
                    return now().saturating_add(
1497
                        u64::try_from(delay.as_millis())
1498
                            .unwrap_or(u64::MAX)
1499
                            .max(fallback_ms),
1500
                    );
1501
                }
1502
            }
1503
        }
1504
        source = current.source();
1505
    }
1506
    now().saturating_add(fallback_ms)
1507
}
1508
1509
fn redact(event: &mut SavedEvent) {
1510
    if let TimelineItem::Event(item) = &mut event.item {
1511
        item.content = TimelineContent::Redacted;
1512
        item.reply_to = None;
1513
        item.reactions.clear();
1514
    }
1515
    event.media.clear();
1516
    event.raw = None;
1517
}
1518
1519
fn media_for_item(
1520
    item: &TimelineItem,
1521
    registrations: &[MediaRegistration],
1522
) -> Vec<MediaRegistration> {
1523
    let TimelineItem::Event(event) = item else {
1524
        return Vec::new();
1525
    };
1526
    let mut ids = Vec::new();
1527
    match &event.content {
1528
        TimelineContent::Message(message) => {
1529
            ids.extend(message.media.iter().map(|media| &media.id));
1530
            ids.extend(
1531
                message
1532
                    .link_previews
1533
                    .iter()
1534
                    .filter_map(|preview| preview.image.as_ref().map(|image| &image.media_id)),
1535
            );
1536
        }
1537
        TimelineContent::Sticker { media, .. } => ids.push(&media.id),
1538
        _ => {}
1539
    }
1540
    registrations
1541
        .iter()
1542
        .filter(|registration| ids.contains(&&registration.id))
1543
        .cloned()
1544
        .collect()
1545
}
1546
1547
pub(crate) async fn watch_updates(mut updates: watch::Receiver<MentionsPage>, events: EventSender) {
1548
    let mut previous: BTreeMap<String, MentionEntry> = BTreeMap::new();
1549
    loop {
1550
        let snapshot = updates.borrow_and_update().clone();
1551
        let current: BTreeMap<_, _> = snapshot
1552
            .entries
1553
            .iter()
1554
            .map(|entry| (id(&entry.key), entry.clone()))
1555
            .collect();
1556
        let upserts = current
1557
            .iter()
1558
            .filter(|(key, entry)| previous.get(*key) != Some(entry))
1559
            .map(|(_, entry)| entry.clone())
1560
            .collect();
1561
        let removed = previous
1562
            .iter()
1563
            .filter(|(key, _)| !current.contains_key(*key))
1564
            .map(|(_, entry)| entry.key.clone())
1565
            .collect();
1566
        previous = current;
1567
        if events
1568
            .send(Event::MentionsChanged(MentionsUpdate {
1569
                snapshot,
1570
                upserts,
1571
                removed,
1572
            }))
1573
            .await
1574
            .is_err()
1575
        {
1576
            return;
1577
        }
1578
        if updates.changed().await.is_err() {
1579
            return;
1580
        }
1581
    }
1582
}