macltok

public
Star0
Watch0
103 commits 3 branches 0 tags 0 releases3.0 MiB
use std::{
    borrow::Cow,
    collections::{HashMap, VecDeque},
    io::Read as _,
    panic::AssertUnwindSafe,
    sync::Arc,
    time::Duration,
};

use futures_util::{
    FutureExt, StreamExt,
    future::{Shared, WeakShared},
};
use matrix_sdk::{
    Client, Room, SendOutsideWasm, SyncOutsideWasm,
    media::{MediaFormat, MediaRequestParameters, MediaThumbnailSettings},
    room::reply::{EnforceThread, Reply as SdkReply},
    ruma::{
        api::{
            Metadata, OutgoingRequest,
            auth_scheme::SendAccessToken,
            client::{authenticated_media, media},
        },
        events::room::{
            EncryptedFileHash, EncryptedFileHashAlgorithm, EncryptedFileInfo, MediaSource,
            message::{
                AddMentions, Relation, ReplyWithinThread, RoomMessageEventContentWithoutRelation,
            },
        },
    },
};
use matrix_sdk_crypto::{AttachmentDecryptor, MediaEncryptionInfo};
use ruma::{MxcUri, OwnedEventId};
use tokio::sync::Mutex;

use super::{executor::EngineBoxFuture, url_preview::UrlPreviews};
use crate::{
    AttachmentEncryption, AttachmentUpload, Error, ErrorKind, LinkPreview, Media, MediaFile,
    MediaId, MediaKind, MediaLocation, MediaSize,
    normalise::{MediaRegistration, MediaScope, MediaSourceInfo},
    upload::prepare_upload,
};

/// a stalled download gives up after this long without a byte
const DOWNLOAD_IDLE: Duration = Duration::from_secs(60);
/// the sdk client gives up on any request after thirty seconds, far too soon for a video
const DOWNLOAD_TOTAL: Duration = Duration::from_secs(60 * 60);

type FlightFuture = EngineBoxFuture<'static, Result<Media, Error>>;
type Flight = Shared<FlightFuture>;

/// owns everything needed to turn a media id or an mxc uri into bytes, so the
/// handle can answer media queries without a trip through the engine
pub(crate) struct MediaService {
    client: Client,
    registry: Mutex<MediaRegistry>,
    flights: Mutex<HashMap<String, WeakShared<FlightFuture>>>,
    max_bytes: usize,
    context_limit: usize,
    url_previews: UrlPreviews,
}

impl MediaService {
    pub(crate) fn new(client: Client, max_bytes: usize, context_limit: usize) -> Self {
        Self {
            url_previews: UrlPreviews::new(client.clone()),
            client,
            registry: Mutex::new(MediaRegistry::default()),
            flights: Mutex::new(HashMap::new()),
            max_bytes,
            context_limit,
        }
    }

    /// one event context session's media, bounded by `EngineConfig::max_context_media`
    pub(crate) fn context_scope(&self) -> MediaScope {
        MediaScope::bounded(self.context_limit)
    }

    pub(crate) async fn url_preview(&self, url: &str) -> Result<Option<LinkPreview>, Error> {
        self.url_previews.preview(url).await
    }

    pub(crate) async fn register(&self, scope: &MediaScope, media: Vec<MediaRegistration>) {
        if !media.is_empty() {
            self.registry.lock().await.insert_all(scope, media);
        }
    }

    pub(crate) async fn release(&self, scope: &MediaScope) {
        self.registry.lock().await.release(scope);
    }

    pub(crate) async fn clear(&self) {
        self.registry.lock().await.clear();
    }

    pub(crate) async fn timeline_media(
        &self,
        media_id: &MediaId,
        size: MediaSize,
    ) -> Result<Media, Error> {
        let chosen = choose_source(self.registered(media_id).await?, size)?;
        if chosen
            .declared_size
            .is_some_and(|size| size > self.max_bytes as u64)
        {
            return Err(too_large("media exceeds the configured byte limit"));
        }
        self.fetch_or_file(chosen.request, chosen.mimetype).await
    }

