macltok

public

Star0

Watch0

7 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::{
            MediaSource,
            message::{AddMentions, ReplyWithinThread},
        },
    },
};
use matrix_sdk_crypto::{AttachmentDecryptor, MediaEncryptionInfo};
use ruma::MxcUri;
use tokio::sync::Mutex;

use super::executor::EngineBoxFuture;
use crate::{
    AttachmentUpload, Error, ErrorKind, Media, MediaId, MediaKind, MediaSize, TimelineId,
    normalise::{MediaRegistration, MediaScope},
    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,
}

impl MediaService {
    pub(crate) fn new(client: Client, max_bytes: usize, max_sources: usize) -> Self {
        Self {
            client,
            registry: Mutex::new(MediaRegistry::new(max_sources)),
            flights: Mutex::new(HashMap::new()),
            max_bytes,
        }
    }

    pub(crate) async fn register(&self, media: Vec<MediaRegistration>) {
        if !media.is_empty() {
            self.registry.lock().await.insert_all(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 media = self
            .registry
            .lock()
            .await
            .get(media_id)
            .cloned()
            .ok_or_else(|| {
                Error::new(
                    ErrorKind::MediaUnknown,
                    "media is not registered in a retained timeline",
                )
            })?;
        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",
                    ));
                }
            },
        };
        if 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(MediaRequestParameters { source, format }, mimetype)
            .await
    }

    pub(crate) async fn mxc_media(&self, uri: &MxcUri, size: MediaSize) -> Result<Media, Error> {
        if !uri.is_valid() {
            return Err(Error::new(ErrorKind::Invalid, "mxc uri is malformed"));
        }
        let format = match size {
            MediaSize::Full => MediaFormat::File,
            MediaSize::Thumbnail { width, height } => thumbnail_format(width, height),
        };
        self.fetch_or_file(
            MediaRequestParameters {
                source: MediaSource::Plain(uri.to_owned()),
                format,
            },
            None,
        )
        .await
    }

    /// 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
    }
}

fn thumbnail_format(width: u32, height: u32) -> MediaFormat {
    MediaFormat::Thumbnail(MediaThumbnailSettings::new(width.into(), height.into()))
}

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)
}

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 (upload, mime, mut config) = prepare_upload(upload, max_bytes)?;
    if let TimelineId::Thread { root_event_id, .. } = &upload.timeline_id {
        config.reply = Some(SdkReply {
            event_id: root_event_id.clone(),
            enforce_thread: EnforceThread::Threaded(ReplyWithinThread::No),
            add_mentions: AddMentions::No,
        });
    }
    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)
        })
}

/// least recently seen goes first. stale order slots are skipped by stamp rather than hunted down
struct MediaRegistry {
    entries: HashMap<MediaId, (MediaRegistration, u64)>,
    order: VecDeque<(MediaId, u64)>,
    next_stamp: u64,
    limit: usize,
}

impl MediaRegistry {
    fn new(limit: usize) -> Self {
        Self {
            entries: HashMap::new(),
            order: VecDeque::new(),
            next_stamp: 0,
            limit,
        }
    }

    fn get(&self, id: &MediaId) -> Option<&MediaRegistration> {
        self.entries.get(id).map(|(media, _)| media)
    }

    fn insert_all(&mut self, media: Vec<MediaRegistration>) {
        for media in media {
            let stamp = self.next_stamp;
            self.next_stamp += 1;
            self.order.push_back((media.id.clone(), stamp));
            self.entries.insert(media.id.clone(), (media, stamp));
        }
        while self.entries.len() > self.limit {
            let Some((id, stamp)) = self.order.pop_front() else {
                break;
            };
            if self.is_current(&id, stamp) {
                self.entries.remove(&id);
            }
        }
        if self.order.len() > self.limit.saturating_mul(4) {
            let entries = &self.entries;
            self.order
                .retain(|(id, stamp)| entries.get(id).is_some_and(|(_, current)| current == stamp));
        }
    }

