init

src/decode.rs

@@ -0,0 +1,129 @@
+use std::collections::HashMap;
+use std::ffi::CStr;
+use std::ptr;
+
+use anyhow::Error;
+use ffmpeg_sys_the_third::AVPictureType::AV_PICTURE_TYPE_NONE;
+use ffmpeg_sys_the_third::{
+    av_buffer_alloc, av_frame_alloc, avcodec_alloc_context3, avcodec_find_decoder,
+    avcodec_free_context, avcodec_get_name, avcodec_open2, avcodec_parameters_to_context,
+    avcodec_receive_frame, avcodec_send_packet, AVCodec, AVCodecContext, AVFrame, AVMediaType,
+    AVPacket, AVStream, AVERROR, AVERROR_EOF,
+};
+use libc::memcpy;
+
+struct CodecContext {
+    pub context: *mut AVCodecContext,
+    pub codec: *const AVCodec,
+}
+
+impl Drop for CodecContext {
+    fn drop(&mut self) {
+        unsafe {
+            avcodec_free_context(&mut self.context);
+            self.codec = ptr::null_mut();
+            self.context = ptr::null_mut();
+        }
+    }
+}
+
+pub struct Decoder {
+    codecs: HashMap<i32, CodecContext>,
+    pts: i64,
+}
+
+unsafe impl Send for Decoder {}
+
+unsafe impl Sync for Decoder {}
+
+impl Decoder {
+    pub fn new() -> Self {
+        Self {
+            codecs: HashMap::new(),
+            pts: 0,
+        }
+    }
+
+    pub unsafe fn decode_pkt(
+        &mut self,
+        pkt: *mut AVPacket,
+        stream: *mut AVStream,
+    ) -> Result<Vec<(*mut AVFrame, *mut AVStream)>, Error> {
+        let stream_index = (*pkt).stream_index;
+        assert_eq!(
+            stream_index,
+            (*stream).index,
+            "Passed stream reference does not match stream_index of packet"
+        );
+
+        let codec_par = (*stream).codecpar;
+        assert_ne!(
+            codec_par,
+            ptr::null_mut(),
+            "Codec parameters are missing from stream"
+        );
+
+        if let std::collections::hash_map::Entry::Vacant(e) = self.codecs.entry(stream_index) {
+            let codec = avcodec_find_decoder((*codec_par).codec_id);
+            if codec.is_null() {
+                return Err(Error::msg(format!(
+                    "Failed to find codec: {}",
+                    CStr::from_ptr(avcodec_get_name((*codec_par).codec_id)).to_str()?
+                )));
+            }
+            let context = avcodec_alloc_context3(ptr::null());
+            if context.is_null() {
+                return Err(Error::msg("Failed to alloc context"));
+            }
+            if avcodec_parameters_to_context(context, (*stream).codecpar) != 0 {
+                return Err(Error::msg("Failed to copy codec parameters to context"));
+            }
+            if avcodec_open2(context, codec, ptr::null_mut()) < 0 {
+                return Err(Error::msg("Failed to open codec"));
+            }
+            e.insert(CodecContext { context, codec });
+        }
+
+        if let Some(ctx) = self.codecs.get_mut(&stream_index) {
+            // subtitles don't need decoding, create a frame from the pkt data
+            if (*ctx.codec).type_ == AVMediaType::AVMEDIA_TYPE_SUBTITLE {
+                let frame = av_frame_alloc();
+                (*frame).pts = (*pkt).pts;
+                (*frame).pkt_dts = (*pkt).dts;
+                (*frame).duration = (*pkt).duration;
+                (*frame).buf[0] = av_buffer_alloc((*pkt).size as usize);
+                (*frame).data[0] = (*(*frame).buf[0]).data;
+                (*frame).linesize[0] = (*pkt).size;
+                memcpy(
+                    (*frame).data[0] as *mut libc::c_void,
+                    (*pkt).data as *const libc::c_void,
+                    (*pkt).size as usize,
+                );
+                return Ok(vec![(frame, stream)]);
+            }
+
+            let mut ret = avcodec_send_packet(ctx.context, pkt);
+            if ret < 0 {
+                return Err(Error::msg(format!("Failed to decode packet {}", ret)));
+            }
+
+            let mut pkgs = Vec::new();
+            while ret >= 0 {
+                let frame = av_frame_alloc();
+                ret = avcodec_receive_frame(ctx.context, frame);
+                if ret < 0 {
+                    if ret == AVERROR_EOF || ret == AVERROR(libc::EAGAIN) {
+                        break;
+                    }
+                    return Err(Error::msg(format!("Failed to decode {}", ret)));
+                }
+
+                (*frame).pict_type = AV_PICTURE_TYPE_NONE; // encoder prints warnings
+                pkgs.push((frame, stream));
+            }
+            Ok(pkgs)
+        } else {
+            Ok(vec![])
+        }
+    }
+}

