From 6763e53d4105cc2f4a207f4c9e45564acba51cca Mon Sep 17 00:00:00 2001 From: kieran Date: Sat, 25 Jan 2025 23:22:39 +0000 Subject: [PATCH] feat: generate thumbnails --- src/bin/main.rs | 4 ++ src/processing/mod.rs | 48 +++++++++++++-- src/routes/mod.rs | 132 +++++++++++++++++++++++++++++++----------- 3 files changed, 145 insertions(+), 39 deletions(-) diff --git a/src/bin/main.rs b/src/bin/main.rs index 03fa101..cbe29a7 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -89,6 +89,10 @@ async fn main() -> Result<(), Error> { { rocket = rocket.mount("/", routes::nip96_routes()); } + #[cfg(feature = "media-compression")] + { + rocket = rocket.mount("/", routes![routes::get_blob_thumb]); + } if let Err(e) = rocket.launch().await { error!("Rocker error {}", e); Err(Error::from(e)) diff --git a/src/processing/mod.rs b/src/processing/mod.rs index 174503b..2dc0f6d 100644 --- a/src/processing/mod.rs +++ b/src/processing/mod.rs @@ -1,8 +1,7 @@ -use std::path::{Path, PathBuf}; - use anyhow::{bail, Error, Result}; use ffmpeg_rs_raw::ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_YUV420P; -use ffmpeg_rs_raw::{Demuxer, DemuxerInfo, Encoder, StreamType, Transcoder}; +use ffmpeg_rs_raw::{Demuxer, DemuxerInfo, Encoder, Muxer, StreamType, Transcoder}; +use std::path::{Path, PathBuf}; use uuid::Uuid; #[cfg(feature = "labels")] @@ -21,7 +20,7 @@ impl WebpProcessor { Self } - pub fn process_file( + pub fn compress( &mut self, input: &Path, mime_type: &str, @@ -68,6 +67,45 @@ impl WebpProcessor { }) } } + + pub fn thumbnail(&mut self, input: &Path, out_path: &Path) -> Result<()> { + use ffmpeg_rs_raw::ffmpeg_sys_the_third::AVCodecID::AV_CODEC_ID_WEBP; + + unsafe { + let mut input = Demuxer::new(input.to_str().unwrap())?; + + let probe = input.probe_input()?; + + let image_stream = probe + .streams + .iter() + .find(|c| c.stream_type == StreamType::Video) + .ok_or(Error::msg("No image found, cant compress"))?; + + let w = 512u16; + let scale = w as f32 / image_stream.width as f32; + let h = (image_stream.height as f32 * scale) as u16; + + let enc = Encoder::new(AV_CODEC_ID_WEBP)? + .with_height(h as i32) + .with_width(w as i32) + .with_pix_fmt(AV_PIX_FMT_YUV420P) + .with_framerate(1.0)? + .open(None)?; + + let mut trans = Transcoder::new_custom_io( + input, + Muxer::builder() + .with_output_path(out_path.to_str().unwrap(), Some("webp"))? + .build()?, + ); + + trans.transcode_stream(image_stream, enc)?; + trans.run(None)?; + + Ok(()) + } + } } pub struct NewFileProcessorResult { @@ -92,7 +130,7 @@ pub fn compress_file( if mime_type.starts_with("image/") { let mut proc = WebpProcessor::new(); - return proc.process_file(stream, mime_type, out_dir); + return proc.compress(stream, mime_type, out_dir); } bail!("No media processor") } diff --git a/src/routes/mod.rs b/src/routes/mod.rs index 18022d3..9a76344 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -1,5 +1,7 @@ use crate::db::{Database, FileUpload}; use crate::filesystem::FileStore; +#[cfg(feature = "media-compression")] +use crate::processing::WebpProcessor; pub use crate::routes::admin::admin_routes; #[cfg(feature = "blossom")] pub use crate::routes::blossom::blossom_routes; @@ -16,6 +18,7 @@ use rocket::http::{ContentType, Header, Status}; use rocket::response::Responder; use rocket::serde::Serialize; use rocket::{Request, Response, State}; +use std::env::temp_dir; use std::io::SeekFrom; use std::ops::Range; use std::pin::{pin, Pin}; @@ -170,49 +173,61 @@ impl AsyncRead for RangeBody { impl<'r> Responder<'r, 'static> for FilePayload { fn respond_to(self, request: &'r Request<'_>) -> rocket::response::Result<'static> { let mut response = Response::new(); + response.set_header(Header::new("cache-control", "max-age=31536000, immutable")); // handle ranges #[cfg(feature = "ranges")] { - response.set_header(Header::new("accept-ranges", "bytes")); - if let Some(r) = request.headers().get("range").next() { - if let Ok(ranges) = parse_range_header(r) { - if ranges.ranges.len() > 1 { - warn!("Multipart ranges are not supported, fallback to non-range request"); - response.set_streamed_body(self.file); - } else { - const MAX_UNBOUNDED_RANGE: u64 = 1024 * 1024; - let single_range = ranges.ranges.first().unwrap(); - let range_start = match single_range.start { - StartPosition::Index(i) => i, - StartPosition::FromLast(i) => self.info.size - i, - }; - let range_end = match single_range.end { - EndPosition::Index(i) => i, - EndPosition::LastByte => { - (range_start + MAX_UNBOUNDED_RANGE).min(self.info.size) - } - }; - let r_len = range_end - range_start; - let r_body = RangeBody::new(self.file, range_start..range_end); - - response.set_status(Status::PartialContent); - response.set_header(Header::new("content-length", r_len.to_string())); - response.set_header(Header::new( - "content-range", - format!("bytes {}-{}/{}", range_start, range_end - 1, self.info.size), - )); - response.set_streamed_body(Box::pin(r_body)); - } - } + const MAX_UNBOUNDED_RANGE: u64 = 1024 * 1024; + // only use range response for files > 1MiB + if self.info.size < MAX_UNBOUNDED_RANGE { + response.set_sized_body(None, self.file); } else { - response.set_streamed_body(self.file); + response.set_header(Header::new("accept-ranges", "bytes")); + if let Some(r) = request.headers().get("range").next() { + if let Ok(ranges) = parse_range_header(r) { + if ranges.ranges.len() > 1 { + warn!( + "Multipart ranges are not supported, fallback to non-range request" + ); + response.set_streamed_body(self.file); + } else { + let single_range = ranges.ranges.first().unwrap(); + let range_start = match single_range.start { + StartPosition::Index(i) => i, + StartPosition::FromLast(i) => self.info.size - i, + }; + let range_end = match single_range.end { + EndPosition::Index(i) => i, + EndPosition::LastByte => { + (range_start + MAX_UNBOUNDED_RANGE).min(self.info.size) + } + }; + let r_len = range_end - range_start; + let r_body = RangeBody::new(self.file, range_start..range_end); + + response.set_status(Status::PartialContent); + response.set_header(Header::new("content-length", r_len.to_string())); + response.set_header(Header::new( + "content-range", + format!( + "bytes {}-{}/{}", + range_start, + range_end - 1, + self.info.size + ), + )); + response.set_streamed_body(Box::pin(r_body)); + } + } + } else { + response.set_sized_body(None, self.file); + } } } #[cfg(not(feature = "ranges"))] { - response.set_streamed_body(self.file); - response.set_header(Header::new("content-length", self.info.size.to_string())); + response.set_sized_body(None, self.file); } if let Ok(ct) = ContentType::from_str(&self.info.mime_type) { @@ -352,6 +367,55 @@ pub async fn head_blob(sha256: &str, fs: &State) -> Status { } } +/// Generate thumbnail for image / video +#[cfg(feature = "media-compression")] +#[rocket::get("/thumb/")] +pub async fn get_blob_thumb( + sha256: &str, + fs: &State, + db: &State, +) -> Result { + let sha256 = if sha256.contains(".") { + sha256.split('.').next().unwrap() + } else { + sha256 + }; + let id = if let Ok(i) = hex::decode(sha256) { + i + } else { + return Err(Status::NotFound); + }; + + if id.len() != 32 { + return Err(Status::NotFound); + } + if let Ok(Some(info)) = db.get_file(&id).await { + let file_path = fs.get(&id); + + let mut thumb_file = temp_dir().join(format!("thumb_{}", sha256)); + thumb_file.set_extension("webp"); + + if !thumb_file.exists() { + let mut p = WebpProcessor::new(); + if p.thumbnail(&file_path, &thumb_file).is_err() { + return Err(Status::InternalServerError); + } + }; + + if let Ok(f) = File::open(&thumb_file).await { + return Ok(FilePayload { + file: f, + info: FileUpload { + size: thumb_file.metadata().unwrap().len(), + mime_type: "image/webp".to_string(), + ..info + }, + }); + } + } + Err(Status::NotFound) +} + /// Legacy URL redirect for void.cat uploads #[rocket::get("/d/")] pub async fn void_cat_redirect(id: &str, settings: &State) -> Option {