atagen

macltok

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