video record

Signed-off-by: 21pages <pages21@163.com>

libs/hbb_common/Cargo.toml

@@ -30,6 +30,7 @@ filetime = "0.2"
 sodiumoxide = "0.2"
 regex = "1.4"
 tokio-socks = { git = "https://github.com/open-trade/tokio-socks" }
+chrono = "0.4"
 
 [target.'cfg(not(any(target_os = "android", target_os = "ios")))'.dependencies]
 mac_address = "1.1"

libs/hbb_common/src/lib.rs

@@ -38,6 +38,8 @@ pub use tokio_socks;
 pub use tokio_socks::IntoTargetAddr;
 pub use tokio_socks::TargetAddr;
 pub mod password_security;
+pub use chrono;
+pub use directories_next;
 
 #[cfg(feature = "quic")]
 pub type Stream = quic::Connection;

libs/scrap/Cargo.toml

@@ -20,6 +20,7 @@ libc = "0.2"
 num_cpus = "1.13"
 lazy_static = "1.4"
 hbb_common = { path = "../hbb_common" }
+webm = "1.0"
 
 [dependencies.winapi]
 version = "0.3"
@@ -37,7 +38,6 @@ ndk = { version = "0.7", features = ["media"], optional = true}
 [target.'cfg(not(target_os = "android"))'.dev-dependencies]
 repng = "0.2"
 docopt = "1.1"
-webm = "1.0"
 serde = {version="1.0", features=["derive"]}
 quest = "0.3"
 

libs/scrap/src/common/hwcodec.rs

@@ -28,7 +28,7 @@ const CFG_KEY_ENCODER: &str = "bestHwEncoders";
 const CFG_KEY_DECODER: &str = "bestHwDecoders";
 
 const DEFAULT_PIXFMT: AVPixelFormat = AVPixelFormat::AV_PIX_FMT_YUV420P;
-const DEFAULT_TIME_BASE: [i32; 2] = [1, 30];
+pub const DEFAULT_TIME_BASE: [i32; 2] = [1, 30];
 const DEFAULT_GOP: i32 = 60;
 const DEFAULT_HW_QUALITY: Quality = Quality_Default;
 const DEFAULT_RC: RateContorl = RC_DEFAULT;

libs/scrap/src/common/mod.rs

@@ -39,6 +39,7 @@ pub use self::convert::*;
 pub const STRIDE_ALIGN: usize = 64; // commonly used in libvpx vpx_img_alloc caller
 pub const HW_STRIDE_ALIGN: usize = 0; // recommended by av_frame_get_buffer
 
+pub mod record;
 mod vpx;
 
 #[inline]

libs/scrap/src/common/record.rs