src/demux.rs

@@ -0,0 +1,318 @@
+use std::ffi::CStr;
+use std::time::Instant;
+use std::{ptr, slice};
+
+use anyhow::Error;
+use ffmpeg_sys_the_third::*;
+
+use crate::get_ffmpeg_error_msg;
+use crate::return_ffmpeg_error;
+use std::fmt::{Display, Formatter};
+use std::io::{Read, Write};
+use std::mem::transmute;
+
+unsafe extern "C" fn read_data(
+    opaque: *mut libc::c_void,
+    dst_buffer: *mut libc::c_uchar,
+    size: libc::c_int,
+) -> libc::c_int {
+    let buffer: *mut Box<dyn Read> = opaque.cast();
+    /// we loop until there is enough data to fill [size]
+    let mut dst_slice: &mut [u8] = slice::from_raw_parts_mut(dst_buffer, size as usize);
+    let mut w_total = 0usize;
+    loop {
+        return match (*buffer).read(dst_slice) {
+            Ok(v) => {
+                w_total += v;
+                if w_total != size as usize {
+                    dst_slice = &mut dst_slice[v..];
+                    continue;
+                }
+                size
+            }
+            Err(e) => {
+                eprintln!("read_data {}", e);
+                AVERROR_EOF
+            }
+        };
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct DemuxerInfo {
+    pub bitrate: usize,
+    pub duration: f32,
+    pub channels: Vec<StreamInfoChannel>,
+}
+
+unsafe impl Send for DemuxerInfo {}
+unsafe impl Sync for DemuxerInfo {}
+
+impl DemuxerInfo {
+    pub fn best_stream(&self, t: StreamChannelType) -> Option<&StreamInfoChannel> {
+        self.channels
+            .iter()
+            .filter(|a| a.channel_type == t)
+            .reduce(|acc, channel| {
+                if channel.best_metric() > acc.best_metric() {
+                    channel
+                } else {
+                    acc
+                }
+            })
+    }
+
+    pub fn best_video(&self) -> Option<&StreamInfoChannel> {
+        self.best_stream(StreamChannelType::Video)
+    }
+
+    pub fn best_audio(&self) -> Option<&StreamInfoChannel> {
+        self.best_stream(StreamChannelType::Audio)
+    }
+
+    pub fn best_subtitle(&self) -> Option<&StreamInfoChannel> {
+        self.best_stream(StreamChannelType::Subtitle)
+    }
+
+    pub unsafe fn is_best_stream(&self, stream: *mut AVStream) -> bool {
+        match (*(*stream).codecpar).codec_type {
+            AVMediaType::AVMEDIA_TYPE_VIDEO => {
+                (*stream).index == self.best_video().map_or(usize::MAX, |r| r.index) as libc::c_int
+            }
+            AVMediaType::AVMEDIA_TYPE_AUDIO => {
+                (*stream).index == self.best_audio().map_or(usize::MAX, |r| r.index) as libc::c_int
+            }
+            AVMediaType::AVMEDIA_TYPE_SUBTITLE => {
+                (*stream).index
+                    == self.best_subtitle().map_or(usize::MAX, |r| r.index) as libc::c_int
+            }
+            _ => false,
+        }
+    }
+}
+
+impl Display for DemuxerInfo {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Demuxer Info:")?;
+        for c in &self.channels {
+            write!(f, "\n{}", c)?;
+        }
+        Ok(())
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub enum StreamChannelType {
+    Video,
+    Audio,
+    Subtitle,
+}
+
+impl Display for StreamChannelType {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "{}",
+            match self {
+                StreamChannelType::Video => "video",
+                StreamChannelType::Audio => "audio",
+                StreamChannelType::Subtitle => "subtitle",
+            }
+        )
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct StreamInfoChannel {
+    pub index: usize,
+    pub channel_type: StreamChannelType,
+    pub codec: usize,
+    pub width: usize,
+    pub height: usize,
+    pub fps: f32,
+    pub sample_rate: usize,
+    pub format: usize,
+}
+
+impl StreamInfoChannel {
+    pub fn best_metric(&self) -> f32 {
+        match self.channel_type {
+            StreamChannelType::Video => self.width as f32 * self.height as f32 * self.fps,
+            StreamChannelType::Audio => self.sample_rate as f32,
+            StreamChannelType::Subtitle => 999. - self.index as f32,
+        }
+    }
+}
+
+impl Display for StreamInfoChannel {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        let codec_name = unsafe { CStr::from_ptr(avcodec_get_name(transmute(self.codec as i32))) };
+        write!(
+            f,
+            "{} #{}: codec={},size={}x{},fps={}",
+            self.channel_type,
+            self.index,
+            codec_name.to_str().unwrap(),
+            self.width,
+            self.height,
+            self.fps
+        )
+    }
+}
+
+pub struct Demuxer<'a> {
+    ctx: *mut AVFormatContext,
+    input: String,
+    started: Instant,
+    buffer: Option<Box<dyn Read + 'a>>,
+}
+
+unsafe impl Send for Demuxer<'_> {}
+
+unsafe impl Sync for Demuxer<'_> {}
+
+impl Demuxer<'_> {
+    pub fn new(input: &str) -> Self {
+        unsafe {
+            let ps = avformat_alloc_context();
+            Self {
+                ctx: ps,
+                input: input.to_string(),
+                started: Instant::now(),
+                buffer: None,
+            }
+        }
+    }
+
+    pub fn new_custom_io<T: Read + 'static>(reader: T) -> Self {
+        unsafe {
+            let ps = avformat_alloc_context();
+            (*ps).flags |= AVFMT_FLAG_CUSTOM_IO;
+
+            let buffer = Box::new(reader);
+            Self {
+                ctx: ps,
+                input: String::new(),
+                started: Instant::now(),
+                buffer: Some(buffer),
+            }
+        }
+    }
+
+    unsafe fn open_input(&mut self) -> libc::c_int {
+        if let Some(mut buffer) = self.buffer.take() {
+            const BUFFER_SIZE: usize = 4096;
+            let pb = avio_alloc_context(
+                av_mallocz(BUFFER_SIZE) as *mut libc::c_uchar,
+                BUFFER_SIZE as libc::c_int,
+                0,
+                ptr::addr_of_mut!(buffer) as _,
+                Some(read_data),
+                None,
+                None,
+            );
+            (*self.ctx).pb = pb;
+            avformat_open_input(
+                &mut self.ctx,
+                ptr::null_mut(),
+                ptr::null_mut(),
+                ptr::null_mut(),
+            )
+        } else {
+            avformat_open_input(
+                &mut self.ctx,
+                format!("{}\0", self.input).as_ptr() as *const libc::c_char,
+                ptr::null_mut(),
+                ptr::null_mut(),
+            )
+        }
+    }
+
+    pub unsafe fn probe_input(&mut self) -> Result<DemuxerInfo, Error> {
+        let ret = self.open_input();
+        return_ffmpeg_error!(ret);
+
+        if avformat_find_stream_info(self.ctx, ptr::null_mut()) < 0 {
+            return Err(Error::msg("Could not find stream info"));
+        }
+        av_dump_format(self.ctx, 0, ptr::null_mut(), 0);
+
+        let mut channel_infos = vec![];
+
+        for n in 0..(*self.ctx).nb_streams as usize {
+            let stream = *(*self.ctx).streams.add(n);
+            match (*(*stream).codecpar).codec_type {
+                AVMediaType::AVMEDIA_TYPE_VIDEO => {
+                    channel_infos.push(StreamInfoChannel {
+                        index: (*stream).index as usize,
+                        codec: (*(*stream).codecpar).codec_id as usize,
+                        channel_type: StreamChannelType::Video,
+                        width: (*(*stream).codecpar).width as usize,
+                        height: (*(*stream).codecpar).height as usize,
+                        fps: av_q2d((*stream).avg_frame_rate) as f32,
+                        format: (*(*stream).codecpar).format as usize,
+                        sample_rate: 0,
+                    });
+                }
+                AVMediaType::AVMEDIA_TYPE_AUDIO => {
+                    channel_infos.push(StreamInfoChannel {
+                        index: (*stream).index as usize,
+                        codec: (*(*stream).codecpar).codec_id as usize,
+                        channel_type: StreamChannelType::Audio,
+                        width: (*(*stream).codecpar).width as usize,
+                        height: (*(*stream).codecpar).height as usize,
+                        fps: 0.0,
+                        format: (*(*stream).codecpar).format as usize,
+                        sample_rate: (*(*stream).codecpar).sample_rate as usize,
+                    });
+                }
+                AVMediaType::AVMEDIA_TYPE_SUBTITLE => {
+                    channel_infos.push(StreamInfoChannel {
+                        index: (*stream).index as usize,
+                        codec: (*(*stream).codecpar).codec_id as usize,
+                        channel_type: StreamChannelType::Subtitle,
+                        width: 0,
+                        height: 0,
+                        fps: 0.0,
+                        format: 0,
+                        sample_rate: 0,
+                    });
+                }
+                AVMediaType::AVMEDIA_TYPE_ATTACHMENT => {}
+                AVMediaType::AVMEDIA_TYPE_NB => {}
+                _ => {}
+            }
+        }
+
+        let info = DemuxerInfo {
+            duration: (*self.ctx).duration as f32 / AV_TIME_BASE as f32,
+            bitrate: (*self.ctx).bit_rate as usize,
+            channels: channel_infos,
+        };
+        Ok(info)
+    }
+
+    pub unsafe fn get_packet(&mut self) -> Result<(*mut AVPacket, *mut AVStream), Error> {
+        let pkt: *mut AVPacket = av_packet_alloc();
+        let ret = av_read_frame(self.ctx, pkt);
+        if ret == AVERROR_EOF {
+            return Ok((ptr::null_mut(), ptr::null_mut()));
+        }
+        if ret < 0 {
+            let msg = get_ffmpeg_error_msg(ret);
+            return Err(Error::msg(msg));
+        }
+        let stream = *(*self.ctx).streams.add((*pkt).stream_index as usize);
+        let pkg = (pkt, stream);
+        Ok(pkg)
+    }
+}
+
+impl Drop for Demuxer<'_> {
+    fn drop(&mut self) {
+        unsafe {
+            avformat_free_context(self.ctx);
+            self.ctx = ptr::null_mut();
+        }
+    }
+}

