atagen

macltok

public
Star0
Watch0
9f95c6462d606393c70c65d478107c0fe4496177
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git
src/engine/url_preview.rsRust166 lines5.5 KB
1
use std::{
2
    collections::{HashMap, VecDeque},
3
    sync::{
4
        Arc,
5
        atomic::{AtomicBool, Ordering},
6
    },
7
    time::Duration,
8
};
9
10
use futures_util::{FutureExt, future::WeakShared};
11
use matrix_sdk::{
12
    Client, HttpError,
13
    ruma::api::{
14
        Metadata,
15
        client::{authenticated_media, media},
16
    },
17
};
18
use ruma::{
19
    api::error::ErrorKind as ApiErrorKind,
20
    serde::{JsonObject, Raw},
21
    time::Instant,
22
};
23
use tokio::sync::Mutex;
24
25
use super::executor::EngineBoxFuture;
26
use crate::{Error, LinkPreview, normalise::server_preview};
27
28
const HIT_FOR: Duration = Duration::from_secs(60 * 60);
29
const MISS_FOR: Duration = Duration::from_secs(10 * 60);
30
const REMEMBERED: usize = 512;
31
32
type Lookup = EngineBoxFuture<'static, Result<Option<LinkPreview>, Error>>;
33
34
/// the homeserver rate limits `preview_url` and a busy timeline repeats links, so answers are
35
/// kept a while, misses too, and asks for one url at once share one request
36
pub(crate) struct UrlPreviews(Arc<Inner>);
37
38
struct Inner {
39
    client: Client,
40
    /// a server that doesn't preview at all says so once and isn't asked again
41
    unsupported: AtomicBool,
42
    cache: Mutex<Cache>,
43
    lookups: Mutex<HashMap<String, WeakShared<Lookup>>>,
44
}
45
46
#[derive(Default)]
47
struct Cache {
48
    entries: HashMap<String, (Instant, Option<LinkPreview>)>,
49
    order: VecDeque<String>,
50
}
51
52
impl UrlPreviews {
53
    pub(crate) fn new(client: Client) -> Self {
54
        Self(Arc::new(Inner {
55
            client,
56
            unsupported: AtomicBool::new(false),
57
            cache: Mutex::new(Cache::default()),
58
            lookups: Mutex::new(HashMap::new()),
59
        }))
60
    }
61
62
    pub(crate) async fn preview(&self, url: &str) -> Result<Option<LinkPreview>, Error> {
63
        if self.0.unsupported.load(Ordering::Relaxed) {
64
            return Ok(None);
65
        }
66
        if let Some(preview) = self.0.cache.lock().await.get(url) {
67
            return Ok(preview);
68
        }
69
        let lookup = {
70
            let mut lookups = self.0.lookups.lock().await;
71
            lookups.retain(|_, lookup| lookup.upgrade().is_some());
72
            match lookups.get(url).and_then(WeakShared::upgrade) {
73
                Some(lookup) => lookup,
74
                None => {
75
                    let inner = self.0.clone();
76
                    let owned = url.to_owned();
77
                    let lookup: Lookup = Box::pin(async move { inner.look_up(owned).await });
78
                    let lookup = lookup.shared();
79
                    if let Some(weak) = lookup.downgrade() {
80
                        lookups.insert(url.to_owned(), weak);
81
                    }
82
                    lookup
83
                }
84
            }
85
        };
86
        lookup.await
87
    }
88
}
89
90
impl Inner {
91
    async fn look_up(&self, url: String) -> Result<Option<LinkPreview>, Error> {
92
        let data = match self.fetch(&url).await {
93
            Ok(data) => data,
94
            Err(error) => {
95
                let Some(api) = error.as_client_api_error() else {
96
                    return Err(Error::from_sdk(
97
                        "asking for a url preview failed",
98
                        error.into(),
99
                    ));
100
                };
101
                let unrecognised = matches!(api.error_kind(), Some(ApiErrorKind::Unrecognized))
102
                    || api.status_code.as_u16() == 405;
103
                if unrecognised {
104
                    self.unsupported.store(true, Ordering::Relaxed);
105
                    return Ok(None);
106
                }
107
                // synapse answers 500 for a page it can't download, 502 when dns fails and 403 for
108
                // its blocklist (synapse/media/url_previewer.py). a real fault hides in the 500s too,
109
                // but a missing card is cheap and only cached ten minutes
110
                if !matches!(api.status_code.as_u16(), 400 | 403 | 404 | 500 | 502) {
111
                    return Err(Error::from_sdk(
112
                        "asking for a url preview failed",
113
                        error.into(),
114
                    ));
115
                }
116
                None
117
            }
118
        };
119
        let preview = data.and_then(|data| server_preview(&url, &data));
120
        self.cache.lock().await.insert(url, preview.clone());
121
        Ok(preview)
122
    }
123
124
    async fn fetch(&self, url: &str) -> Result<Option<Raw<JsonObject>>, HttpError> {
125
        let versions = self.client.supported_versions().await?;
126
        let data = if authenticated_media::get_media_preview::v1::Request::PATH_BUILDER
127
            .is_supported(&versions)
128
        {
129
            self.client
130
                .send(authenticated_media::get_media_preview::v1::Request::new(
131
                    url.to_owned(),
132
                ))
133
                .await?
134
                .data
135
        } else {
136
            #[allow(deprecated)]
137
            let request = media::get_media_preview::v3::Request::new(url.to_owned());
138
            self.client.send(request).await?.data
139
        };
140
        Ok(data.map(Raw::from_json))
141
    }
142
}
143
144
impl Cache {
145
    fn get(&mut self, url: &str) -> Option<Option<LinkPreview>> {
146
        let (at, preview) = self.entries.get(url)?;
147
        let lasts = if preview.is_some() { HIT_FOR } else { MISS_FOR };
148
        (at.elapsed() < lasts).then(|| preview.clone())
149
    }
150
151
    fn insert(&mut self, url: String, preview: Option<LinkPreview>) {
152
        if self
153
            .entries
154
            .insert(url.clone(), (Instant::now(), preview))
155
            .is_none()
156
        {
157
            self.order.push_back(url);
158
        }
159
        while self.entries.len() > REMEMBERED {
160
            let Some(oldest) = self.order.pop_front() else {
161
                break;
162
            };
163
            self.entries.remove(&oldest);
164
        }
165
    }
166
}