From da2448fae94c84ef969d046a87464aa529a18842 Mon Sep 17 00:00:00 2001 From: kieran Date: Mon, 13 May 2024 12:13:30 +0100 Subject: [PATCH] Compute blurhash --- Cargo.lock | 7 + Cargo.toml | 1 + src/filesystem.rs | 63 ++++--- src/processing/blurhash.rs | 42 +++++ src/processing/image.rs | 214 --------------------- src/processing/mod.rs | 57 +++--- src/processing/webp.rs | 373 +++++++++++++++++++++++++++++++++++++ src/routes/blossom.rs | 22 +-- src/routes/nip96.rs | 29 +-- 9 files changed, 521 insertions(+), 287 deletions(-) create mode 100644 src/processing/blurhash.rs delete mode 100644 src/processing/image.rs create mode 100644 src/processing/webp.rs diff --git a/Cargo.lock b/Cargo.lock index 6d8f8ec..13ad357 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -300,6 +300,12 @@ dependencies = [ "generic-array", ] +[[package]] +name = "blurhash" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "565b78e03039f24994c5bc87ff793987be98a9ff59fa4851b72bc2e630001c9d" + [[package]] name = "bumpalo" version = "3.16.0" @@ -3168,6 +3174,7 @@ version = "0.1.0" dependencies = [ "anyhow", "base64 0.21.7", + "blurhash", "chrono", "config", "ffmpeg-sys-the-third", diff --git a/Cargo.toml b/Cargo.toml index a896550..179f27a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,3 +22,4 @@ 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" +blurhash = "0.2.1" diff --git a/src/filesystem.rs b/src/filesystem.rs index 5fa18d4..31a95af 100644 --- a/src/filesystem.rs +++ b/src/filesystem.rs @@ -11,7 +11,7 @@ use sha2::{Digest, Sha256}; use tokio::fs::File; use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeekExt}; -use crate::processing::{FileProcessor, MediaProcessor}; +use crate::processing::{FileProcessor, FileProcessorResult, MediaProcessor}; use crate::settings::Settings; #[derive(Clone)] @@ -19,6 +19,10 @@ pub struct FileSystemResult { pub path: PathBuf, pub sha256: Vec, pub size: u64, + pub mime_type: String, + pub width: Option, + pub height: Option, + pub blur_hash: Option, } pub struct FileStore { @@ -46,8 +50,9 @@ impl FileStore { { let random_id = uuid::Uuid::new_v4(); - let (n, hash, tmp_path) = { - let tmp_path = FileStore::map_temp(random_id); + let mut mime_type = mime_type.to_string(); + let (n, hash, tmp_path, width, height, blur_hash) = { + let mut tmp_path = FileStore::map_temp(random_id); let mut file = File::options() .create(true) .write(true) @@ -59,29 +64,37 @@ impl FileStore { info!("File saved to temp path: {}", tmp_path.to_str().unwrap()); let start = SystemTime::now(); - let new_temp = { + let proc_result = { let mut p_lock = self.processor.lock().expect("asd"); - p_lock.process_file(tmp_path.clone(), mime_type)? + 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 - ); + if let FileProcessorResult::NewFile(new_temp) = proc_result { + mime_type = new_temp.mime_type; + let old_size = tmp_path.metadata()?.len(); + let new_size = new_temp.result.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) + // delete old temp + fs::remove_file(tmp_path)?; + file = File::options() + .create(true) + .write(true) + .read(true) + .open(new_temp.result.clone()) + .await?; + let n = file.metadata().await?.len(); + let hash = FileStore::hash_file(&mut file).await?; + (n, hash, new_temp.result, Some(new_temp.width), Some(new_temp.height), Some(new_temp.blur_hash)) + } else { + let n = file.metadata().await?.len(); + let hash = FileStore::hash_file(&mut file).await?; + (n, hash, tmp_path, None, None, None) + } }; let dst_path = self.map_path(&hash); fs::create_dir_all(dst_path.parent().unwrap())?; @@ -94,6 +107,10 @@ impl FileStore { size: n, sha256: hash, path: dst_path, + mime_type, + width, + height, + blur_hash, }) } } diff --git a/src/processing/blurhash.rs b/src/processing/blurhash.rs new file mode 100644 index 0000000..224aa6e --- /dev/null +++ b/src/processing/blurhash.rs @@ -0,0 +1,42 @@ +use std::{ptr, slice}; +use std::intrinsics::transmute; +use std::time::SystemTime; + +use anyhow::Error; +use blurhash::encode; +use ffmpeg_sys_the_third::{av_frame_alloc, av_frame_free, AVFrame, sws_freeContext, sws_getContext, sws_scale_frame}; +use ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_RGBA; +use log::info; + +pub unsafe fn make_blur_hash(frame: *mut AVFrame, detail: u32) -> Result { + let start = SystemTime::now(); + let sws_ctx = sws_getContext((*frame).width, + (*frame).height, + transmute((*frame).format), + (*frame).width, + (*frame).height, + AV_PIX_FMT_RGBA, + 0, ptr::null_mut(), ptr::null_mut(), ptr::null_mut()); + if sws_ctx.is_null() { + return Err(Error::msg("Failed to create sws context")); + } + + let mut dst_frame = av_frame_alloc(); + let ret = sws_scale_frame(sws_ctx, dst_frame, frame); + if ret < 0 { + return Err(Error::msg("Failed to scale frame (blurhash)")); + } + + let pic_slice = slice::from_raw_parts_mut((*dst_frame).data[0], ((*frame).width * (*frame).height * 4) as usize); + let bh = encode(detail, detail, + (*frame).width as u32, + (*frame).height as u32, + pic_slice, + )?; + + av_frame_free(&mut dst_frame); + sws_freeContext(sws_ctx); + + info!("Generated blurhash in {}ms", SystemTime::now().duration_since(start).unwrap().as_millis()); + Ok(bh) +} \ No newline at end of file diff --git a/src/processing/image.rs b/src/processing/image.rs deleted file mode 100644 index 308715b..0000000 --- a/src/processing/image.rs +++ /dev/null @@ -1,214 +0,0 @@ -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 index a9230c4..69267b9 100644 --- a/src/processing/mod.rs +++ b/src/processing/mod.rs @@ -1,45 +1,48 @@ -use std::collections::HashMap; use std::path::PathBuf; use anyhow::Error; -use crate::processing::image::ImageProcessor; -mod image; +use crate::processing::webp::WebpProcessor; + +mod webp; +mod blurhash; + +pub(crate) enum FileProcessorResult { + NewFile(NewFileProcessorResult), + Skip, +} + +pub(crate) struct NewFileProcessorResult { + pub result: PathBuf, + pub mime_type: String, + pub width: usize, + pub height: usize, + pub blur_hash: String, +} pub(crate) trait FileProcessor { - fn process_file(&mut self, in_file: PathBuf, mime_type: &str) -> Result; + fn process_file(&mut self, in_file: PathBuf, mime_type: &str) -> Result; } -pub(crate) struct MediaProcessor { - processors: HashMap>, -} +pub(crate) struct MediaProcessor {} impl MediaProcessor { pub fn new() -> Self { - Self { - processors: HashMap::new() - } + Self {} } } 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")); - } + fn process_file(&mut self, in_file: PathBuf, mime_type: &str) -> Result { + let proc = if mime_type.starts_with("image/") { + Some(WebpProcessor::new()) + } else { + None }; - - proc.process_file(in_file, mime_type) + if let Some(mut proc) = proc { + proc.process_file(in_file, mime_type) + } else { + Ok(FileProcessorResult::Skip) + } } } \ No newline at end of file diff --git a/src/processing/webp.rs b/src/processing/webp.rs new file mode 100644 index 0000000..12e0247 --- /dev/null +++ b/src/processing/webp.rs @@ -0,0 +1,373 @@ +use std::collections::HashMap; +use std::mem::transmute; +use std::path::PathBuf; +use std::ptr; + +use anyhow::Error; +use ffmpeg_sys_the_third::{AV_CODEC_FLAG_GLOBAL_HEADER, av_dump_format, av_find_best_stream, 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_PROFILE_H264_HIGH, av_read_frame, av_write_trailer, AVCodec, avcodec_alloc_context3, 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, AVCodecContext, AVCodecID, AVERROR, AVERROR_EOF, AVERROR_STREAM_NOT_FOUND, AVFMT_GLOBALHEADER, 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, AVMediaType, AVPacket, sws_freeContext, sws_getContext, sws_scale_frame, SwsContext}; +use ffmpeg_sys_the_third::AVMediaType::{AVMEDIA_TYPE_AUDIO, AVMEDIA_TYPE_VIDEO}; +use ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_YUV420P; +use libc::EAGAIN; + +use crate::processing::{FileProcessor, FileProcessorResult, NewFileProcessorResult}; +use crate::processing::blurhash::make_blur_hash; + +/// Image converter to WEBP +pub struct WebpProcessor { + encoders: HashMap, + decoders: HashMap, + scalers: HashMap, + stream_map: HashMap, + width: Option, + height: Option, + blur_hash: Option, +} + +unsafe impl Sync for WebpProcessor {} + +unsafe impl Send for WebpProcessor {} + +impl WebpProcessor { + pub fn new() -> Self { + Self { + encoders: HashMap::new(), + decoders: HashMap::new(), + scalers: HashMap::new(), + stream_map: HashMap::new(), + width: None, + height: None, + blur_hash: None, + } + } + + unsafe fn transcode_pkt(&mut self, pkt: *mut AVPacket, in_fmt: *mut AVFormatContext, out_fmt: *mut AVFormatContext) -> Result<(), Error> { + let idx = (*pkt).stream_index as usize; + let out_idx = match self.stream_map.get(&idx) { + Some(i) => i, + None => return Ok(()) + }; + let in_stream = *(*in_fmt).streams.add(idx); + let out_stream = *(*out_fmt).streams.add(*out_idx); + av_packet_rescale_ts(pkt, (*in_stream).time_base, (*out_stream).time_base); + + let dec_ctx = self.decoders.get_mut(&idx).expect("Missing decoder config"); + let enc_ctx = self.encoders.get_mut(&out_idx).expect("Missing encoder config"); + + let ret = avcodec_send_packet(*dec_ctx, pkt); + if ret < 0 { + return Err(Error::msg("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 { + return Err(Error::msg("Frame read error")); + } + + let frame_out = match self.scalers.get_mut(&out_idx) { + Some(sws) => { + av_frame_copy_props(frame_out, frame); + let ret = sws_scale_frame(*sws, frame_out, frame); + if ret < 0 { + return Err(Error::msg("Failed to scale frame")); + } + frame_out + } + None => frame + }; + + // take blur_hash from first video frame + if (*(*out_stream).codecpar).codec_type == AVMEDIA_TYPE_VIDEO && self.blur_hash.is_none() { + self.blur_hash = Some(make_blur_hash(frame_out, 9)?); + } + + let ret = avcodec_send_frame(*enc_ctx, frame_out); + if ret < 0 { + return Err(Error::msg("Failed to encode frame")); + } + av_packet_unref(pkt); + loop { + let ret = avcodec_receive_packet(*enc_ctx, pkt); + if ret == AVERROR_EOF || ret == AVERROR(EAGAIN) { + break; + } else if ret < 0 { + return Err(Error::msg("Frame read error")); + } + + av_packet_rescale_ts(pkt, (*in_stream).time_base, (*out_stream).time_base); + let ret = av_interleaved_write_frame(out_fmt, pkt); + if ret < 0 { + return Err(Error::msg("Failed to encode frame")); + } + } + } + av_frame_free(&mut frame_out); + av_frame_free(&mut frame); + Ok(()) + } + + unsafe fn setup_decoder(&mut self, in_fmt: *mut AVFormatContext, av_type: AVMediaType) -> Result { + let mut decoder: *const AVCodec = ptr::null_mut(); + let stream_idx = av_find_best_stream(in_fmt, av_type, -1, -1, &mut decoder, 0); + if stream_idx == AVERROR_STREAM_NOT_FOUND { + return Ok(stream_idx); + } + let decoder_ctx = avcodec_alloc_context3(decoder); + if decoder_ctx.is_null() { + return Err(Error::msg("Failed to open video decoder")); + } + + let in_stream = *(*in_fmt).streams.add(stream_idx as usize); + let ret = avcodec_parameters_to_context(decoder_ctx, (*in_stream).codecpar); + if ret < 0 { + return Err(Error::msg("Failed to copy codec params to decoder")); + } + + let ret = avcodec_open2(decoder_ctx, decoder, ptr::null_mut()); + if ret < 0 { + return Err(Error::msg("Failed to open decoder")); + } + + self.decoders.insert(stream_idx as usize, decoder_ctx); + Ok(stream_idx) + } + + unsafe fn setup_encoder(&mut self, in_fmt: *mut AVFormatContext, out_fmt: *mut AVFormatContext, in_idx: i32) -> Result<(), Error> { + let in_stream = *(*in_fmt).streams.add(in_idx as usize); + let stream_type = (*(*in_stream).codecpar).codec_type; + let out_codec = match stream_type { + AVMEDIA_TYPE_VIDEO => avcodec_find_encoder((*(*out_fmt).oformat).video_codec), + AVMEDIA_TYPE_AUDIO => avcodec_find_encoder((*(*out_fmt).oformat).audio_codec), + _ => ptr::null_mut() + }; + // not mapped ignore + if out_codec.is_null() { + return Ok(()); + } + let stream = avformat_new_stream(out_fmt, out_codec); + + let encoder_ctx = avcodec_alloc_context3(out_codec); + if encoder_ctx.is_null() { + return Err(Error::msg("Failed to create encoder context")); + } + + match stream_type { + AVMEDIA_TYPE_VIDEO => { + (*encoder_ctx).width = (*(*in_stream).codecpar).width; + (*encoder_ctx).height = (*(*in_stream).codecpar).height; + (*encoder_ctx).pix_fmt = AV_PIX_FMT_YUV420P; + (*encoder_ctx).time_base = (*in_stream).time_base; + (*encoder_ctx).framerate = (*in_stream).avg_frame_rate; + if (*out_codec).id == AVCodecID::AV_CODEC_ID_H264 { + (*encoder_ctx).profile = AV_PROFILE_H264_HIGH; + (*encoder_ctx).level = 50; + (*encoder_ctx).qmin = 20; + (*encoder_ctx).qmax = 30; + } + (*stream).time_base = (*encoder_ctx).time_base; + (*stream).avg_frame_rate = (*encoder_ctx).framerate; + (*stream).r_frame_rate = (*encoder_ctx).framerate; + } + AVMEDIA_TYPE_AUDIO => { + (*encoder_ctx).sample_rate = (*(*in_stream).codecpar).sample_rate; + (*encoder_ctx).sample_fmt = transmute((*(*in_stream).codecpar).format); + (*encoder_ctx).ch_layout = (*(*in_stream).codecpar).ch_layout; + (*encoder_ctx).time_base = (*in_stream).time_base; + (*stream).time_base = (*encoder_ctx).time_base; + } + _ => {} + } + + if (*(*out_fmt).oformat).flags & AVFMT_GLOBALHEADER == AVFMT_GLOBALHEADER { + (*encoder_ctx).flags |= AV_CODEC_FLAG_GLOBAL_HEADER as libc::c_int; + } + + let ret = avcodec_open2(encoder_ctx, out_codec, ptr::null_mut()); + if ret < 0 { + return Err(Error::msg("Failed to open encoder")); + } + + let ret = avcodec_parameters_from_context((*stream).codecpar, encoder_ctx); + if ret < 0 { + return Err(Error::msg("Failed to open encoder")); + } + + let out_idx = (*stream).index as usize; + // setup scaler if pix_fmt doesnt match + if stream_type == AVMEDIA_TYPE_VIDEO && + (*(*in_stream).codecpar).format != (*(*stream).codecpar).format { + 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() { + return Err(Error::msg("Failed to create sws context")); + } + self.scalers.insert(out_idx, sws_ctx); + } + + self.encoders.insert(out_idx, encoder_ctx); + self.stream_map.insert(in_idx as usize, out_idx); + Ok(()) + } + + unsafe fn flush_output(&mut self, out_fmt: *mut AVFormatContext) -> Result<(), Error> { + let mut pkt = av_packet_alloc(); + for encoder in self.encoders.values() { + 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 { + return Err(Error::msg("Frame read error")); + } + + let ret = av_interleaved_write_frame(out_fmt, pkt); + if ret < 0 { + return Err(Error::msg("Failed to encode frame")); + } + } + } + av_packet_free(&mut pkt); + Ok(()) + } + + unsafe fn free(&mut self) -> Result<(), Error> { + for decoders in self.decoders.values_mut() { + avcodec_free_context(&mut *decoders); + } + self.decoders.clear(); + for encoders in self.encoders.values_mut() { + avcodec_free_context(&mut *encoders); + } + self.encoders.clear(); + for scaler in self.scalers.values_mut() { + sws_freeContext(*scaler); + } + self.scalers.clear(); + + Ok(()) + } +} + +impl Drop for WebpProcessor { + fn drop(&mut self) { + unsafe { self.free().unwrap(); } + } +} + +impl FileProcessor for WebpProcessor { + 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 { + return Err(Error::msg("Failed to create input context")); + } + + let ret = avformat_find_stream_info(dec_fmt, ptr::null_mut()); + if ret < 0 { + return Err(Error::msg("Failed to probe input")); + } + let in_video_stream = self.setup_decoder(dec_fmt, AVMEDIA_TYPE_VIDEO)?; + let in_audio_stream = self.setup_decoder(dec_fmt, AVMEDIA_TYPE_AUDIO)?; + + let out_format = if mime_type.starts_with("image/") { + av_guess_format("webp\0".as_ptr() as *const libc::c_char, + ptr::null_mut(), + ptr::null_mut()) + } else if mime_type.starts_with("video/") { + av_guess_format("matroska\0".as_ptr() as *const libc::c_char, + ptr::null_mut(), + ptr::null_mut()) + } else { + return Err(Error::msg("Mime type not supported")); + }; + + 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 { + return Err(Error::msg("Failed to create output context")); + } + + let ret = avio_open(&mut (*out_fmt).pb, (*out_fmt).url, AVIO_FLAG_WRITE); + if ret < 0 { + return Err(Error::msg("Failed to open output IO")); + } + + if in_video_stream != AVERROR_STREAM_NOT_FOUND { + self.setup_encoder(dec_fmt, out_fmt, in_video_stream)?; + let video_stream = *(*dec_fmt).streams.add(in_video_stream as usize); + self.width = Some((*(*video_stream).codecpar).width as usize); + self.height = Some((*(*video_stream).codecpar).height as usize); + } + if in_audio_stream != AVERROR_STREAM_NOT_FOUND { + self.setup_encoder(dec_fmt, out_fmt, in_audio_stream)?; + } + + let ret = avformat_init_output(out_fmt, ptr::null_mut()); + if ret < 0 { + return Err(Error::msg("Failed to write output")); + } + + av_dump_format(dec_fmt, 0, ptr::null_mut(), 0); + av_dump_format(out_fmt, 0, ptr::null_mut(), 1); + + let ret = avformat_write_header(out_fmt, ptr::null_mut()); + if ret < 0 { + return Err(Error::msg("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; + } + self.transcode_pkt(pkt, dec_fmt, out_fmt)?; + } + + // flush encoder + self.flush_output(out_fmt)?; + + av_write_trailer(out_fmt); + av_packet_free(&mut pkt); + + self.free()?; + + avformat_close_input(&mut dec_fmt); + avformat_free_context(dec_fmt); + avformat_free_context(out_fmt); + + Ok(FileProcessorResult::NewFile( + NewFileProcessorResult { + result: out_path, + mime_type: "image/webp".to_string(), + width: self.width.unwrap_or(0), + height: self.height.unwrap_or(0), + blur_hash: match &self.blur_hash { + Some(s) => s.clone(), + None => "".to_string() + }, + })) + } + } +} \ No newline at end of file diff --git a/src/routes/blossom.rs b/src/routes/blossom.rs index 66f3575..9abf1c8 100644 --- a/src/routes/blossom.rs +++ b/src/routes/blossom.rs @@ -1,21 +1,21 @@ use std::sync::{Mutex, RwLock}; -use chrono::Utc; -use log::{error}; -use nostr::prelude::hex; -use nostr::{Tag}; -use rocket::data::{ByteUnit}; -use rocket::http::{Status}; +use chrono::Utc; +use log::error; +use nostr::prelude::hex; +use nostr::Tag; +use rocket::{Data, Route, routes, State}; +use rocket::data::ByteUnit; +use rocket::http::Status; use rocket::response::Responder; use rocket::serde::json::Json; -use rocket::{routes, Data, Route, State}; use serde::{Deserialize, Serialize}; use crate::auth::blossom::BlossomAuth; use crate::blob::BlobDescriptor; use crate::db::{Database, FileUpload}; use crate::filesystem::FileStore; -use crate::routes::{delete_file}; +use crate::routes::delete_file; use crate::settings::Settings; #[derive(Serialize, Deserialize)] @@ -106,7 +106,7 @@ async fn upload( let mime_type = auth .content_type .unwrap_or("application/octet-stream".to_string()); - + match fs .put(data.open(ByteUnit::from(settings.max_upload_bytes)), &mime_type) .await @@ -116,7 +116,7 @@ async fn upload( let user_id = match db.upsert_user(&pubkey_vec).await { Ok(u) => u, Err(e) => { - return BlossomResponse::error(format!("Failed to save file (db): {}", e)) + return BlossomResponse::error(format!("Failed to save file (db): {}", e)); } }; let f = FileUpload { @@ -124,7 +124,7 @@ async fn upload( user_id, name: name.unwrap_or("".to_string()), size: blob.size, - mime_type, + mime_type: blob.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 495fa64..ee3aae2 100644 --- a/src/routes/nip96.rs +++ b/src/routes/nip96.rs @@ -178,10 +178,7 @@ async fn upload( None => "".to_string(), }, size: blob.size, - mime_type: match &form.media_type { - Some(c) => c.to_string(), - None => "application/octet-stream".to_string(), - }, + mime_type: blob.mime_type, created: Utc::now(), }; if let Err(e) = db.add_file(&file_upload).await { @@ -189,17 +186,25 @@ async fn upload( } let hex_id = hex::encode(&file_upload.id); + let mut tags = vec![ + vec![ + "url".to_string(), + format!("{}/{}", &settings.public_url, &hex_id), + ], + vec!["x".to_string(), hex_id], + vec!["m".to_string(), file_upload.mime_type], + ]; + if let Some(bh) = blob.blur_hash { + tags.push(vec!["blurhash".to_string(), bh]); + } + if let (Some(w), Some(h)) = (blob.width, blob.height) { + tags.push(vec!["dim".to_string(), format!("{}x{}", w, h)]) + } + Nip96Response::UploadResult(Json(Nip96UploadResult { status: "success".to_string(), nip94_event: Some(Nip94Event { - tags: vec![ - vec![ - "url".to_string(), - format!("{}/{}", &settings.public_url, &hex_id), - ], - vec!["x".to_string(), hex_id], - vec!["m".to_string(), file_upload.mime_type], - ], + tags, }), ..Default::default() }))