split thumbnailing related into unit

Signed-off-by: Jason Volk <jason@zemos.net>

src/service/media/data.rs

@@ -12,6 +12,13 @@ pub(crate) struct Data {
 	url_previews: Arc<Map>,
 }
 
+#[derive(Debug)]
+pub(super) struct Metadata {
+	pub(super) content_disposition: Option<String>,
+	pub(super) content_type: Option<String>,
+	pub(super) key: Vec<u8>,
+}
+
 impl Data {
 	pub(super) fn new(db: &Arc<Database>) -> Self {
 		Self {
@@ -104,9 +111,7 @@ impl Data {
 		Ok(keys)
 	}
 
-	pub(super) fn search_file_metadata(
+	pub(super) fn search_file_metadata(&self, mxc: &str, width: u32, height: u32) -> Result<Metadata> {
-		&self, mxc: &str, width: u32, height: u32,
-	) -> Result<(Option<String>, Option<String>, Vec<u8>)> {
 		let mut prefix = mxc.as_bytes().to_vec();
 		prefix.push(0xFF);
 		prefix.extend_from_slice(&width.to_be_bytes());
@@ -141,7 +146,12 @@ impl Data {
 					.map_err(|_| Error::bad_database("Content Disposition in mediaid_file is invalid unicode."))?,
 			)
 		};
-		Ok((content_disposition, content_type, key))
+
+		Ok(Metadata {
+			content_disposition,
+			content_type,
+			key,
+		})
 	}
 
 	/// Gets all the media keys in our database (this includes all the metadata

src/service/media/mod.rs

@@ -1,13 +1,13 @@
 mod data;
 mod tests;
+mod thumbnail;
 
-use std::{collections::HashMap, io::Cursor, num::Saturating as Sat, path::PathBuf, sync::Arc, time::SystemTime};
+use std::{collections::HashMap, path::PathBuf, sync::Arc, time::SystemTime};
 
 use async_trait::async_trait;
 use base64::{engine::general_purpose, Engine as _};
-use conduit::{checked, debug, debug_error, error, utils, Error, Result, Server};
+use conduit::{debug, debug_error, err, error, utils, Result, Server};
-use data::Data;
+use data::{Data, Metadata};
-use image::imageops::FilterType;
 use ruma::{OwnedMxcUri, OwnedUserId};
 use serde::Serialize;
 use tokio::{
@@ -107,30 +107,14 @@ impl Service {
 		}
 	}
 
-	/// Uploads or replaces a file thumbnail.
-	#[allow(clippy::too_many_arguments)]
-	pub async fn upload_thumbnail(
-		&self, sender_user: Option<OwnedUserId>, mxc: &str, content_disposition: Option<&str>,
-		content_type: Option<&str>, width: u32, height: u32, file: &[u8],
-	) -> Result<()> {
-		let key = if let Some(user) = sender_user {
-			self.db
-				.create_file_metadata(Some(user.as_str()), mxc, width, height, content_disposition, content_type)?
-		} else {
-			self.db
-				.create_file_metadata(None, mxc, width, height, content_disposition, content_type)?
-		};
-
-		//TODO: Dangling metadata in database if creation fails
-		let mut f = self.create_media_file(&key).await?;
-		f.write_all(file).await?;
-
-		Ok(())
-	}
-
 	/// Downloads a file.
 	pub async fn get(&self, mxc: &str) -> Result<Option<FileMeta>> {
-		if let Ok((content_disposition, content_type, key)) = self.db.search_file_metadata(mxc, 0, 0) {
+		if let Ok(Metadata {
+			content_disposition,
+			content_type,
+			key,
+		}) = self.db.search_file_metadata(mxc, 0, 0)
+		{
 			let mut content = Vec::new();
 			let path = self.get_media_file(&key);
 			BufReader::new(fs::File::open(path).await?)
@@ -249,129 +233,6 @@ impl Service {
 		Ok(deletion_count)
 	}
 
-	/// Returns width, height of the thumbnail and whether it should be cropped.
-	/// Returns None when the server should send the original file.
-	pub fn thumbnail_properties(&self, width: u32, height: u32) -> Option<(u32, u32, bool)> {
-		match (width, height) {
-			(0..=32, 0..=32) => Some((32, 32, true)),
-			(0..=96, 0..=96) => Some((96, 96, true)),
-			(0..=320, 0..=240) => Some((320, 240, false)),
-			(0..=640, 0..=480) => Some((640, 480, false)),
-			(0..=800, 0..=600) => Some((800, 600, false)),
-			_ => None,
-		}
-	}
-
-	/// Downloads a file's thumbnail.
-	///
-	/// Here's an example on how it works:
-	///
-	/// - Client requests an image with width=567, height=567
-	/// - Server rounds that up to (800, 600), so it doesn't have to save too
-	///   many thumbnails
-	/// - Server rounds that up again to (958, 600) to fix the aspect ratio
-	///   (only for width,height>96)
-	/// - Server creates the thumbnail and sends it to the user
-	///
-	/// For width,height <= 96 the server uses another thumbnailing algorithm
-	/// which crops the image afterwards.
-	pub async fn get_thumbnail(&self, mxc: &str, width: u32, height: u32) -> Result<Option<FileMeta>> {
-		let (width, height, crop) = self
-			.thumbnail_properties(width, height)
-			.unwrap_or((0, 0, false)); // 0, 0 because that's the original file
-
-		if let Ok((content_disposition, content_type, key)) = self.db.search_file_metadata(mxc, width, height) {
-			// Using saved thumbnail
-			let mut content = Vec::new();
-			let path = self.get_media_file(&key);
-			fs::File::open(path)
-				.await?
-				.read_to_end(&mut content)
-				.await?;
-
-			Ok(Some(FileMeta {
-				content: Some(content),
-				content_type,
-				content_disposition,
-			}))
-		} else if let Ok((content_disposition, content_type, key)) = self.db.search_file_metadata(mxc, 0, 0) {
-			// Generate a thumbnail
-			let mut content = Vec::new();
-			let path = self.get_media_file(&key);
-			fs::File::open(path)
-				.await?
-				.read_to_end(&mut content)
-				.await?;
-
-			if let Ok(image) = image::load_from_memory(&content) {
-				let original_width = image.width();
-				let original_height = image.height();
-				if width > original_width || height > original_height {
-					return Ok(Some(FileMeta {
-						content: Some(content),
-						content_type,
-						content_disposition,
-					}));
-				}
-
-				let thumbnail = if crop {
-					image.resize_to_fill(width, height, FilterType::CatmullRom)
-				} else {
-					let (exact_width, exact_height) = {
-						let ratio = Sat(original_width) * Sat(height);
-						let nratio = Sat(width) * Sat(original_height);
-
-						let use_width = nratio <= ratio;
-						let intermediate = if use_width {
-							Sat(original_height) * Sat(checked!(width / original_width)?)
-						} else {
-							Sat(original_width) * Sat(checked!(height / original_height)?)
-						};
-
-						if use_width {
-							(width, intermediate.0)
-						} else {
-							(intermediate.0, height)
-						}
-					};
-
-					image.thumbnail_exact(exact_width, exact_height)
-				};
-
-				let mut thumbnail_bytes = Vec::new();
-				thumbnail.write_to(&mut Cursor::new(&mut thumbnail_bytes), image::ImageFormat::Png)?;
-
-				// Save thumbnail in database so we don't have to generate it again next time
-				let thumbnail_key = self.db.create_file_metadata(
-					None,
-					mxc,
-					width,
-					height,
-					content_disposition.as_deref(),
-					content_type.as_deref(),
-				)?;
-
-				let mut f = self.create_media_file(&thumbnail_key).await?;
-				f.write_all(&thumbnail_bytes).await?;
-
-				Ok(Some(FileMeta {
-					content: Some(thumbnail_bytes),
-					content_type,
-					content_disposition,
-				}))
-			} else {
-				// Couldn't parse file to generate thumbnail, send original
-				Ok(Some(FileMeta {
-					content: Some(content),
-					content_type,
-					content_disposition,
-				}))
-			}
-		} else {
-			Ok(None)
-		}
-	}
-
 	pub async fn get_url_preview(&self, url: &str) -> Option<UrlPreviewData> { self.db.get_url_preview(url) }
 
 	/// TODO: use this?

src/service/media/thumbnail.rs

@@ -0,0 +1,171 @@
+use std::{io::Cursor, num::Saturating as Sat};
+
+use conduit::{checked, Result};
+use image::{imageops::FilterType, DynamicImage};
+use ruma::OwnedUserId;
+use tokio::{
+	fs,
+	io::{AsyncReadExt, AsyncWriteExt},
+};
+
+use super::{data::Metadata, FileMeta};
+
+impl super::Service {
+	/// Uploads or replaces a file thumbnail.
+	#[allow(clippy::too_many_arguments)]
+	pub async fn upload_thumbnail(
+		&self, sender_user: Option<OwnedUserId>, mxc: &str, content_disposition: Option<&str>,
+		content_type: Option<&str>, width: u32, height: u32, file: &[u8],
+	) -> Result<()> {
+		let key = if let Some(user) = sender_user {
+			self.db
+				.create_file_metadata(Some(user.as_str()), mxc, width, height, content_disposition, content_type)?
+		} else {
+			self.db
+				.create_file_metadata(None, mxc, width, height, content_disposition, content_type)?
+		};
+
+		//TODO: Dangling metadata in database if creation fails
+		let mut f = self.create_media_file(&key).await?;
+		f.write_all(file).await?;
+
+		Ok(())
+	}
+
+	/// Downloads a file's thumbnail.
+	///
+	/// Here's an example on how it works:
+	///
+	/// - Client requests an image with width=567, height=567
+	/// - Server rounds that up to (800, 600), so it doesn't have to save too
+	///   many thumbnails
+	/// - Server rounds that up again to (958, 600) to fix the aspect ratio
+	///   (only for width,height>96)
+	/// - Server creates the thumbnail and sends it to the user
+	///
+	/// For width,height <= 96 the server uses another thumbnailing algorithm
+	/// which crops the image afterwards.
+	#[tracing::instrument(skip(self), name = "thumbnail", level = "debug")]
+	pub async fn get_thumbnail(&self, mxc: &str, width: u32, height: u32) -> Result<Option<FileMeta>> {
+		// 0, 0 because that's the original file
+		let (width, height, crop) = thumbnail_properties(width, height).unwrap_or((0, 0, false));
+
+		if let Ok(metadata) = self.db.search_file_metadata(mxc, width, height) {
+			self.get_thumbnail_saved(metadata).await
+		} else if let Ok(metadata) = self.db.search_file_metadata(mxc, 0, 0) {
+			self.get_thumbnail_generate(mxc, width, height, crop, metadata)
+				.await
+		} else {
+			Ok(None)
+		}
+	}
+
+	/// Using saved thumbnail
+	#[tracing::instrument(skip(self), name = "saved", level = "debug")]
+	async fn get_thumbnail_saved(&self, data: Metadata) -> Result<Option<FileMeta>> {
+		let mut content = Vec::new();
+		let path = self.get_media_file(&data.key);
+		fs::File::open(path)
+			.await?
+			.read_to_end(&mut content)
+			.await?;
+
+		Ok(Some(into_filemeta(data, content)))
+	}
+
+	/// Generate a thumbnail
+	#[tracing::instrument(skip(self), name = "generate", level = "debug")]
+	async fn get_thumbnail_generate(
+		&self, mxc: &str, width: u32, height: u32, crop: bool, data: Metadata,
+	) -> Result<Option<FileMeta>> {
+		let mut content = Vec::new();
+		let path = self.get_media_file(&data.key);
+		fs::File::open(path)
+			.await?
+			.read_to_end(&mut content)
+			.await?;
+
+		let Ok(image) = image::load_from_memory(&content) else {
+			// Couldn't parse file to generate thumbnail, send original
+			return Ok(Some(into_filemeta(data, content)));
+		};
+
+		if width > image.width() || height > image.height() {
+			return Ok(Some(into_filemeta(data, content)));
+		}
+
+		let mut thumbnail_bytes = Vec::new();
+		let thumbnail = thumbnail_generate(&image, width, height, crop)?;
+		thumbnail.write_to(&mut Cursor::new(&mut thumbnail_bytes), image::ImageFormat::Png)?;
+
+		// Save thumbnail in database so we don't have to generate it again next time
+		let thumbnail_key = self.db.create_file_metadata(
+			None,
+			mxc,
+			width,
+			height,
+			data.content_disposition.as_deref(),
+			data.content_type.as_deref(),
+		)?;
+
+		let mut f = self.create_media_file(&thumbnail_key).await?;
+		f.write_all(&thumbnail_bytes).await?;
+
+		Ok(Some(into_filemeta(data, thumbnail_bytes)))
+	}
+}
+
+fn thumbnail_generate(image: &DynamicImage, width: u32, height: u32, crop: bool) -> Result<DynamicImage> {
+	let thumbnail = if crop {
+		image.resize_to_fill(width, height, FilterType::CatmullRom)
+	} else {
+		let (exact_width, exact_height) = thumbnail_dimension(image, width, height)?;
+		image.thumbnail_exact(exact_width, exact_height)
+	};
+
+	Ok(thumbnail)
+}
+
+fn thumbnail_dimension(image: &DynamicImage, width: u32, height: u32) -> Result<(u32, u32)> {
+	let original_width = image.width();
+	let original_height = image.height();
+
+	let ratio = Sat(original_width) * Sat(height);
+	let nratio = Sat(width) * Sat(original_height);
+	let use_width = nratio <= ratio;
+
+	let intermediate = if use_width {
+		Sat(original_height) * Sat(checked!(width / original_width)?)
+	} else {
+		Sat(original_width) * Sat(checked!(height / original_height)?)
+	};
+
+	let dims = if use_width {
+		(width, intermediate.0)
+	} else {
+		(intermediate.0, height)
+	};
+
+	Ok(dims)
+}
+
+/// Returns width, height of the thumbnail and whether it should be cropped.
+/// Returns None when the server should send the original file.
+fn thumbnail_properties(width: u32, height: u32) -> Option<(u32, u32, bool)> {
+	match (width, height) {
+		(0..=32, 0..=32) => Some((32, 32, true)),
+		(0..=96, 0..=96) => Some((96, 96, true)),
+		(0..=320, 0..=240) => Some((320, 240, false)),
+		(0..=640, 0..=480) => Some((640, 480, false)),
+		(0..=800, 0..=600) => Some((800, 600, false)),
+		_ => None,
+	}
+}
+
+fn into_filemeta(data: Metadata, content: Vec<u8>) -> FileMeta {
+	FileMeta {
+		content: Some(content),
+		content_type: data.content_type,
+		content_disposition: data.content_disposition,
+	}
+}