src/lib.rs

@@ -0,0 +1,25 @@
+use ffmpeg_sys_the_third::av_make_error_string;
+use std::ffi::CStr;
+
+mod decode;
+mod demux;
+mod resample;
+mod scale;
+
+#[macro_export(crate)]
+macro_rules! return_ffmpeg_error {
+    ($x:expr) => {
+        if $x < 0 {
+            return Err(Error::msg(get_ffmpeg_error_msg($x)));
+        }
+    };
+}
+
+fn get_ffmpeg_error_msg(ret: libc::c_int) -> String {
+    unsafe {
+        const BUF_SIZE: usize = 512;
+        let mut buf: [libc::c_char; BUF_SIZE] = [0; BUF_SIZE];
+        av_make_error_string(buf.as_mut_ptr(), BUF_SIZE, ret);
+        String::from(CStr::from_ptr(buf.as_ptr()).to_str().unwrap())
+    }
+}

src/main.rs

@@ -0,0 +1,3 @@
+fn main() {
+    println!("Hello, world!");
+}

src/resample.rs

@@ -0,0 +1,75 @@
+use crate::get_ffmpeg_error_msg;
+use crate::return_ffmpeg_error;
+use anyhow::Error;
+use ffmpeg_sys_the_third::{
+    av_channel_layout_default, av_frame_alloc, av_frame_copy_props, av_frame_free,
+    swr_alloc_set_opts2, swr_convert_frame, swr_init, AVChannelLayout, AVFrame, AVSampleFormat,
+    SwrContext,
+};
+use libc::malloc;
+use std::mem::transmute;
+use std::ptr;
+
+pub struct Resample {
+    format: AVSampleFormat,
+    sample_rate: u32,
+    channels: usize,
+    ctx: *mut SwrContext,
+}
+
+impl Resample {
+    pub fn new(format: AVSampleFormat, rate: u32, channels: usize) -> Self {
+        Self {
+            format,
+            channels,
+            sample_rate: rate,
+            ctx: ptr::null_mut(),
+        }
+    }
+
+    unsafe fn setup_swr(&mut self, frame: *mut AVFrame) -> Result<(), Error> {
+        if !self.ctx.is_null() {
+            return Ok(());
+        }
+        let layout = malloc(size_of::<AVChannelLayout>()) as *mut AVChannelLayout;
+        av_channel_layout_default(layout, self.channels as libc::c_int);
+
+        let ret = swr_alloc_set_opts2(
+            &mut self.ctx,
+            layout,
+            self.format,
+            self.sample_rate as libc::c_int,
+            &(*frame).ch_layout,
+            transmute((*frame).format),
+            (*frame).sample_rate,
+            0,
+            ptr::null_mut(),
+        );
+        return_ffmpeg_error!(ret);
+
+        let ret = swr_init(self.ctx);
+        return_ffmpeg_error!(ret);
+
+        Ok(())
+    }
+
+    pub unsafe fn process_frame(&mut self, frame: *mut AVFrame) -> Result<*mut AVFrame, Error> {
+        self.setup_swr(frame)?;
+
+        let mut out_frame = av_frame_alloc();
+        av_frame_copy_props(out_frame, frame);
+        (*out_frame).sample_rate = self.sample_rate as libc::c_int;
+        (*out_frame).format = transmute(self.format);
+        (*out_frame).time_base = (*frame).time_base;
+
+        av_channel_layout_default(&mut (*out_frame).ch_layout, self.channels as libc::c_int);
+
+        let ret = swr_convert_frame(self.ctx, out_frame, frame);
+        if ret < 0 {
+            av_frame_free(&mut out_frame);
+            return Err(Error::msg(get_ffmpeg_error_msg(ret)));
+        }
+
+        Ok(out_frame)
+    }
+}