    pub(crate) async fn media_location(
        &self,
        media_id: &MediaId,
        size: MediaSize,
    ) -> Result<MediaLocation, Error> {
        let media = self.registered(media_id).await?;
        let chosen = choose_source(media.clone(), size)?;
        if is_local(&chosen.request.source) {
            return Ok(MediaLocation::Local);
        }
        if !matches!(chosen.request.format, MediaFormat::Thumbnail(_)) {
            return Ok(MediaLocation::File(
                media_file(
                    &self.client,
                    chosen.request,
                    chosen.mimetype,
                    chosen.declared_size,
                )
                .await?,
            ));
        }
        let url = media_http_request(&self.client, &chosen.request)
            .await?
            .uri()
            .to_string();
        let fallback = MediaRequestParameters {
            source: media.full.source,
            format: MediaFormat::File,
        };
        Ok(MediaLocation::ServerThumbnail {
            url,
            fallback: media_file(
                &self.client,
                fallback,
                media.full.mimetype,
                media.full.declared_size,
            )
            .await?,
        })
    }

    async fn registered(&self, media_id: &MediaId) -> Result<MediaRegistration, Error> {
        if let Some(uri) = media_id.as_plain() {
            return plain_registration(media_id, uri);
        }
        self.registry
            .lock()
            .await
            .get(media_id)
            .cloned()
            .ok_or_else(|| {
                Error::new(
                    ErrorKind::MediaUnknown,
                    "media is not registered in a retained timeline",
                )
            })
    }

    /// plenty of servers won't thumbnail everything, and the file itself is the next best
    /// thing. the byte cap still holds either way
    async fn fetch_or_file(
        &self,
        request: MediaRequestParameters,
        mimetype: Option<String>,
    ) -> Result<Media, Error> {
        if !matches!(request.format, MediaFormat::Thumbnail(_)) {
            return self.fetch(request, mimetype).await;
        }
        match self.fetch(request.clone(), mimetype.clone()).await {
            Err(error) if error.kind == ErrorKind::MediaUnavailable => {
                let file = MediaRequestParameters {
                    source: request.source,
                    format: MediaFormat::File,
                };
                self.fetch(file, mimetype).await
            }
            result => result,
        }
    }

    async fn fetch(
        &self,
        request: MediaRequestParameters,
        mimetype: Option<String>,
    ) -> Result<Media, Error> {
        let key = flight_key(&request);
        let flight = {
            let mut flights = self.flights.lock().await;
            flights.retain(|_, flight| flight.upgrade().is_some());
            match flights.get(&key).and_then(WeakShared::upgrade) {
                Some(flight) => flight,
                None => {
                    let prepared = AssertUnwindSafe(fetch(
                        self.client.clone(),
                        request,
                        mimetype,
                        self.max_bytes,
                    ))
                    .catch_unwind()
                    .map(|result| {
                        result.unwrap_or_else(|_| {
                            Err(Error::new(ErrorKind::Other, "media fetch panicked"))
                        })
                    });
                    let boxed: FlightFuture = Box::pin(prepared);
                    let flight: Flight = boxed.shared();
                    if let Some(weak) = flight.downgrade() {
                        flights.insert(key, weak);
                    }
                    flight
                }
            }
        };
        flight.await
    }
}

struct ChosenSource {
    request: MediaRequestParameters,
    mimetype: Option<String>,
    declared_size: Option<u64>,
}

fn plain_registration(media_id: &MediaId, uri: &MxcUri) -> Result<MediaRegistration, Error> {
    if !uri.is_valid() {
        return Err(Error::new(ErrorKind::Invalid, "mxc uri is malformed"));
    }
    Ok(MediaRegistration {
        id: media_id.clone(),
        kind: MediaKind::Image,
        full: MediaSourceInfo {
            source: MediaSource::Plain(uri.to_owned()),
            mimetype: None,
            declared_size: None,
        },
        thumbnail: None,
    })
}

fn choose_source(media: MediaRegistration, size: MediaSize) -> Result<ChosenSource, Error> {
    let (source, format, mimetype, declared_size) = match size {
        MediaSize::Full => (
            media.full.source,
            MediaFormat::File,
            media.full.mimetype,
            media.full.declared_size,
        ),
        MediaSize::Thumbnail { width, height } => match media.thumbnail {
            Some(thumbnail) => (
                thumbnail.source,
                MediaFormat::File,
                thumbnail.mimetype,
                thumbnail.declared_size,
            ),
            // the server can only resize what it can read
            None if matches!(media.kind, MediaKind::Image | MediaKind::Video)
                && matches!(media.full.source, MediaSource::Plain(_)) =>
            {
                // the server picks the thumbnail's type, so the original's would be a lie
                let format = if is_local(&media.full.source) {
                    MediaFormat::File
                } else {
                    thumbnail_format(width, height)
                };
                (media.full.source, format, None, None)
            }
            None => {
                return Err(Error::new(
                    ErrorKind::MediaUnavailable,
                    "the requested media size is unavailable",
                ));
            }
        },
    };
    Ok(ChosenSource {
        request: MediaRequestParameters { source, format },
        mimetype,
        declared_size,
    })
}

