From 7d86f944586e37534f54fdcc14cb176a2f8f935b Mon Sep 17 00:00:00 2001 From: kieran Date: Fri, 10 May 2024 11:54:40 +0100 Subject: [PATCH] Compress images --- .gitignore | 3 +- Cargo.lock | 90 +++++++++++++++++ Cargo.toml | 2 + Dockerfile | 27 +++++ src/filesystem.rs | 61 +++++++++--- src/main.rs | 1 + src/processing/image.rs | 214 ++++++++++++++++++++++++++++++++++++++++ src/processing/mod.rs | 45 +++++++++ src/routes/blossom.rs | 11 ++- src/routes/nip96.rs | 11 ++- ui/index.html | 60 ++++++++++- 11 files changed, 500 insertions(+), 25 deletions(-) create mode 100644 src/processing/image.rs create mode 100644 src/processing/mod.rs diff --git a/.gitignore b/.gitignore index b143b9c..f3a364a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ target/ -data/ \ No newline at end of file +data/ +.idea/ \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 55a14e7..6d8f8ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -195,6 +195,26 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "383d29d513d8764dcdc42ea295d979eb99c3c9f00607b3692cf68a431f7dca72" +[[package]] +name = "bindgen" +version = "0.64.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4243e6031260db77ede97ad86c27e501d646a27ab57b59a574f725d98ab1fb4" +dependencies = [ + "bitflags 1.3.2", + "cexpr", + "clang-sys", + "lazy_static", + "lazycell", + "peeking_take_while", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 1.0.109", +] + [[package]] name = "bip39" version = "2.0.0" @@ -319,6 +339,15 @@ version = "1.0.95" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d32a725bc159af97c3e629873bb9f88fb8cf8a4867175f76dc987815ea07c83b" +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + [[package]] name = "cfg-if" version = "1.0.0" @@ -375,6 +404,17 @@ dependencies = [ "zeroize", ] +[[package]] +name = "clang-sys" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67523a3b4be3ce1989d607a828d036249522dd9c1c8de7f4dd2dae43a37369d1" +dependencies = [ + "glob", + "libc", + "libloading", +] + [[package]] name = "config" version = "0.14.0" @@ -653,6 +693,20 @@ version = "2.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "658bd65b1cf4c852a3cc96f18a8ce7b5640f6b703f905c7d74532294c2a63984" +[[package]] +name = "ffmpeg-sys-the-third" +version = "1.1.1+ffmpeg-6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94a4b2e9c02074c0ee85661b23b3ac849bad6afc554b503c183975f5e2e0d3de" +dependencies = [ + "bindgen", + "cc", + "libc", + "num_cpus", + "pkg-config", + "vcpkg", +] + [[package]] name = "figment" version = "0.10.18" @@ -1225,12 +1279,28 @@ dependencies = [ "spin 0.5.2", ] +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + [[package]] name = "libc" version = "0.2.153" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c198f91728a82281a64e1f4f9eeb25d82cb32a5de251c6bd1b5154d63a8e7bd" +[[package]] +name = "libloading" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c2a198fb6b0eada2a8df47933734e6d35d350665a33a3593d7164fa52c75c19" +dependencies = [ + "cfg-if", + "windows-targets 0.52.5", +] + [[package]] name = "libm" version = "0.2.8" @@ -1603,6 +1673,12 @@ dependencies = [ "syn 2.0.60", ] +[[package]] +name = "peeking_take_while" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" + [[package]] name = "pem-rfc7468" version = "0.7.0" @@ -2085,6 +2161,12 @@ version = "0.1.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d626bb9dae77e28219937af045c257c28bfd3f69333c512553507f5f9798cb76" +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustix" version = "0.38.34" @@ -2289,6 +2371,12 @@ dependencies = [ "lazy_static", ] +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + [[package]] name = "signal-hook-registry" version = "1.4.2" @@ -3082,7 +3170,9 @@ dependencies = [ "base64 0.21.7", "chrono", "config", + "ffmpeg-sys-the-third", "hex", + "libc", "log", "nostr", "pretty_env_logger", diff --git a/Cargo.toml b/Cargo.toml index ac1afa7..a896550 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,3 +20,5 @@ sha2 = "0.10.8" sqlx = { version = "0.7.4", features = ["mysql", "runtime-tokio", "chrono"] } config = { version = "0.14.0", features = ["toml"] } chrono = { version = "0.4.38", features = ["serde"] } +ffmpeg-sys-the-third = { version = "1.1.1",features = ["default"] } +libc = "0.2.153" diff --git a/Dockerfile b/Dockerfile index 249fb8d..f14f3e7 100644 --- a/Dockerfile +++ b/Dockerfile @@ -3,11 +3,38 @@ ARG IMAGE=rust:bookworm FROM $IMAGE as build WORKDIR /app/src COPY . . +ENV FFMPEG_DIR=/app/ffmpeg +RUN apt update && \ + apt install -y \ + build-essential \ + libx264-dev \ + libwebp-dev \ + nasm \ + libclang-dev && \ + rm -rf /var/lib/apt/lists/* +RUN git clone --depth=1 https://git.ffmpeg.org/ffmpeg.git && \ + cd ffmpeg && \ + ./configure \ + --prefix=$FFMPEG_DIR \ + --disable-programs \ + --disable-doc \ + --disable-network \ + --enable-gpl \ + --enable-version3 \ + --enable-libx264 \ + --enable-libwebp \ + --disable-static \ + --enable-shared && \ + make -j8 && make install RUN cargo install --path . --root /app/build FROM $IMAGE as runner WORKDIR /app +RUN apt update && \ + apt install -y libx264-164 libwebp7 && \ + rm -rf /var/lib/apt/lists/* COPY --from=build /app/build . COPY --from=build /app/src/ui ui COPY --from=build /app/src/config.toml . +COPY --from=build /app/ffmpeg/lib/ /lib ENTRYPOINT ["/app/bin/void_cat"] \ No newline at end of file diff --git a/src/filesystem.rs b/src/filesystem.rs index 26c8b21..5fa18d4 100644 --- a/src/filesystem.rs +++ b/src/filesystem.rs @@ -1,7 +1,9 @@ use std::env::temp_dir; -use std::io::{SeekFrom}; +use std::fs; +use std::io::SeekFrom; use std::path::{Path, PathBuf}; -use std::{fs}; +use std::sync::{Arc, Mutex}; +use std::time::SystemTime; use anyhow::Error; use log::info; @@ -9,6 +11,7 @@ use sha2::{Digest, Sha256}; use tokio::fs::File; use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeekExt}; +use crate::processing::{FileProcessor, MediaProcessor}; use crate::settings::Settings; #[derive(Clone)] @@ -20,12 +23,14 @@ pub struct FileSystemResult { pub struct FileStore { path: String, + processor: Arc>, } impl FileStore { pub fn new(settings: Settings) -> Self { Self { path: settings.storage_dir, + processor: Arc::new(Mutex::new(MediaProcessor::new())), } } @@ -35,23 +40,49 @@ impl FileStore { } /// Store a new file - pub async fn put(&self, mut stream: TStream) -> Result - where - TStream: AsyncRead + Unpin, + pub async fn put(&self, mut stream: TStream, mime_type: &str) -> Result + where + TStream: AsyncRead + Unpin, { let random_id = uuid::Uuid::new_v4(); - let tmp_path = FileStore::map_temp(random_id); - let mut file = File::options() - .create(true) - .write(true) - .read(true) - .open(tmp_path.clone()) - .await?; - let n = tokio::io::copy(&mut stream, &mut file).await?; + let (n, hash, tmp_path) = { + let tmp_path = FileStore::map_temp(random_id); + let mut file = File::options() + .create(true) + .write(true) + .read(true) + .open(tmp_path.clone()) + .await?; + tokio::io::copy(&mut stream, &mut file).await?; - info!("File saved to temp path: {}", tmp_path.to_str().unwrap()); - let hash = FileStore::hash_file(&mut file).await?; + info!("File saved to temp path: {}", tmp_path.to_str().unwrap()); + + let start = SystemTime::now(); + let new_temp = { + let mut p_lock = self.processor.lock().expect("asd"); + p_lock.process_file(tmp_path.clone(), mime_type)? + }; + let old_size = tmp_path.metadata()?.len(); + let new_size = new_temp.metadata()?.len(); + info!("Compressed media: ratio={:.2}x, old_size={:.3}kb, new_size={:.3}kb, duration={:.2}ms", + old_size as f32 / new_size as f32, + old_size as f32 / 1024.0, + new_size as f32 / 1024.0, + SystemTime::now().duration_since(start).unwrap().as_micros() as f64 / 1000.0 + ); + + let mut file = File::options() + .create(true) + .write(true) + .read(true) + .open(new_temp.clone()) + .await?; + let n = file.metadata().await?.len(); + let hash = FileStore::hash_file(&mut file).await?; + fs::remove_file(tmp_path)?; + (n, hash, new_temp) + }; let dst_path = self.map_path(&hash); fs::create_dir_all(dst_path.parent().unwrap())?; if let Err(e) = fs::copy(&tmp_path, &dst_path) { diff --git a/src/main.rs b/src/main.rs index b6a2897..f32f1f9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -18,6 +18,7 @@ mod db; mod filesystem; mod routes; mod settings; +mod processing; #[rocket::main] async fn main() -> Result<(), Error> { diff --git a/src/processing/image.rs b/src/processing/image.rs new file mode 100644 index 0000000..308715b --- /dev/null +++ b/src/processing/image.rs @@ -0,0 +1,214 @@ +use std::mem::transmute; +use std::path::PathBuf; +use std::ptr; + +use anyhow::Error; +use ffmpeg_sys_the_third::{av_dump_format, av_frame_alloc, av_frame_copy_props, av_frame_free, av_guess_format, av_interleaved_write_frame, av_packet_alloc, av_packet_free, av_packet_rescale_ts, av_packet_unref, av_read_frame, av_write_trailer, avcodec_alloc_context3, avcodec_find_decoder, avcodec_find_encoder, avcodec_free_context, avcodec_open2, avcodec_parameters_from_context, avcodec_parameters_to_context, avcodec_receive_frame, avcodec_receive_packet, avcodec_send_frame, avcodec_send_packet, AVERROR, AVERROR_EOF, avformat_alloc_output_context2, avformat_close_input, avformat_find_stream_info, avformat_free_context, avformat_init_output, avformat_new_stream, avformat_open_input, avformat_write_header, AVFormatContext, AVIO_FLAG_WRITE, avio_open, sws_getContext, sws_scale_frame}; +use ffmpeg_sys_the_third::AVPictureType::AV_PICTURE_TYPE_NONE; +use ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_YUV420P; +use libc::EAGAIN; + +use crate::processing::FileProcessor; + +pub struct ImageProcessor {} + +impl ImageProcessor { + pub fn new() -> Self { + Self {} + } +} + +impl FileProcessor for ImageProcessor { + fn process_file(&mut self, in_file: PathBuf, mime_type: &str) -> Result { + unsafe { + let mut out_path = in_file.clone(); + out_path.set_extension("_compressed"); + + let mut dec_fmt: *mut AVFormatContext = ptr::null_mut(); + let ret = avformat_open_input(&mut dec_fmt, + format!("{}\0", in_file.into_os_string().into_string().unwrap()).as_ptr() as *const libc::c_char, + ptr::null_mut(), + ptr::null_mut()); + if ret < 0 { + panic!("Failed to create input context") + } + + let ret = avformat_find_stream_info(dec_fmt, ptr::null_mut()); + if ret < 0 { + panic!("Failed to probe input") + } + + let in_stream = *(*dec_fmt).streams.add(0); + let decoder = avcodec_find_decoder((*(*in_stream).codecpar).codec_id); + let mut dec_ctx = avcodec_alloc_context3(decoder); + if dec_ctx.is_null() { + panic!("Failed to open decoder") + } + + let ret = avcodec_parameters_to_context(dec_ctx, (*in_stream).codecpar); + if ret < 0 { + panic!("Failed to copy codec params to decoder") + } + + let ret = avcodec_open2(dec_ctx, decoder, ptr::null_mut()); + if ret < 0 { + panic!("Failed to open decoder") + } + + let out_format = av_guess_format("webp\0".as_ptr() as *const libc::c_char, + ptr::null_mut(), + "image/webp\0".as_ptr() as *const libc::c_char); + let out_filename = format!("{}\0", out_path.clone().into_os_string().into_string().unwrap()); + let mut out_fmt: *mut AVFormatContext = ptr::null_mut(); + let ret = avformat_alloc_output_context2(&mut out_fmt, + out_format, + ptr::null_mut(), + out_filename.as_ptr() as *const libc::c_char); + if ret < 0 { + panic!("Failed to create output context") + } + + let ret = avio_open(&mut (*out_fmt).pb, (*out_fmt).url, AVIO_FLAG_WRITE); + if ret < 0 { + panic!("Failed to open output IO") + } + + let out_codec = avcodec_find_encoder((*(*out_fmt).oformat).video_codec); + let stream = avformat_new_stream(out_fmt, out_codec); + + let mut encoder = avcodec_alloc_context3(out_codec); + if encoder.is_null() { + panic!("Failed to create encoder context") + } + (*encoder).width = (*(*in_stream).codecpar).width; + (*encoder).height = (*(*in_stream).codecpar).height; + (*encoder).pix_fmt = AV_PIX_FMT_YUV420P; + (*encoder).time_base = (*in_stream).time_base; + (*encoder).framerate = (*in_stream).avg_frame_rate; + (*stream).time_base = (*encoder).time_base; + (*stream).avg_frame_rate = (*encoder).framerate; + (*stream).r_frame_rate = (*encoder).framerate; + + let ret = avcodec_open2(encoder, out_codec, ptr::null_mut()); + if ret < 0 { + panic!("Failed to open encoder"); + } + + let ret = avcodec_parameters_from_context((*stream).codecpar, encoder); + if ret < 0 { + panic!("Failed to open encoder"); + } + + let ret = avformat_init_output(out_fmt, ptr::null_mut()); + if ret < 0 { + panic!("Failed to write output"); + } + + av_dump_format(out_fmt, 0, ptr::null_mut(), 1); + + let sws_ctx = sws_getContext((*(*in_stream).codecpar).width, + (*(*in_stream).codecpar).height, + transmute((*(*in_stream).codecpar).format), + (*(*stream).codecpar).width, + (*(*stream).codecpar).height, + transmute((*(*stream).codecpar).format), + 0, ptr::null_mut(), ptr::null_mut(), ptr::null_mut()); + if sws_ctx.is_null() { + panic!("Failed to create sws context"); + } + + let ret = avformat_write_header(out_fmt, ptr::null_mut()); + if ret < 0 { + panic!("Failed to write header to output"); + } + + let mut pkt = av_packet_alloc(); + loop { + let ret = av_read_frame(dec_fmt, pkt); + if ret < 0 { + break; + } + + let in_stream = *(*dec_fmt).streams.add((*pkt).stream_index as usize); + let out_stream = *(*out_fmt).streams; + av_packet_rescale_ts(pkt, (*in_stream).time_base, (*out_stream).time_base); + + let ret = avcodec_send_packet(dec_ctx, pkt); + if ret < 0 { + panic!("Failed to decode packet"); + } + + + let mut frame = av_frame_alloc(); + let mut frame_out = av_frame_alloc(); + loop { + let ret = avcodec_receive_frame(dec_ctx, frame); + if ret == AVERROR_EOF || ret == AVERROR(EAGAIN) { + break; + } else if ret < 0 { + panic!("Frame read error") + } + + av_frame_copy_props(frame_out, frame); + let ret = sws_scale_frame(sws_ctx, frame_out, frame); + if ret < 0 { + panic!("Failed to scale frame") + } + + (*frame_out).pict_type = AV_PICTURE_TYPE_NONE; + (*frame_out).time_base = (*in_stream).time_base; + + let ret = avcodec_send_frame(encoder, frame_out); + if ret < 0 { + panic!("Failed to encode frame") + } + av_packet_unref(pkt); + loop { + let ret = avcodec_receive_packet(encoder, pkt); + if ret == AVERROR_EOF || ret == AVERROR(EAGAIN) { + break; + } else if ret < 0 { + panic!("Frame read error") + } + + let ret = av_interleaved_write_frame(out_fmt, pkt); + if ret < 0 { + panic!("Failed to encode frame") + } + } + } + av_frame_free(&mut frame_out); + av_frame_free(&mut frame); + } + + // flush encoder + avcodec_send_frame(encoder, ptr::null_mut()); + loop { + let ret = avcodec_receive_packet(encoder, pkt); + if ret == AVERROR_EOF || ret == AVERROR(EAGAIN) { + break; + } else if ret < 0 { + panic!("Frame read error") + } + + let ret = av_interleaved_write_frame(out_fmt, pkt); + if ret < 0 { + panic!("Failed to encode frame") + } + } + + av_interleaved_write_frame(out_fmt, ptr::null_mut()); + av_write_trailer(out_fmt); + + av_packet_free(&mut pkt); + avcodec_free_context(&mut dec_ctx); + avcodec_free_context(&mut encoder); + + avformat_close_input(&mut dec_fmt); + avformat_free_context(dec_fmt); + avformat_free_context(out_fmt); + + Ok(out_path) + } + } +} \ No newline at end of file diff --git a/src/processing/mod.rs b/src/processing/mod.rs new file mode 100644 index 0000000..a9230c4 --- /dev/null +++ b/src/processing/mod.rs @@ -0,0 +1,45 @@ +use std::collections::HashMap; +use std::path::PathBuf; + +use anyhow::Error; +use crate::processing::image::ImageProcessor; + +mod image; + +pub(crate) trait FileProcessor { + fn process_file(&mut self, in_file: PathBuf, mime_type: &str) -> Result; +} + +pub(crate) struct MediaProcessor { + processors: HashMap>, +} + +impl MediaProcessor { + pub fn new() -> Self { + Self { + processors: HashMap::new() + } + } +} + +impl FileProcessor for MediaProcessor { + fn process_file(&mut self, in_file: PathBuf, mime_type: &str) -> Result { + if !self.processors.contains_key(mime_type) { + if mime_type.starts_with("image/") { + let ix = ImageProcessor::new(); + self.processors.insert(mime_type.to_string(), Box::new(ix)); + } else if mime_type.starts_with("video/") { + + } + } + + let proc = match self.processors.get_mut(mime_type) { + Some(p) => p, + None => { + return Err(Error::msg("Not supported mime type")); + } + }; + + proc.process_file(in_file, mime_type) + } +} \ No newline at end of file diff --git a/src/routes/blossom.rs b/src/routes/blossom.rs index ff8b798..66f3575 100644 --- a/src/routes/blossom.rs +++ b/src/routes/blossom.rs @@ -1,3 +1,4 @@ +use std::sync::{Mutex, RwLock}; use chrono::Utc; use log::{error}; use nostr::prelude::hex; @@ -102,8 +103,12 @@ async fn upload( if size.is_none() { return BlossomResponse::error("Invalid request, no size tag"); } + let mime_type = auth + .content_type + .unwrap_or("application/octet-stream".to_string()); + match fs - .put(data.open(ByteUnit::from(settings.max_upload_bytes))) + .put(data.open(ByteUnit::from(settings.max_upload_bytes)), &mime_type) .await { Ok(blob) => { @@ -119,9 +124,7 @@ async fn upload( user_id, name: name.unwrap_or("".to_string()), size: blob.size, - mime_type: auth - .content_type - .unwrap_or("application/octet-stream".to_string()), + mime_type, created: Utc::now(), }; if let Err(e) = db.add_file(&f).await { diff --git a/src/routes/nip96.rs b/src/routes/nip96.rs index 4b120ff..495fa64 100644 --- a/src/routes/nip96.rs +++ b/src/routes/nip96.rs @@ -1,11 +1,13 @@ use std::collections::HashMap; +use std::sync::Mutex; use chrono::Utc; +use rocket::{FromForm, Responder, Route, routes, State}; +use rocket::data::ByteUnit; use rocket::form::Form; use rocket::fs::TempFile; use rocket::serde::json::Json; use rocket::serde::Serialize; -use rocket::{routes, FromForm, Responder, Route, State}; use crate::auth::nip98::Nip98Auth; use crate::db::{Database, FileUpload}; @@ -156,7 +158,12 @@ async fn upload( Ok(f) => f, Err(e) => return Nip96Response::error(&format!("Could not open file: {}", e)), }; - match fs.put(file).await { + let mime_type = form.media_type + .unwrap_or("application/octet-stream"); + match fs + .put(file, mime_type) + .await + { Ok(blob) => { let pubkey_vec = auth.event.pubkey.to_bytes().to_vec(); let user_id = match db.upsert_user(&pubkey_vec).await { diff --git a/ui/index.html b/ui/index.html index 099876b..ced2280 100644 --- a/ui/index.html +++ b/ui/index.html @@ -9,10 +9,64 @@ color: white; } + -

- Welcome to void_cat_rs -

+

+ Welcome to void_cat_rs +

+
+

+ Upload a file using NIP-96 +

+ + You must have a nostr extension for this to work + +
+
+ + +
+

 
 
\ No newline at end of file