src/scale.rs

@@ -0,0 +1,95 @@
+use crate::get_ffmpeg_error_msg;
+use std::mem::transmute;
+use std::ptr;
+
+use crate::return_ffmpeg_error;
+use anyhow::Error;
+use ffmpeg_sys_the_third::{
+    av_frame_alloc, av_frame_copy_props, sws_freeContext, sws_getContext, sws_scale_frame, AVFrame,
+    AVPixelFormat, SwsContext, SWS_BILINEAR,
+};
+
+pub struct Scaler {
+    width: u16,
+    height: u16,
+    format: AVPixelFormat,
+    ctx: *mut SwsContext,
+}
+
+unsafe impl Send for Scaler {}
+
+unsafe impl Sync for Scaler {}
+
+impl Drop for Scaler {
+    fn drop(&mut self) {
+        unsafe {
+            sws_freeContext(self.ctx);
+            self.ctx = ptr::null_mut();
+        }
+    }
+}
+
+impl Scaler {
+    pub fn new(format: AVPixelFormat) -> Self {
+        Self {
+            width: 0,
+            height: 0,
+            format,
+            ctx: ptr::null_mut(),
+        }
+    }
+
+    unsafe fn setup_scaler(
+        &mut self,
+        frame: *const AVFrame,
+        width: u16,
+        height: u16,
+    ) -> Result<(), Error> {
+        if !self.ctx.is_null() && self.width == width && self.height == height {
+            return Ok(());
+        }
+
+        // clear previous context, before re-creating
+        if !self.ctx.is_null() {
+            sws_freeContext(self.ctx);
+            self.ctx = ptr::null_mut();
+        }
+
+        self.width = width;
+        self.height = height;
+        self.ctx = sws_getContext(
+            (*frame).width,
+            (*frame).height,
+            transmute((*frame).format),
+            self.width as libc::c_int,
+            self.height as libc::c_int,
+            transmute(self.format),
+            SWS_BILINEAR,
+            ptr::null_mut(),
+            ptr::null_mut(),
+            ptr::null_mut(),
+        );
+        if self.ctx.is_null() {
+            return Err(Error::msg("Failed to create scalar context"));
+        }
+        Ok(())
+    }
+
+    pub unsafe fn process_frame(
+        &mut self,
+        frame: *mut AVFrame,
+        width: u16,
+        height: u16,
+    ) -> Result<*mut AVFrame, Error> {
+        self.setup_scaler(frame, width, height)?;
+
+        let dst_frame = av_frame_alloc();
+        let ret = av_frame_copy_props(dst_frame, frame);
+        return_ffmpeg_error!(ret);
+
+        let ret = sws_scale_frame(self.ctx, dst_frame, frame);
+        return_ffmpeg_error!(ret);
+
+        Ok(dst_frame)
+    }
+}