async fn media_file(
    client: &Client,
    request: MediaRequestParameters,
    mimetype: Option<String>,
    declared_size: Option<u64>,
) -> Result<MediaFile, Error> {
    let url = media_http_request(client, &request)
        .await?
        .uri()
        .to_string();
    let encryption = match &request.source {
        MediaSource::Plain(_) => None,
        MediaSource::Encrypted(file) => {
            let (EncryptedFileInfo::V2(info), Some(EncryptedFileHash::Sha256(sha256))) = (
                &file.info,
                file.hashes.get(&EncryptedFileHashAlgorithm::Sha256),
            ) else {
                return Err(Error::new(
                    ErrorKind::MediaUnavailable,
                    "encrypted media is not v2 with a sha256 hash",
                ));
            };
            Some(AttachmentEncryption {
                key: info.k.encode(),
                iv: info.iv.encode(),
                sha256: sha256.encode(),
            })
        }
    };
    Ok(MediaFile {
        url,
        mimetype,
        declared_size,
        encryption,
    })
}

/// servers without spec v1.11 animated thumbnails ignore the flag and send a still frame
fn thumbnail_format(width: u32, height: u32) -> MediaFormat {
    let mut settings = MediaThumbnailSettings::new(width.into(), height.into());
    settings.animated = true;
    MediaFormat::Thumbnail(settings)
}

fn is_local(source: &MediaSource) -> bool {
    let uri = match source {
        MediaSource::Plain(uri) => uri,
        MediaSource::Encrypted(file) => &file.url,
    };
    uri.server_name()
        .is_ok_and(|server| server.as_str() == "send-queue.localhost")
}

/// plain and encrypted sources share a cache key in the sdk, but must never share a flight
fn flight_key(request: &MediaRequestParameters) -> String {
    let uri = match &request.source {
        MediaSource::Plain(uri) => format!("plain|{uri}"),
        MediaSource::Encrypted(file) => format!("encrypted|{}", file.url),
    };
    match &request.format {
        MediaFormat::File => uri.to_owned(),
        MediaFormat::Thumbnail(settings) => format!(
            "{uri}|{}|{}x{}|{}",
            settings.method, settings.width, settings.height, settings.animated
        ),
    }
}

fn too_large(message: &'static str) -> Error {
    Error::new(ErrorKind::MediaTooLarge, message)
}

pub(crate) async fn download_file(
    client: &Client,
    source: MediaSource,
    mimetype: Option<String>,
    max_bytes: usize,
) -> Result<Media, Error> {
    let request = MediaRequestParameters {
        source,
        format: MediaFormat::File,
    };
    fetch(client.clone(), request, mimetype, max_bytes).await
}

async fn fetch(
    client: Client,
    request: MediaRequestParameters,
    mimetype: Option<String>,
    max_bytes: usize,
) -> Result<Media, Error> {
    if let Some(cached) = cached_media(&client, &request).await? {
        if cached.len() > max_bytes {
            if let Ok(store) = client.media_store().lock().await {
                let _ = store.remove_media_content(&request).await;
            }
            return Err(too_large("cached media exceeds the configured byte limit"));
        }
        // plain and encrypted share a cache key, so a plain fetch of an encrypted file's uri
        // leaves ciphertext here. its hash gives it away
        if is_ciphertext(&cached, &request.source) {
            let bytes = maybe_decrypt(cached, &request, max_bytes)?;
            cache(&client, &request, &bytes).await;
            return Ok(media(bytes, mimetype));
        }
        return Ok(media(cached, mimetype));
    }
    let (bytes, served_mimetype) = download_bounded(&client, &request, max_bytes).await?;
    // the sdk caches plaintext under the encrypted key, so we do too or nobody can read anything
    let bytes = maybe_decrypt(bytes, &request, max_bytes)?;
    cache(&client, &request, &bytes).await;
    Ok(media(bytes, mimetype.or(served_mimetype)))
}

