split thumbnailing related into unit

Signed-off-by: Jason Volk <jason@zemos.net>
This commit is contained in:
Jason Volk 2024-07-15 07:21:10 +00:00
parent 167559bb27
commit b903b46d16
3 changed files with 195 additions and 153 deletions

View file

@ -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(
&self, mxc: &str, width: u32, height: u32,
) -> Result<(Option<String>, Option<String>, Vec<u8>)> {
pub(super) fn search_file_metadata(&self, mxc: &str, width: u32, height: u32) -> Result<Metadata> {
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

View file

@ -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<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?

View file

@ -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,
}
}