macltok

public

Star0

Watch0

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