Kieran/void-cat-rs
2
0
Fork 1
Code Issues 4 Pull Requests 1 Packages Projects Releases Wiki Activity

Compute blurhash
All checks were successful
continuous-integration/drone/push Build is passing
Details

Browse Source
This commit is contained in:
kieran 2024-05-13 12:13:30 +01:00
parent f60b510cc0
commit da2448fae9
Signed by: Kieran
GPG Key ID: DE71CEB3925BE941
9 changed files with 521 additions and 287 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
Cargo.lock generated
View File

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

1
Cargo.toml
View File

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

63
src/filesystem.rs
View File

@ -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<u8>,
pub size: u64,
pub mime_type: String,
pub width: Option<usize>,
pub height: Option<usize>,
pub blur_hash: Option<String>,
}
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,
})
}
}

42
src/processing/blurhash.rs Normal file
View File

@ -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<String, Error> {
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)
}

214
src/processing/image.rs
View File

@ -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<PathBuf, Error> {
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)
}
}
}

57
src/processing/mod.rs
View File

@ -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<PathBuf, Error>;
fn process_file(&mut self, in_file: PathBuf, mime_type: &str) -> Result<FileProcessorResult, Error>;
}
pub(crate) struct MediaProcessor {
processors: HashMap<String, Box<dyn FileProcessor + Sync + Send>>,
}
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<PathBuf, Error> {
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<FileProcessorResult, Error> {
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)
}
}
}

373
src/processing/webp.rs Normal file
View File

@ -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<usize, *mut AVCodecContext>,
decoders: HashMap<usize, *mut AVCodecContext>,
scalers: HashMap<usize, *mut SwsContext>,
stream_map: HashMap<usize, usize>,
width: Option<usize>,
height: Option<usize>,
blur_hash: Option<String>,
}
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<i32, Error> {
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<FileProcessorResult, Error> {
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()
},
}))
}
}
}

22
src/routes/blossom.rs
View File

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

29
src/routes/nip96.rs
View File

@ -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()
}))