fix(media): Validate preview image downloads

Cargo.lock

@@ -1137,6 +1137,7 @@ dependencies = [
  "askama",
  "assign",
  "async-trait",
+ "axum",
  "base64 0.23.1",
  "bytes",
  "conduwuit_core",
@@ -1167,6 +1168,7 @@ dependencies = [
  "resolvematrix",
  "ruma",
  "ruminuwuity",
+ "rustls",
  "rustyline-async",
  "sd-notify",
  "serde",
@@ -1174,6 +1176,7 @@ dependencies = [
  "serde_json",
  "serde_urlencoded",
  "sha2 0.11.0",
+ "tempfile",
  "termimad",
  "tokio",
  "tracing",
@@ -3791,7 +3794,7 @@ version = "5.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d"
 dependencies = [
- "base64 0.21.7",
+ "base64 0.22.1",
  "chrono",
  "getrandom 0.2.17",
  "http",
@@ -5582,7 +5585,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b8050abb251097357e24aff63ba2c52a6309ecb7d23a5474023df960a02694d8"
 dependencies = [
  "annotate-snippets",
- "base64 0.21.7",
+ "base64 0.23.1",
  "encoding_rs_io",
  "granit-parser",
  "nohash-hasher",
@@ -6082,6 +6085,19 @@ version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
 
+[[package]]
+name = "tempfile"
+version = "3.27.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
+dependencies = [
+ "fastrand",
+ "getrandom 0.4.3",
+ "once_cell",
+ "rustix",
+ "windows-sys 0.61.2",
+]
+
 [[package]]
 name = "tendril"
 version = "0.4.3"

changelog.d/2308.bugfix

@@ -0,0 +1 @@
+Fixed missing MIME types on URL preview images and rejected failed or invalid image downloads. Contributed by @eleboucher.

src/service/Cargo.toml

@@ -124,6 +124,11 @@ serde_urlencoded.workspace = true
 openidconnect.workspace = true
 rand.workspace = true
 
+[dev-dependencies]
+axum.workspace = true
+rustls = { workspace = true, features = ["ring"] }
+tempfile = "3"
+
 [target.'cfg(all(unix, target_os = "linux"))'.dependencies]
 sd-notify.workspace = true
 sd-notify.optional = true

src/service/media/mod.rs

@@ -36,6 +36,8 @@ pub struct FileMeta {
 
 pub struct Service {
 	url_preview_mutex: MutexMap<String, ()>,
+	#[cfg(feature = "url_preview")]
+	url_preview_image_semaphore: Arc<tokio::sync::Semaphore>,
 	pub(super) db: Data,
 	services: Services,
 }
@@ -62,6 +64,8 @@ impl crate::Service for Service {
 	fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {
 		Ok(Arc::new(Self {
 			url_preview_mutex: MutexMap::new(),
+			#[cfg(feature = "url_preview")]
+			url_preview_image_semaphore: Arc::new(tokio::sync::Semaphore::new(2)),
 			db: Data::new(args.db),
 			services: Services {
 				server: args.server.clone(),

src/service/media/preview.rs

@@ -139,13 +139,18 @@ impl Service {
 
 		let mut preview_data = preview_data.unwrap_or_default();
 
-		let image = self
-			.services
-			.client
-			.url_preview
-			.get(url)
-			.send()
-			.await?
+		let response = self.services.client.url_preview.get(url).send().await?;
+		if !response.status().is_success() {
+			return Err!(Request(Unknown("Image download returned HTTP {}", response.status())));
+		}
+
+		let permit = self
+			.url_preview_image_semaphore
+			.clone()
+			.acquire_owned()
+			.await
+			.map_err(|error| err!(Request(Unknown("Image decoder unavailable: {error}"))))?;
+		let image = response
 			.limit_read(
 				self.services
 					.server
@@ -156,27 +161,45 @@ impl Service {
 			)
 			.await?;
 
+		let (image, content_type, width, height) = self
+			.services
+			.server
+			.runtime()
+			.spawn_blocking(move || -> Result<_> {
+				// Keep the slot until decoding finishes, even if the request is
+				// cancelled.
+				let _permit = permit;
+				let mut reader =
+					ImageReader::new(std::io::Cursor::new(&image)).with_guessed_format()?;
+				let mut limits = image::Limits::default();
+				limits.max_image_width = Some(8192);
+				limits.max_image_height = Some(8192);
+				limits.max_alloc = Some(64 * 1024 * 1024);
+				reader.limits(limits);
+				let content_type = reader
+					.format()
+					.ok_or_else(|| {
+						err!(Request(Unknown("Downloaded file is not a supported image")))
+					})?
+					.to_mime_type();
+				let decoded = reader.decode().map_err(|error| {
+					err!(Request(Unknown("Failed to decode preview image: {error}")))
+				})?;
+				Ok((image, content_type, decoded.width(), decoded.height()))
+			})
+			.await??;
+
 		let mxc = Mxc {
 			server_name: self.services.globals.server_name(),
 			media_id: &random_string(super::MXC_LENGTH),
 		};
 
-		self.create(&mxc, None, None, None, &image).await?;
+		self.create(&mxc, None, None, Some(content_type), &image)
+			.await?;
 
 		preview_data.image = Some(mxc.to_string());
-		if preview_data.image_height.is_none() || preview_data.image_width.is_none() {
-			let cursor = std::io::Cursor::new(&image);
-			let (width, height) = match ImageReader::new(cursor).with_guessed_format() {
-				| Err(_) => (None, None),
-				| Ok(reader) => match reader.into_dimensions() {
-					| Err(_) => (None, None),
-					| Ok((width, height)) => (Some(width), Some(height)),
-				},
-			};
-
-			preview_data.image_width = width;
-			preview_data.image_height = height;
-		}
+		preview_data.image_width = Some(width);
+		preview_data.image_height = Some(height);
 
 		Ok(preview_data)
 	}

src/service/media/tests.rs

@@ -1,5 +1,181 @@
 #![cfg(test)]
 
+#[cfg(feature = "url_preview")]
+mod url_preview {
+	use std::{io::Cursor, sync::Arc};
+
+	use axum::{Router, body::Body, http::StatusCode, routing::get};
+	use conduwuit::{
+		Config, Server,
+		log::{Log, LogLevelReloadHandles},
+	};
+	use image::{DynamicImage, ImageFormat};
+	use serde_json::json;
+	use tokio::net::TcpListener;
+
+	use crate::{
+		Services,
+		media::{mxc::Mxc, preview::UrlPreviewData},
+	};
+
+	async fn services() -> (tempfile::TempDir, Arc<Services>) {
+		_ = rustls::crypto::ring::default_provider().install_default();
+		let directory = tempfile::tempdir().unwrap();
+		let config: Config = serde_json::from_value(json!({
+			"server_name": "example.com",
+			"database_path": directory.path(),
+			"db_pool_affinity": false,
+			"db_pool_workers": 1,
+			"rocksdb_parallelism_threads": 1,
+			"db_cache_capacity_mb": 8,
+			"db_write_buffer_capacity_mb": 8,
+			"max_request_size": 4096,
+		}))
+		.unwrap();
+		let server =
+			Arc::new(Server::new(config, Some(tokio::runtime::Handle::current()), Log {
+				reload: LogLevelReloadHandles::default(),
+				capture: Arc::default(),
+			}));
+		let services = Services::build(server).await.unwrap();
+		services.media.create_media_dir().await.unwrap();
+		(directory, services)
+	}
+
+	fn image(format: ImageFormat) -> Vec<u8> {
+		let mut bytes = Cursor::new(Vec::new());
+		DynamicImage::new_rgb8(2, 3)
+			.write_to(&mut bytes, format)
+			.unwrap();
+		bytes.into_inner()
+	}
+
+	struct ImageServer {
+		url: String,
+		task: tokio::task::JoinHandle<()>,
+	}
+
+	impl ImageServer {
+		async fn new(
+			services: &Services,
+			status: StatusCode,
+			content_type: Option<&str>,
+			body: Vec<u8>,
+		) -> Self {
+			let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+			let url = format!("http://{}/image", listener.local_addr().unwrap());
+			let mut response = http::Response::builder().status(status);
+			if let Some(content_type) = content_type {
+				response = response.header(http::header::CONTENT_TYPE, content_type);
+			}
+			let response = response.body(body).unwrap();
+			let app = Router::new().route(
+				"/image",
+				get(move || {
+					let response = response.clone();
+					async move { response.map(Body::from) }
+				}),
+			);
+			let task = services
+				.server
+				.runtime()
+				.spawn(async move { axum::serve(listener, app).await.unwrap() });
+			Self { url, task }
+		}
+	}
+
+	impl Drop for ImageServer {
+		fn drop(&mut self) { self.task.abort(); }
+	}
+
+	#[tokio::test]
+	async fn preview_images_store_mime_type_and_dimensions() {
+		let (_directory, services) = services().await;
+		for format in [ImageFormat::Png, ImageFormat::Jpeg, ImageFormat::Gif, ImageFormat::WebP] {
+			for content_type in [
+				Some(format.to_mime_type()),
+				None,
+				Some("application/octet-stream"),
+				Some("image/svg+xml"),
+			] {
+				let bytes = image(format);
+				let server =
+					ImageServer::new(&services, StatusCode::OK, content_type, bytes.clone())
+						.await;
+				let preview = services
+					.media
+					.download_image(
+						&server.url,
+						Some(UrlPreviewData {
+							title: Some("Example preview".into()),
+							image_width: Some(99),
+							image_height: Some(99),
+							..UrlPreviewData::default()
+						}),
+					)
+					.await
+					.unwrap();
+				assert_eq!((preview.image_width, preview.image_height), (Some(2), Some(3)));
+				assert_eq!(preview.title.as_deref(), Some("Example preview"));
+				let uri = preview.image.unwrap();
+				let mxc = Mxc::try_from(uri.as_str()).unwrap();
+				let stored = services.media.get(&mxc).await.unwrap().unwrap();
+				assert_eq!(stored.content_type.as_deref(), Some(format.to_mime_type()));
+				assert_eq!(stored.content.unwrap(), bytes);
+			}
+		}
+	}
+
+	#[tokio::test]
+	async fn preview_images_reject_failed_or_invalid_downloads_before_storage() {
+		let (_directory, services) = services().await;
+		let png = image(ImageFormat::Png);
+		let cases = [
+			(StatusCode::FOUND, png.clone()),
+			(StatusCode::FORBIDDEN, png.clone()),
+			(StatusCode::NOT_FOUND, b"<html>Not found</html>".to_vec()),
+			(StatusCode::INTERNAL_SERVER_ERROR, png.clone()),
+			(StatusCode::OK, b"<html>Access denied</html>".to_vec()),
+			(StatusCode::OK, Vec::new()),
+			(StatusCode::OK, png[..33].to_vec()),
+			(StatusCode::OK, png[..png.len() - 20].to_vec()),
+			(StatusCode::OK, vec![0; 4097]),
+		];
+		for (status, bytes) in cases {
+			let server = ImageServer::new(&services, status, Some("image/png"), bytes).await;
+			assert!(
+				services
+					.media
+					.download_image(&server.url, None)
+					.await
+					.is_err(),
+				"accepted failed or invalid image with status {status}",
+			);
+			assert!(services.media.get_all_mxcs().await.unwrap().is_empty());
+		}
+	}
+
+	#[tokio::test]
+	async fn preview_images_reject_excessive_decoded_dimensions_and_memory() {
+		let (_directory, services) = services().await;
+		for (width, height) in [(8193_u16, 1_u16), (5000, 5000)] {
+			let mut bytes = image(ImageFormat::Gif);
+			bytes[6..8].copy_from_slice(&width.to_le_bytes());
+			bytes[8..10].copy_from_slice(&height.to_le_bytes());
+			let server =
+				ImageServer::new(&services, StatusCode::OK, Some("image/gif"), bytes).await;
+			assert!(
+				services
+					.media
+					.download_image(&server.url, None)
+					.await
+					.is_err()
+			);
+			assert!(services.media.get_all_mxcs().await.unwrap().is_empty());
+		}
+	}
+}
+
 #[tokio::test]
 #[cfg(disable)] //TODO: fixme
 async fn long_file_names_works() {