From b903b46d16352e536e7fbecfbb2ab1f97ace8d4a Mon Sep 17 00:00:00 2001 From: Jason Volk Date: Mon, 15 Jul 2024 07:21:10 +0000 Subject: [PATCH] split thumbnailing related into unit Signed-off-by: Jason Volk --- src/service/media/data.rs | 18 +++- src/service/media/mod.rs | 159 ++---------------------------- src/service/media/thumbnail.rs | 171 +++++++++++++++++++++++++++++++++ 3 files changed, 195 insertions(+), 153 deletions(-) create mode 100644 src/service/media/thumbnail.rs diff --git a/src/service/media/data.rs b/src/service/media/data.rs index 186d502d..e5856bbf 100644 --- a/src/service/media/data.rs +++ b/src/service/media/data.rs @@ -12,6 +12,13 @@ pub(crate) struct Data { url_previews: Arc, } +#[derive(Debug)] +pub(super) struct Metadata { + pub(super) content_disposition: Option, + pub(super) content_type: Option, + pub(super) key: Vec, +} + impl Data { pub(super) fn new(db: &Arc) -> Self { Self { @@ -104,9 +111,7 @@ impl Data { Ok(keys) } - pub(super) fn search_file_metadata( - &self, mxc: &str, width: u32, height: u32, - ) -> Result<(Option, Option, Vec)> { + pub(super) fn search_file_metadata(&self, mxc: &str, width: u32, height: u32) -> Result { 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 diff --git a/src/service/media/mod.rs b/src/service/media/mod.rs index 2f687925..c969d1e8 100644 --- a/src/service/media/mod.rs +++ b/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 data::Data; -use image::imageops::FilterType; +use conduit::{debug, debug_error, err, error, utils, Result, Server}; +use data::{Data, Metadata}; 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, 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> { - 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> { - 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 { self.db.get_url_preview(url) } /// TODO: use this? diff --git a/src/service/media/thumbnail.rs b/src/service/media/thumbnail.rs new file mode 100644 index 00000000..50be4ee1 --- /dev/null +++ b/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, 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> { + // 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> { + 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> { + 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 { + 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) -> FileMeta { + FileMeta { + content: Some(content), + content_type: data.content_type, + content_disposition: data.content_disposition, + } +}