async fn cache(client: &Client, request: &MediaRequestParameters, bytes: &[u8]) {
    if let Ok(store) = client.media_store().lock().await {
        let _ = store
            .add_media_content(
                request,
                bytes.to_vec(),
                matrix_sdk::media::store::IgnoreMediaRetentionPolicy::No,
            )
            .await;
    }
}

fn is_ciphertext(bytes: &[u8], source: &MediaSource) -> bool {
    use sha2::{Digest, Sha256};
    let MediaSource::Encrypted(file) = source else {
        return false;
    };
    file.hashes
        .get(&ruma::events::room::EncryptedFileHashAlgorithm::Sha256)
        .is_some_and(|hash| hash.as_bytes() == Sha256::digest(bytes).as_slice())
}

fn media(bytes: Vec<u8>, mimetype: Option<String>) -> Media {
    let mimetype = mimetype
        .or_else(|| sniff_image(&bytes).map(str::to_owned))
        .unwrap_or_else(|| "application/octet-stream".to_owned());
    Media {
        bytes: Arc::from(bytes),
        mimetype,
    }
}

/// the cache forgets content types, and avatars never declare one. images are the common case
fn sniff_image(bytes: &[u8]) -> Option<&'static str> {
    match bytes {
        [0x89, b'P', b'N', b'G', ..] => Some("image/png"),
        [0xff, 0xd8, 0xff, ..] => Some("image/jpeg"),
        [b'G', b'I', b'F', b'8', ..] => Some("image/gif"),
        [
            b'R',
            b'I',
            b'F',
            b'F',
            _,
            _,
            _,
            _,
            b'W',
            b'E',
            b'B',
            b'P',
            ..,
        ] => Some("image/webp"),
        _ => None,
    }
}

async fn cached_media(
    client: &Client,
    request: &MediaRequestParameters,
) -> Result<Option<Vec<u8>>, Error> {
    fn cache_failed(
        source: impl std::error::Error + SendOutsideWasm + SyncOutsideWasm + 'static,
    ) -> Error {
        Error::with_source(ErrorKind::Store, "reading media cache failed", source)
    }
    client
        .media_store()
        .lock()
        .await
        .map_err(cache_failed)?
        .get_media_content(request)
        .await
        .map_err(cache_failed)
}

fn maybe_decrypt(
    bytes: Vec<u8>,
    request: &MediaRequestParameters,
    max_bytes: usize,
) -> Result<Vec<u8>, Error> {
    let MediaSource::Encrypted(file) = &request.source else {
        return Ok(bytes);
    };
    fn decrypt_failed(
        source: impl std::error::Error + SendOutsideWasm + SyncOutsideWasm + 'static,
    ) -> Error {
        Error::with_source(
            ErrorKind::MediaUnavailable,
            "media decryption failed",
            source,
        )
    }
    let info: MediaEncryptionInfo = (**file).clone().into();
    let mut cursor = std::io::Cursor::new(bytes);
    let mut decryptor = AttachmentDecryptor::new(&mut cursor, info).map_err(decrypt_failed)?;
    let mut decrypted = Vec::new();
    let mut chunk = [0_u8; 8192];
    loop {
        let read = decryptor.read(&mut chunk).map_err(decrypt_failed)?;
        if read == 0 {
            return Ok(decrypted);
        }
        if read > max_bytes.saturating_sub(decrypted.len()) {
            return Err(too_large(
                "decrypted media exceeds the configured byte limit",
            ));
        }
        decrypted.extend_from_slice(&chunk[..read]);
    }
}

