feat: generate thumbnails

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))

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")
 }

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"));
+            const MAX_UNBOUNDED_RANGE: u64 = 1024 * 1024;
-            if let Some(r) = request.headers().get("range").next() {
+            // only use range response for files > 1MiB
-                if let Ok(ranges) = parse_range_header(r) {
+            if self.info.size < MAX_UNBOUNDED_RANGE {
-                    if ranges.ranges.len() > 1 {
+                response.set_sized_body(None, self.file);
-                        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));
-                    }
-                }
             } 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_sized_body(None, self.file);
-            response.set_header(Header::new("content-length", self.info.size.to_string()));
         }
 
         if let Ok(ct) = ContentType::from_str(&self.info.mime_type) {
@@ -352,6 +367,55 @@ pub async fn head_blob(sha256: &str, fs: &State<FileStore>) -> Status {
     }
 }
 
+/// Generate thumbnail for image / video
+#[cfg(feature = "media-compression")]
+#[rocket::get("/thumb/<sha256>")]
+pub async fn get_blob_thumb(
+    sha256: &str,
+    fs: &State<FileStore>,
+    db: &State<Database>,
+) -> Result<FilePayload, Status> {
+    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/<id>")]
 pub async fn void_cat_redirect(id: &str, settings: &State<Settings>) -> Option<NamedFile> {