    fn is_current(&self, id: &MediaId, stamp: u64) -> bool {
        self.entries
            .get(id)
            .is_some_and(|(_, current)| *current == stamp)
    }

    fn release(&mut self, scope: &MediaScope) {
        self.entries.retain(|id, _| !scope.owns(id));
        let entries = &self.entries;
        self.order.retain(|(id, _)| entries.contains_key(id));
    }

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

1	use std::{
2	borrow::Cow,
3	collections::{HashMap, VecDeque},
4	io::Read as _,
5	panic::AssertUnwindSafe,
6	sync::Arc,
7	time::Duration,
8	};
9
10	use futures_util::{
11	FutureExt, StreamExt,
12	future::{Shared, WeakShared},
13	};
14	use matrix_sdk::{
15	Client, Room, SendOutsideWasm, SyncOutsideWasm,
16	media::{MediaFormat, MediaRequestParameters, MediaThumbnailSettings},
17	room::reply::{EnforceThread, Reply as SdkReply},
18	ruma::{
19	api::{
20	Metadata, OutgoingRequest,
21	auth_scheme::SendAccessToken,
22	client::{authenticated_media, media},
23	},
24	events::room::{
25	MediaSource,
26	message::{AddMentions, ReplyWithinThread},
27	},
28	},
29	};
30	use matrix_sdk_crypto::{AttachmentDecryptor, MediaEncryptionInfo};
31	use ruma::MxcUri;
32	use tokio::sync::Mutex;
33
34	use super::executor::EngineBoxFuture;
35	use crate::{
36	AttachmentUpload, Error, ErrorKind, Media, MediaId, MediaKind, MediaSize, TimelineId,
37	normalise::{MediaRegistration, MediaScope},
38	upload::prepare_upload,
39	};
40
41	/// a stalled download gives up after this long without a byte
42	const DOWNLOAD_IDLE: Duration = Duration::from_secs(60);
43	/// the sdk client gives up on any request after thirty seconds, far too soon for a video
44	const DOWNLOAD_TOTAL: Duration = Duration::from_secs(60 * 60);
45
46	type FlightFuture = EngineBoxFuture<'static, Result<Media, Error>>;
47	type Flight = Shared<FlightFuture>;
48
49	/// owns everything needed to turn a media id or an mxc uri into bytes, so the
50	/// handle can answer media queries without a trip through the engine
51	pub(crate) struct MediaService {
52	client: Client,
53	registry: Mutex<MediaRegistry>,
54	flights: Mutex<HashMap<String, WeakShared<FlightFuture>>>,
55	max_bytes: usize,
56	}
57
58	impl MediaService {
59	pub(crate) fn new(client: Client, max_bytes: usize, max_sources: usize) -> Self {
60	Self {
61	client,
62	registry: Mutex::new(MediaRegistry::new(max_sources)),
63	flights: Mutex::new(HashMap::new()),
64	max_bytes,
65	}
66	}
67
68	pub(crate) async fn register(&self, media: Vec<MediaRegistration>) {
69	if !media.is_empty() {
70	self.registry.lock().await.insert_all(media);
71	}
72	}
73
74	pub(crate) async fn release(&self, scope: &MediaScope) {
75	self.registry.lock().await.release(scope);
76	}
77
78	pub(crate) async fn clear(&self) {
79	self.registry.lock().await.clear();
80	}
81
82	pub(crate) async fn timeline_media(
83	&self,
84	media_id: &MediaId,
85	size: MediaSize,
86	) -> Result<Media, Error> {
87	let media = self
88	.registry
89	.lock()
90	.await
91	.get(media_id)
92	.cloned()
93	.ok_or_else(\|\| {
94	Error::new(
95	ErrorKind::MediaUnknown,
96	"media is not registered in a retained timeline",
97	)
98	})?;
99	let (source, format, mimetype, declared_size) = match size {
100	MediaSize::Full => (
101	media.full.source,
102	MediaFormat::File,
103	media.full.mimetype,
104	media.full.declared_size,
105	),
106	MediaSize::Thumbnail { width, height } => match media.thumbnail {
107	Some(thumbnail) => (
108	thumbnail.source,
109	MediaFormat::File,
110	thumbnail.mimetype,
111	thumbnail.declared_size,
112	),
113	// the server can only resize what it can read
114	None if matches!(media.kind, MediaKind::Image \| MediaKind::Video)
115	&& matches!(media.full.source, MediaSource::Plain(_)) =>
116	{
117	// the server picks the thumbnail's type, so the original's would be a lie
118	let format = if is_local(&media.full.source) {
119	MediaFormat::File
120	} else {
121	thumbnail_format(width, height)
122	};
123	(media.full.source, format, None, None)
124	}
125	None => {
126	return Err(Error::new(
127	ErrorKind::MediaUnavailable,
128	"the requested media size is unavailable",
129	));
130	}
131	},
132	};
133	if declared_size.is_some_and(\|size\| size > self.max_bytes as u64) {
134	return Err(too_large("media exceeds the configured byte limit"));
135	}
136	self.fetch_or_file(MediaRequestParameters { source, format }, mimetype)
137	.await
138	}
139
140	pub(crate) async fn mxc_media(&self, uri: &MxcUri, size: MediaSize) -> Result<Media, Error> {
141	if !uri.is_valid() {
142	return Err(Error::new(ErrorKind::Invalid, "mxc uri is malformed"));
143	}
144	let format = match size {
145	MediaSize::Full => MediaFormat::File,
146	MediaSize::Thumbnail { width, height } => thumbnail_format(width, height),
147	};
148	self.fetch_or_file(
149	MediaRequestParameters {
150	source: MediaSource::Plain(uri.to_owned()),
151	format,
152	},
153	None,
154	)
155	.await
156	}
157
158	/// plenty of servers won't thumbnail everything, and the file itself is the next best
159	/// thing. the byte cap still holds either way
160	async fn fetch_or_file(
161	&self,
162	request: MediaRequestParameters,
163	mimetype: Option<String>,
164	) -> Result<Media, Error> {
165	if !matches!(request.format, MediaFormat::Thumbnail(_)) {
166	return self.fetch(request, mimetype).await;
167	}
168	match self.fetch(request.clone(), mimetype.clone()).await {
169	Err(error) if error.kind == ErrorKind::MediaUnavailable => {
170	let file = MediaRequestParameters {
171	source: request.source,
172	format: MediaFormat::File,
173	};
174	self.fetch(file, mimetype).await
175	}
176	result => result,
177	}
178	}
179
180	async fn fetch(
181	&self,
182	request: MediaRequestParameters,
183	mimetype: Option<String>,
184	) -> Result<Media, Error> {
185	let key = flight_key(&request);
186	let flight = {
187	let mut flights = self.flights.lock().await;
188	flights.retain(\|_, flight\| flight.upgrade().is_some());
189	match flights.get(&key).and_then(WeakShared::upgrade) {
190	Some(flight) => flight,
191	None => {
192	let prepared = AssertUnwindSafe(fetch(
193	self.client.clone(),
194	request,
195	mimetype,
196	self.max_bytes,
197	))
198	.catch_unwind()
199	.map(\|result\| {
200	result.unwrap_or_else(\|_\| {
201	Err(Error::new(ErrorKind::Other, "media fetch panicked"))
202	})
203	});
204	let boxed: FlightFuture = Box::pin(prepared);
205	let flight: Flight = boxed.shared();
206	if let Some(weak) = flight.downgrade() {
207	flights.insert(key, weak);
208	}
209	flight
210	}
211	}
212	};
213	flight.await
214	}
215	}
216
217	fn thumbnail_format(width: u32, height: u32) -> MediaFormat {
218	MediaFormat::Thumbnail(MediaThumbnailSettings::new(width.into(), height.into()))
219	}
220
221	fn is_local(source: &MediaSource) -> bool {
222	let uri = match source {
223	MediaSource::Plain(uri) => uri,
224	MediaSource::Encrypted(file) => &file.url,
225	};
226	uri.server_name()
227	.is_ok_and(\|server\| server.as_str() == "send-queue.localhost")
228	}
229
230	/// plain and encrypted sources share a cache key in the sdk, but must never share a flight
231	fn flight_key(request: &MediaRequestParameters) -> String {
232	let uri = match &request.source {
233	MediaSource::Plain(uri) => format!("plain\|{uri}"),
234	MediaSource::Encrypted(file) => format!("encrypted\|{}", file.url),
235	};
236	match &request.format {
237	MediaFormat::File => uri.to_owned(),
238	MediaFormat::Thumbnail(settings) => format!(
239	"{uri}\|{}\|{}x{}\|{}",
240	settings.method, settings.width, settings.height, settings.animated
241	),
242	}
243	}
244
245	fn too_large(message: &'static str) -> Error {
246	Error::new(ErrorKind::MediaTooLarge, message)
247	}
248
249	async fn fetch(
250	client: Client,
251	request: MediaRequestParameters,
252	mimetype: Option<String>,
253	max_bytes: usize,
254	) -> Result<Media, Error> {
255	if let Some(cached) = cached_media(&client, &request).await? {
256	if cached.len() > max_bytes {
257	if let Ok(store) = client.media_store().lock().await {
258	let _ = store.remove_media_content(&request).await;
259	}
260	return Err(too_large("cached media exceeds the configured byte limit"));
261	}
262	// plain and encrypted share a cache key, so a plain fetch of an encrypted file's uri
263	// leaves ciphertext here. its hash gives it away
264	if is_ciphertext(&cached, &request.source) {
265	let bytes = maybe_decrypt(cached, &request, max_bytes)?;
266	cache(&client, &request, &bytes).await;
267	return Ok(media(bytes, mimetype));
268	}
269	return Ok(media(cached, mimetype));
270	}
271	let (bytes, served_mimetype) = download_bounded(&client, &request, max_bytes).await?;
272	// the sdk caches plaintext under the encrypted key, so we do too or nobody can read anything
273	let bytes = maybe_decrypt(bytes, &request, max_bytes)?;
274	cache(&client, &request, &bytes).await;
275	Ok(media(bytes, mimetype.or(served_mimetype)))
276	}
277
278	async fn cache(client: &Client, request: &MediaRequestParameters, bytes: &[u8]) {
279	if let Ok(store) = client.media_store().lock().await {
280	let _ = store
281	.add_media_content(
282	request,
283	bytes.to_vec(),
284	matrix_sdk::media::store::IgnoreMediaRetentionPolicy::No,
285	)
286	.await;
287	}
288	}
289
290	fn is_ciphertext(bytes: &[u8], source: &MediaSource) -> bool {
291	use sha2::{Digest, Sha256};
292	let MediaSource::Encrypted(file) = source else {
293	return false;
294	};
295	file.hashes
296	.get(&ruma::events::room::EncryptedFileHashAlgorithm::Sha256)
297	.is_some_and(\|hash\| hash.as_bytes() == Sha256::digest(bytes).as_slice())
298	}
299
300	fn media(bytes: Vec<u8>, mimetype: Option<String>) -> Media {
301	let mimetype = mimetype
302	.or_else(\|\| sniff_image(&bytes).map(str::to_owned))
303	.unwrap_or_else(\|\| "application/octet-stream".to_owned());
304	Media {
305	bytes: Arc::from(bytes),
306	mimetype,
307	}
308	}
309
310	/// the cache forgets content types, and avatars never declare one. images are the common case
311	fn sniff_image(bytes: &[u8]) -> Option<&'static str> {
312	match bytes {
313	[0x89, b'P', b'N', b'G', ..] => Some("image/png"),
314	[0xff, 0xd8, 0xff, ..] => Some("image/jpeg"),
315	[b'G', b'I', b'F', b'8', ..] => Some("image/gif"),
316	[
317	b'R',
318	b'I',
319	b'F',
320	b'F',
321	_,
322	_,
323	_,
324	_,
325	b'W',
326	b'E',
327	b'B',
328	b'P',
329	..,
330	] => Some("image/webp"),
331	_ => None,
332	}
333	}
334
335	async fn cached_media(
336	client: &Client,
337	request: &MediaRequestParameters,
338	) -> Result<Option<Vec<u8>>, Error> {
339	fn cache_failed(
340	source: impl std::error::Error + SendOutsideWasm + SyncOutsideWasm + 'static,
341	) -> Error {
342	Error::with_source(ErrorKind::Store, "reading media cache failed", source)
343	}
344	client
345	.media_store()
346	.lock()
347	.await
348	.map_err(cache_failed)?
349	.get_media_content(request)
350	.await
351	.map_err(cache_failed)
352	}
353
354	fn maybe_decrypt(
355	bytes: Vec<u8>,
356	request: &MediaRequestParameters,
357	max_bytes: usize,
358	) -> Result<Vec<u8>, Error> {
359	let MediaSource::Encrypted(file) = &request.source else {
360	return Ok(bytes);
361	};
362	fn decrypt_failed(
363	source: impl std::error::Error + SendOutsideWasm + SyncOutsideWasm + 'static,
364	) -> Error {
365	Error::with_source(
366	ErrorKind::MediaUnavailable,
367	"media decryption failed",
368	source,
369	)
370	}
371	let info: MediaEncryptionInfo = (**file).clone().into();
372	let mut cursor = std::io::Cursor::new(bytes);
373	let mut decryptor = AttachmentDecryptor::new(&mut cursor, info).map_err(decrypt_failed)?;
374	let mut decrypted = Vec::new();
375	let mut chunk = [0_u8; 8192];
376	loop {
377	let read = decryptor.read(&mut chunk).map_err(decrypt_failed)?;
378	if read == 0 {
379	return Ok(decrypted);
380	}
381	if read > max_bytes.saturating_sub(decrypted.len()) {
382	return Err(too_large(
383	"decrypted media exceeds the configured byte limit",
384	));
385	}
386	decrypted.extend_from_slice(&chunk[..read]);
387	}
388	}
389
390	/// the sdk's own download has no size cap, so we do it by hand
391	async fn download_bounded(
392	client: &Client,
393	request: &MediaRequestParameters,
394	max_bytes: usize,
395	) -> Result<(Vec<u8>, Option<String>), Error> {
396	if is_local(&request.source) {
397	return Err(Error::new(
398	ErrorKind::MediaUnavailable,
399	"local media is not in the media cache",
400	));
401	}
402	let network_failed =
403	\|source\| Error::with_source(ErrorKind::Network, "media download failed", source);
404	let mut refreshed = false;
405	let response = loop {
406	let http_request = media_http_request(client, request).await?;
407	let (parts, _) = http_request.into_parts();
408	let mut builder = client
409	.http_client()
410	.get(parts.uri.to_string())
411	.timeout(DOWNLOAD_TOTAL);
412	for (name, value) in &parts.headers {
413	builder = builder.header(name.as_str(), value.as_bytes());
414	}
415	let response = builder.send().await.map_err(network_failed)?;
416	if response.status() != matrix_sdk::reqwest::StatusCode::UNAUTHORIZED {
417	break response;
418	}
419	if refreshed
420	\|\| client
421	.session_tokens()
422	.is_none_or(\|tokens\| tokens.refresh_token.is_none())
423	{
424	return Err(Error::new(
425	ErrorKind::Unauthenticated,
426	"media download was not authorized",
427	));
428	}
429	client.refresh_access_token().await.map_err(\|source\| {
430	Error::from_refresh_token("refreshing media access token failed", source)
431	})?;
432	refreshed = true;
433	};
434	let status = response.status();
435	if !status.is_success() {
436	return Err(Error::new(
437	ErrorKind::MediaUnavailable,
438	format!("media download failed with status {status}"),
439	));
440	}
441	if response
442	.content_length()
443	.is_some_and(\|length\| length > max_bytes as u64)
444	{
445	return Err(too_large("media exceeds the configured byte limit"));
446	}
447	let mimetype = response
448	.headers()
449	.get(http::header::CONTENT_TYPE)
450	.and_then(\|value\| value.to_str().ok())
451	.map(str::to_owned);
452	let mut stream = response.bytes_stream();
453	let mut bytes = Vec::new();
454	loop {
455	let chunk = matrix_sdk::timeout::timeout(stream.next(), DOWNLOAD_IDLE)
456	.await
457	.map_err(\|_\| Error::new(ErrorKind::Network, "media download stalled"))?;
458	let Some(chunk) = chunk else { break };
459	let chunk = chunk.map_err(network_failed)?;
460	if chunk.len() > max_bytes.saturating_sub(bytes.len()) {
461	return Err(too_large(
462	"downloaded media exceeds the configured byte limit",
463	));
464	}
465	bytes.extend_from_slice(&chunk);
466	}
467	Ok((bytes, mimetype))
468	}
469
470	async fn media_http_request(
471	client: &Client,
472	request: &MediaRequestParameters,
473	) -> Result<http::Request<Vec<u8>>, Error> {
474	let versions = client.supported_versions().await.map_err(\|source\| {
475	Error::with_source(
476	ErrorKind::Network,
477	"reading supported server versions failed",
478	source,
479	)
480	})?;
481	let authenticated =
482	authenticated_media::get_content::v1::Request::PATH_BUILDER.is_supported(&versions);
483	let base = client.homeserver().to_string();
484	let token = client.access_token();
485	let invalid_uri =
486	\|source\| Error::with_source(ErrorKind::Invalid, "media URI is invalid", source);
487	let (uri, thumbnail) = match (&request.source, &request.format) {
488	(MediaSource::Encrypted(file), _) => (&file.url, None),
489	(MediaSource::Plain(uri), MediaFormat::File) => (uri, None),
490	(MediaSource::Plain(uri), MediaFormat::Thumbnail(settings)) => (uri, Some(settings)),
491	};
492	// each endpoint has its own auth input type, so no generic helper can take them all
493	macro_rules! into_http {
494	($endpoint:expr, $token:expr) => {
495	$endpoint
496	.try_into_http_request::<Vec<u8>>(&base, $token, Cow::Borrowed(&versions))
497	.map_err(\|source\| {
498	Error::with_source(ErrorKind::Other, "media request is invalid", source)
499	})?
500	};
501	}
502	let authed = SendAccessToken::IfRequired(token.as_deref().unwrap_or(""));
503	#[allow(deprecated)]
504	let mut http_request = match (thumbnail, authenticated) {
505	(None, true) => into_http!(
506	authenticated_media::get_content::v1::Request::from_uri(uri).map_err(invalid_uri)?,
507	authed
508	),
509	(None, false) => into_http!(
510	media::get_content::v3::Request::from_url(uri).map_err(invalid_uri)?,
511	SendAccessToken::None
512	),
513	(Some(settings), true) => {
514	let mut endpoint = authenticated_media::get_content_thumbnail::v1::Request::from_uri(
515	uri,
516	settings.width,
517	settings.height,
518	)
519	.map_err(invalid_uri)?;
520	endpoint.method = Some(settings.method.clone());
521	endpoint.animated = Some(settings.animated);
522	into_http!(endpoint, authed)
523	}
524	(Some(settings), false) => {
525	let mut endpoint = media::get_content_thumbnail::v3::Request::from_url(
526	uri,
527	settings.width,
528	settings.height,
529	)
530	.map_err(invalid_uri)?;
531	endpoint.method = Some(settings.method.clone());
532	endpoint.animated = Some(settings.animated);
533	into_http!(endpoint, SendAccessToken::None)
534	}
535	};
536	// legacy endpoints don't ask for a token but plenty of servers want one anyway
537	if let Some(token) = token.as_deref() {
538	let value = format!("Bearer {token}").parse().map_err(\|_\| {
539	Error::new(ErrorKind::Unauthenticated, "stored access token is invalid")
540	})?;
541	http_request
542	.headers_mut()
543	.insert(http::header::AUTHORIZATION, value);
544	}
545	Ok(http_request)
546	}
547
548	pub(crate) async fn queue_upload(
549	room: Room,
550	upload: AttachmentUpload,
551	max_bytes: usize,
552	) -> Result<(), Error> {
553	let (upload, mime, mut config) = prepare_upload(upload, max_bytes)?;
554	if let TimelineId::Thread { root_event_id, .. } = &upload.timeline_id {
555	config.reply = Some(SdkReply {
556	event_id: root_event_id.clone(),
557	enforce_thread: EnforceThread::Threaded(ReplyWithinThread::No),
558	add_mentions: AddMentions::No,
559	});
560	}
561	room.send_queue()
562	.send_attachment(upload.filename, mime, upload.bytes, config)
563	.await
564	.map(drop)
565	.map_err(\|source\| {
566	use matrix_sdk::send_queue::RoomSendQueueError;
567	let kind = match &source {
568	RoomSendQueueError::RoomNotJoined \| RoomSendQueueError::RoomDisappeared => {
569	ErrorKind::RoomUnknown
570	}
571	RoomSendQueueError::StorageError(_) => ErrorKind::Store,
572	_ => ErrorKind::MediaUnavailable,
573	};
574	Error::with_source(kind, "media upload could not be queued", source)
575	})
576	}
577
578	/// least recently seen goes first. stale order slots are skipped by stamp rather than hunted down
579	struct MediaRegistry {
580	entries: HashMap<MediaId, (MediaRegistration, u64)>,
581	order: VecDeque<(MediaId, u64)>,
582	next_stamp: u64,
583	limit: usize,
584	}
585
586	impl MediaRegistry {
587	fn new(limit: usize) -> Self {
588	Self {
589	entries: HashMap::new(),
590	order: VecDeque::new(),
591	next_stamp: 0,
592	limit,
593	}
594	}
595
596	fn get(&self, id: &MediaId) -> Option<&MediaRegistration> {
597	self.entries.get(id).map(\|(media, _)\| media)
598	}
599
600	fn insert_all(&mut self, media: Vec<MediaRegistration>) {
601	for media in media {
602	let stamp = self.next_stamp;
603	self.next_stamp += 1;
604	self.order.push_back((media.id.clone(), stamp));
605	self.entries.insert(media.id.clone(), (media, stamp));
606	}
607	while self.entries.len() > self.limit {
608	let Some((id, stamp)) = self.order.pop_front() else {
609	break;
610	};
611	if self.is_current(&id, stamp) {
612	self.entries.remove(&id);
613	}
614	}
615	if self.order.len() > self.limit.saturating_mul(4) {
616	let entries = &self.entries;
617	self.order
618	.retain(\|(id, stamp)\| entries.get(id).is_some_and(\|(_, current)\| current == stamp));
619	}
620	}
621
622	fn is_current(&self, id: &MediaId, stamp: u64) -> bool {
623	self.entries
624	.get(id)
625	.is_some_and(\|(_, current)\| *current == stamp)
626	}
627
628	fn release(&mut self, scope: &MediaScope) {
629	self.entries.retain(\|id, _\| !scope.owns(id));
630	let entries = &self.entries;
631	self.order.retain(\|(id, _)\| entries.contains_key(id));
632	}
633
634	fn clear(&mut self) {
635	self.entries.clear();
636	self.order.clear();
637	}
638	}