/// the sdk's own download has no size cap, so we do it by hand
async fn download_bounded(
    client: &Client,
    request: &MediaRequestParameters,
    max_bytes: usize,
) -> Result<(Vec<u8>, Option<String>), Error> {
    if is_local(&request.source) {
        return Err(Error::new(
            ErrorKind::MediaUnavailable,
            "local media is not in the media cache",
        ));
    }
    let network_failed =
        |source| Error::with_source(ErrorKind::Network, "media download failed", source);
    let mut refreshed = false;
    let response = loop {
        let http_request = media_http_request(client, request).await?;
        let (parts, _) = http_request.into_parts();
        let mut builder = client
            .http_client()
            .get(parts.uri.to_string())
            .timeout(DOWNLOAD_TOTAL);
        for (name, value) in &parts.headers {
            builder = builder.header(name.as_str(), value.as_bytes());
        }
        let response = builder.send().await.map_err(network_failed)?;
        if response.status() != matrix_sdk::reqwest::StatusCode::UNAUTHORIZED {
            break response;
        }
        if refreshed
            || client
                .session_tokens()
                .is_none_or(|tokens| tokens.refresh_token.is_none())
        {
            return Err(Error::new(
                ErrorKind::Unauthenticated,
                "media download was not authorized",
            ));
        }
        client.refresh_access_token().await.map_err(|source| {
            Error::from_refresh_token("refreshing media access token failed", source)
        })?;
        refreshed = true;
    };
    let status = response.status();
    if !status.is_success() {
        return Err(Error::new(
            ErrorKind::MediaUnavailable,
            format!("media download failed with status {status}"),
        ));
    }
    if response
        .content_length()
        .is_some_and(|length| length > max_bytes as u64)
    {
        return Err(too_large("media exceeds the configured byte limit"));
    }
    let mimetype = response
        .headers()
        .get(http::header::CONTENT_TYPE)
        .and_then(|value| value.to_str().ok())
        .map(str::to_owned);
    let mut stream = response.bytes_stream();
    let mut bytes = Vec::new();
    loop {
        let chunk = matrix_sdk::timeout::timeout(stream.next(), DOWNLOAD_IDLE)
            .await
            .map_err(|_| Error::new(ErrorKind::Network, "media download stalled"))?;
        let Some(chunk) = chunk else { break };
        let chunk = chunk.map_err(network_failed)?;
        if chunk.len() > max_bytes.saturating_sub(bytes.len()) {
            return Err(too_large(
                "downloaded media exceeds the configured byte limit",
            ));
        }
        bytes.extend_from_slice(&chunk);
    }
    Ok((bytes, mimetype))
}

async fn media_http_request(
    client: &Client,
    request: &MediaRequestParameters,
) -> Result<http::Request<Vec<u8>>, Error> {
    let versions = client.supported_versions().await.map_err(|source| {
        Error::with_source(
            ErrorKind::Network,
            "reading supported server versions failed",
            source,
        )
    })?;
    let authenticated =
        authenticated_media::get_content::v1::Request::PATH_BUILDER.is_supported(&versions);
    let base = client.homeserver().to_string();
    let token = client.access_token();
    let invalid_uri =
        |source| Error::with_source(ErrorKind::Invalid, "media URI is invalid", source);
    let (uri, thumbnail) = match (&request.source, &request.format) {
        (MediaSource::Encrypted(file), _) => (&file.url, None),
        (MediaSource::Plain(uri), MediaFormat::File) => (uri, None),
        (MediaSource::Plain(uri), MediaFormat::Thumbnail(settings)) => (uri, Some(settings)),
    };
    // each endpoint has its own auth input type, so no generic helper can take them all
    macro_rules! into_http {
        ($endpoint:expr, $token:expr) => {
            $endpoint
                .try_into_http_request::<Vec<u8>>(&base, $token, Cow::Borrowed(&versions))
                .map_err(|source| {
                    Error::with_source(ErrorKind::Other, "media request is invalid", source)
                })?
        };
    }
    let authed = SendAccessToken::IfRequired(token.as_deref().unwrap_or(""));
    #[allow(deprecated)]
    let mut http_request = match (thumbnail, authenticated) {
        (None, true) => into_http!(
            authenticated_media::get_content::v1::Request::from_uri(uri).map_err(invalid_uri)?,
            authed
        ),
        (None, false) => into_http!(
            media::get_content::v3::Request::from_url(uri).map_err(invalid_uri)?,
            SendAccessToken::None
        ),
        (Some(settings), true) => {
            let mut endpoint = authenticated_media::get_content_thumbnail::v1::Request::from_uri(
                uri,
                settings.width,
                settings.height,
            )
            .map_err(invalid_uri)?;
            endpoint.method = Some(settings.method.clone());
            endpoint.animated = Some(settings.animated);
            into_http!(endpoint, authed)
        }
        (Some(settings), false) => {
            let mut endpoint = media::get_content_thumbnail::v3::Request::from_url(
                uri,
                settings.width,
                settings.height,
            )
            .map_err(invalid_uri)?;
            endpoint.method = Some(settings.method.clone());
            endpoint.animated = Some(settings.animated);
            into_http!(endpoint, SendAccessToken::None)
        }
    };
    // legacy endpoints don't ask for a token but plenty of servers want one anyway
    if let Some(token) = token.as_deref() {
        let value = format!("Bearer {token}").parse().map_err(|_| {
            Error::new(ErrorKind::Unauthenticated, "stored access token is invalid")
        })?;
        http_request
            .headers_mut()
            .insert(http::header::AUTHORIZATION, value);
    }
    Ok(http_request)
}