@@ -0,0 +1,297 @@
+#[cfg(feature = "hwcodec")]
+use hbb_common::anyhow::anyhow;
+use hbb_common::{
+    bail, chrono,
+    config::Config,
+    directories_next,
+    message_proto::{message, video_frame, EncodedVideoFrame, Message},
+    ResultType,
+};
+#[cfg(feature = "hwcodec")]
+use hwcodec::mux::{MuxContext, Muxer};
+use std::{
+    fs::{File, OpenOptions},
+    io,
+    time::Instant,
+};
+use std::{
+    ops::{Deref, DerefMut},
+    path::PathBuf,
+};
+use webm::mux::{self, Segment, Track, VideoTrack, Writer};
+
+const MIN_SECS: u64 = 1;
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum RecodeCodecID {
+    VP9,
+    H264,
+    H265,
+}
+
+#[derive(Debug, Clone)]
+pub struct RecorderContext {
+    pub id: String,
+    pub filename: String,
+    pub width: usize,
+    pub height: usize,
+    pub codec_id: RecodeCodecID,
+}
+
+impl RecorderContext {
+    pub fn set_filename(&mut self) -> ResultType<()> {
+        let mut dir = Config::get_option("video-save-directory");
+        if !dir.is_empty() {
+            if !PathBuf::from(&dir).exists() {
+                std::fs::create_dir_all(&dir)?;
+            }
+        } else {
+            dir = Self::default_save_directory();
+            if !dir.is_empty() && !PathBuf::from(&dir).exists() {
+                std::fs::create_dir_all(&dir)?;
+            }
+        }
+        let file = self.id.clone()
+            + &chrono::Local::now().format("_%Y%m%d%H%M%S").to_string()
+            + if self.codec_id == RecodeCodecID::VP9 {
+                ".webm"
+            } else {
+                ".mp4"
+            };
+        self.filename = PathBuf::from(&dir).join(file).to_string_lossy().to_string();
+        Ok(())
+    }
+
+    pub fn default_save_directory() -> String {
+        if let Some(user) = directories_next::UserDirs::new() {
+            if let Some(video_dir) = user.video_dir() {
+                return video_dir.join("RustDesk").to_string_lossy().to_string();
+            }
+        }
+        "".to_owned()
+    }
+}
+
+unsafe impl Send for Recorder {}
+unsafe impl Sync for Recorder {}
+
+pub trait RecorderApi {
+    fn new(ctx: RecorderContext) -> ResultType<Self>
+    where
+        Self: Sized;
+    fn write_video(&mut self, frame: &EncodedVideoFrame) -> bool;
+}
+
+pub struct Recorder {
+    pub inner: Box<dyn RecorderApi>,
+    ctx: RecorderContext,
+}
+
+impl Deref for Recorder {
+    type Target = Box<dyn RecorderApi>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl DerefMut for Recorder {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner
+    }
+}
+
+impl Recorder {
+    pub fn new(mut ctx: RecorderContext) -> ResultType<Self> {
+        ctx.set_filename()?;
+        let recorder = match ctx.codec_id {
+            RecodeCodecID::VP9 => Recorder {
+                inner: Box::new(WebmRecorder::new(ctx.clone())?),
+                ctx,
+            },
+            #[cfg(feature = "hwcodec")]
+            _ => Recorder {
+                inner: Box::new(HwRecorder::new(ctx.clone())?),
+                ctx,
+            },
+            #[cfg(not(feature = "hwcodec"))]
+            _ => bail!("unsupported codec type"),
+        };
+        Ok(recorder)
+    }
+
+    fn change(&mut self, mut ctx: RecorderContext) -> ResultType<()> {
+        ctx.set_filename()?;
+        self.inner = match ctx.codec_id {
+            RecodeCodecID::VP9 => Box::new(WebmRecorder::new(ctx.clone())?),
+            #[cfg(feature = "hwcodec")]
+            _ => Box::new(HwRecorder::new(ctx.clone())?),
+            #[cfg(not(feature = "hwcodec"))]
+            _ => bail!("unsupported codec type"),
+        };
+        self.ctx = ctx;
+        Ok(())
+    }
+
+    pub fn write_message(&mut self, msg: &Message) {
+        if let Some(message::Union::VideoFrame(vf)) = &msg.union {
+            if let Some(frame) = &vf.union {
+                self.write_frame(frame).ok();
+            }
+        }
+    }
+
+    pub fn write_frame(&mut self, frame: &video_frame::Union) -> ResultType<()> {
+        match frame {
+            video_frame::Union::Vp9s(vp9s) => {
+                if self.ctx.codec_id != RecodeCodecID::VP9 {
+                    self.change(RecorderContext {
+                        codec_id: RecodeCodecID::VP9,
+                        ..self.ctx.clone()
+                    })?;
+                }
+                vp9s.frames.iter().map(|f| self.write_video(f)).count();
+            }
+            #[cfg(feature = "hwcodec")]
+            video_frame::Union::H264s(h264s) => {
+                if self.ctx.codec_id != RecodeCodecID::H264 {
+                    self.change(RecorderContext {
+                        codec_id: RecodeCodecID::H264,
+                        ..self.ctx.clone()
+                    })?;
+                }
+                if self.ctx.codec_id == RecodeCodecID::H264 {
+                    h264s.frames.last().map(|f| self.write_video(f));
+                }
+            }
+            #[cfg(feature = "hwcodec")]
+            video_frame::Union::H265s(h265s) => {
+                if self.ctx.codec_id != RecodeCodecID::H265 {
+                    self.change(RecorderContext {
+                        codec_id: RecodeCodecID::H265,
+                        ..self.ctx.clone()
+                    })?;
+                }
+                if self.ctx.codec_id == RecodeCodecID::H265 {
+                    h265s.frames.last().map(|f| self.write_video(f));
+                }
+            }
+            _ => bail!("unsupported frame type"),
+        }
+        Ok(())
+    }
+}
+
+struct WebmRecorder {
+    vt: VideoTrack,
+    webm: Option<Segment<Writer<File>>>,
+    ctx: RecorderContext,
+    key: bool,
+    written: bool,
+    start: Instant,
+}
+
+impl RecorderApi for WebmRecorder {
+    fn new(ctx: RecorderContext) -> ResultType<Self> {
+        let out = match {
+            OpenOptions::new()
+                .write(true)
+                .create_new(true)
+                .open(&ctx.filename)
+        } {
+            Ok(file) => file,
+            Err(ref e) if e.kind() == io::ErrorKind::AlreadyExists => File::create(&ctx.filename)?,
+            Err(e) => return Err(e.into()),
+        };
+        let mut webm = match mux::Segment::new(mux::Writer::new(out)) {
+            Some(v) => v,
+            None => bail!("Failed to create webm mux"),
+        };
+        let vt = webm.add_video_track(
+            ctx.width as _,
+            ctx.height as _,
+            None,
+            mux::VideoCodecId::VP9,
+        );
+        Ok(WebmRecorder {
+            vt,
+            webm: Some(webm),
+            ctx,
+            key: false,
+            written: false,
+            start: Instant::now(),
+        })
+    }
+
+    fn write_video(&mut self, frame: &EncodedVideoFrame) -> bool {
+        if frame.key {
+            self.key = true;
+        }
+        if self.key {
+            let ok = self
+                .vt
+                .add_frame(&frame.data, frame.pts as u64 * 1_000_000, frame.key);
+            if ok {
+                self.written = true;
+            }
+            ok
+        } else {
+            false
+        }
+    }
+}
+
+impl Drop for WebmRecorder {
+    fn drop(&mut self) {
+        std::mem::replace(&mut self.webm, None).map_or(false, |webm| webm.finalize(None));
+        if !self.written || self.start.elapsed().as_secs() < MIN_SECS {
+            std::fs::remove_file(&self.ctx.filename).ok();
+        }
+    }
+}
+
+#[cfg(feature = "hwcodec")]
+struct HwRecorder {
+    muxer: Muxer,
+    ctx: RecorderContext,
+    written: bool,
+    start: Instant,
+}
+
+#[cfg(feature = "hwcodec")]
+impl RecorderApi for HwRecorder {
+    fn new(ctx: RecorderContext) -> ResultType<Self> {
+        let muxer = Muxer::new(MuxContext {
+            filename: ctx.filename.clone(),
+            width: ctx.width,
+            height: ctx.height,
+            is265: ctx.codec_id == RecodeCodecID::H265,
+            framerate: crate::hwcodec::DEFAULT_TIME_BASE[1] as _,
+        })
+        .map_err(|_| anyhow!("Failed to create hardware muxer"))?;
+        Ok(HwRecorder {
+            muxer,
+            ctx,
+            written: false,
+            start: Instant::now(),
+        })
+    }
+
+    fn write_video(&mut self, frame: &EncodedVideoFrame) -> bool {
+        let ok = self.muxer.write_video(&frame.data, frame.pts).is_ok();
+        if ok {
+            self.written = true;
+        }
+        ok
+    }
+}
+
+#[cfg(feature = "hwcodec")]
+impl Drop for HwRecorder {
+    fn drop(&mut self) {
+        self.muxer.write_tail().ok();
+        if !self.written || self.start.elapsed().as_secs() < MIN_SECS {
+            std::fs::remove_file(&self.ctx.filename).ok();
+        }
+    }
+}