atagen

macltok

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