pub(crate) async fn queue_upload(
    room: Room,
    upload: AttachmentUpload,
    max_bytes: usize,
) -> Result<(), Error> {
    let (mut upload, mime, mut config) = prepare_upload(upload, max_bytes)?;
    let thread_root = upload.timeline_id.thread_root().cloned();
    config.reply = match (upload.reply_to.take(), thread_root) {
        (Some(event_id), thread_root) => {
            let reply = || SdkReply {
                event_id: event_id.clone(),
                enforce_thread: if thread_root.is_some() {
                    EnforceThread::Threaded(ReplyWithinThread::Yes)
                } else {
                    EnforceThread::MaybeThreaded
                },
                add_mentions: if upload.notify_replied_to {
                    AddMentions::Yes
                } else {
                    AddMentions::No
                },
            };
            if let Some(thread_root) = &thread_root {
                check_reply_thread(&room, reply(), thread_root).await?;
            }
            Some(reply())
        }
        (None, Some(root_event_id)) => Some(SdkReply {
            event_id: root_event_id,
            enforce_thread: EnforceThread::Threaded(ReplyWithinThread::No),
            add_mentions: AddMentions::No,
        }),
        (None, None) => None,
    };
    room.send_queue()
        .send_attachment(upload.filename, mime, upload.bytes, config)
        .await
        .map(drop)
        .map_err(|source| {
            use matrix_sdk::send_queue::RoomSendQueueError;
            let kind = match &source {
                RoomSendQueueError::RoomNotJoined | RoomSendQueueError::RoomDisappeared => {
                    ErrorKind::RoomUnknown
                }
                RoomSendQueueError::StorageError(_) => ErrorKind::Store,
                _ => ErrorKind::MediaUnavailable,
            };
            Error::with_source(kind, "media upload could not be queued", source)
        })
}

/// the send queue builds the relation itself, so a reply that would land in another thread has to be caught before it is queued
async fn check_reply_thread(
    room: &Room,
    reply: SdkReply,
    thread_root: &OwnedEventId,
) -> Result<(), Error> {
    let content = room
        .make_reply_event(
            RoomMessageEventContentWithoutRelation::text_plain(""),
            reply,
        )
        .await
        .map_err(|source| {
            Error::with_source(ErrorKind::Timeline, "reply operation failed", source)
        })?;
    if matches!(
        content.relates_to.as_ref(),
        Some(Relation::Thread(thread)) if &thread.event_id == thread_root
    ) {
        Ok(())
    } else {
        Err(Error::new(
            ErrorKind::Invalid,
            "reply target does not belong to the selected thread",
        ))
    }
}

/// every entry belongs to the scopes that registered it and lives until the last of them
/// lets go. there is no global cap, a bounded scope only ever trims itself
#[derive(Default)]
struct MediaRegistry {
    entries: HashMap<MediaId, Entry>,
    scopes: HashMap<u64, ScopeEntries>,
}

struct Entry {
    media: MediaRegistration,
    owners: Vec<u64>,
}

/// least recently used first. stale order slots are skipped by stamp rather than hunted down
#[derive(Default)]
struct ScopeEntries {
    stamps: HashMap<MediaId, u64>,
    order: VecDeque<(MediaId, u64)>,
    next_stamp: u64,
    limit: Option<usize>,
}

impl ScopeEntries {
    fn touch(&mut self, id: &MediaId) {
        let Some(limit) = self.limit else {
            self.stamps.insert(id.clone(), 0);
            return;
        };
        let stamp = self.next_stamp;
        self.next_stamp += 1;
        self.stamps.insert(id.clone(), stamp);
        self.order.push_back((id.clone(), stamp));
        if self.order.len() > limit.saturating_mul(4) {
            let stamps = &self.stamps;
            self.order
                .retain(|(id, stamp)| stamps.get(id) == Some(stamp));
        }
    }

    /// what fell out past the limit
    fn trim(&mut self) -> Vec<MediaId> {
        let Some(limit) = self.limit else {
            return Vec::new();
        };
        let mut dropped = Vec::new();
        while self.stamps.len() > limit {
            let Some((id, stamp)) = self.order.pop_front() else {
                break;
            };
            if self.stamps.get(&id) == Some(&stamp) {
                self.stamps.remove(&id);
                dropped.push(id);
            }
        }
        dropped
    }
}

impl MediaRegistry {
    fn get(&mut self, id: &MediaId) -> Option<&MediaRegistration> {
        let entry = self.entries.get(id)?;
        for owner in &entry.owners {
            if let Some(scope) = self.scopes.get_mut(owner) {
                scope.touch(id);
            }
        }
        Some(&entry.media)
    }

    fn insert_all(&mut self, scope: &MediaScope, media: Vec<MediaRegistration>) {
        let owner = scope.nonce();
        let entries = self.scopes.entry(owner).or_insert_with(|| ScopeEntries {
            limit: scope.limit(),
            ..ScopeEntries::default()
        });
        for media in media {
            let id = media.id.clone();
            entries.touch(&id);
            let entry = self.entries.entry(id).or_insert_with(|| Entry {
                media: media.clone(),
                owners: Vec::new(),
            });
            entry.media = media;
            if !entry.owners.contains(&owner) {
                entry.owners.push(owner);
            }
        }
        for id in entries.trim() {
            Self::disown(&mut self.entries, &id, owner);
        }
    }

    fn disown(entries: &mut HashMap<MediaId, Entry>, id: &MediaId, owner: u64) {
        let Some(entry) = entries.get_mut(id) else {
            return;
        };
        entry.owners.retain(|current| *current != owner);
        if entry.owners.is_empty() {
            entries.remove(id);
        }
    }

    fn release(&mut self, scope: &MediaScope) {
        let Some(released) = self.scopes.remove(&scope.nonce()) else {
            return;
        };
        for id in released.stamps.keys() {
            Self::disown(&mut self.entries, id, scope.nonce());
        }
    }

    fn clear(&mut self) {
        self.entries.clear();
        self.scopes.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::MediaKind;

    fn media(name: &str) -> MediaRegistration {
        let uri = ruma::OwnedMxcUri::from(format!("mxc://example.org/{name}"));
        MediaRegistration {
            id: MediaId::new(format!("test#{name}")),
            kind: MediaKind::Image,
            full: MediaSourceInfo {
                source: MediaSource::Plain(uri),
                mimetype: None,
                declared_size: None,
            },
            thumbnail: None,
        }
    }

    fn has(registry: &mut MediaRegistry, name: &str) -> bool {
        registry.get(&media(name).id).is_some()
    }

    #[test]
    fn shared_media_lives_until_every_owner_lets_go() {
        let mut registry = MediaRegistry::default();
        let (timeline, context) = (MediaScope::new(), MediaScope::bounded(8));
        registry.insert_all(&timeline, vec![media("shared"), media("row")]);
        registry.insert_all(&context, vec![media("shared"), media("page")]);
        registry.release(&context);
        assert!(has(&mut registry, "shared"));
        assert!(has(&mut registry, "row"));
        assert!(!has(&mut registry, "page"));
        registry.release(&timeline);
        assert!(!has(&mut registry, "shared"));
        assert!(registry.entries.is_empty());
    }

    #[test]
    fn heavy_context_paging_trims_only_its_own_least_recent_entries() {
        let mut registry = MediaRegistry::default();
        let (timeline, context) = (MediaScope::new(), MediaScope::bounded(4));
        registry.insert_all(&timeline, vec![media("on-screen"), media("both")]);
        registry.insert_all(&context, vec![media("both"), media("kept")]);
        for page in 0..50 {
            assert!(has(&mut registry, "kept"));
            registry.insert_all(&context, vec![media(&format!("page-{page}"))]);
        }
        assert!(has(&mut registry, "on-screen"));
        assert!(has(&mut registry, "both"));
        assert!(has(&mut registry, "kept"));
        assert!(!has(&mut registry, "page-0"));
        assert!(registry.scopes[&context.nonce()].stamps.len() <= 4);
        registry.release(&context);
        assert!(has(&mut registry, "both"));
        assert!(!has(&mut registry, "kept"));
    }
}