From 510cffb305f17088059744696e21d662415fa8dd Mon Sep 17 00:00:00 2001
From: 21pages <pages21@163.com>
Date: Mon, 7 Aug 2023 15:18:34 +0800
Subject: [PATCH] av1 record, set zero codec private

Signed-off-by: 21pages <pages21@163.com>
---
 Cargo.lock                      | 10 ++++------
 libs/scrap/Cargo.toml           |  2 +-
 libs/scrap/src/common/record.rs | 29 +++++++++++++++++++++++++----
 src/server/video_service.rs     |  6 ++----
 4 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 519476e93..b606cdb21 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -6712,18 +6712,16 @@ dependencies = [
 
 [[package]]
 name = "webm"
-version = "1.0.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ecb047148a12ef1fd8ab26302bca7e82036f005c3073b48e17cc1b44ec577136"
+version = "1.1.0"
+source = "git+https://github.com/21pages/rust-webm#d2c4d3ac133c7b0e4c0f656da710b48391981e64"
 dependencies = [
  "webm-sys",
 ]
 
 [[package]]
 name = "webm-sys"
-version = "1.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0ded6ec82ccf51fe265b0b2b1579cac839574ed910c17baac58e807f8a9de7f3"
+version = "1.0.4"
+source = "git+https://github.com/21pages/rust-webm#d2c4d3ac133c7b0e4c0f656da710b48391981e64"
 dependencies = [
  "cc",
 ]
diff --git a/libs/scrap/Cargo.toml b/libs/scrap/Cargo.toml
index 4956403a9..d8b176597 100644
--- a/libs/scrap/Cargo.toml
+++ b/libs/scrap/Cargo.toml
@@ -19,7 +19,7 @@ cfg-if = "1.0"
 num_cpus = "1.15"
 lazy_static = "1.4"
 hbb_common = { path = "../hbb_common" }
-webm = "1.0"
+webm = { git = "https://github.com/21pages/rust-webm" }
 
 [dependencies.winapi]
 version = "0.3"
diff --git a/libs/scrap/src/common/record.rs b/libs/scrap/src/common/record.rs
index 9de70ae14..09d5efd3f 100644
--- a/libs/scrap/src/common/record.rs
+++ b/libs/scrap/src/common/record.rs
@@ -51,7 +51,10 @@ impl RecorderContext {
             + &self.id.clone()
             + &chrono::Local::now().format("_%Y%m%d%H%M%S_").to_string()
             + &self.format.to_string()
-            + if self.format == CodecFormat::VP9 || self.format == CodecFormat::VP8 {
+            + if self.format == CodecFormat::VP9
+                || self.format == CodecFormat::VP8
+                || self.format == CodecFormat::AV1
+            {
                 ".webm"
             } else {
                 ".mp4"
@@ -103,7 +106,7 @@ impl Recorder {
     pub fn new(mut ctx: RecorderContext) -> ResultType<Self> {
         ctx.set_filename()?;
         let recorder = match ctx.format {
-            CodecFormat::VP8 | CodecFormat::VP9 => Recorder {
+            CodecFormat::VP8 | CodecFormat::VP9 | CodecFormat::AV1 => Recorder {
                 inner: Box::new(WebmRecorder::new(ctx.clone())?),
                 ctx,
             },
@@ -122,7 +125,9 @@ impl Recorder {
     fn change(&mut self, mut ctx: RecorderContext) -> ResultType<()> {
         ctx.set_filename()?;
         self.inner = match ctx.format {
-            CodecFormat::VP8 | CodecFormat::VP9 => Box::new(WebmRecorder::new(ctx.clone())?),
+            CodecFormat::VP8 | CodecFormat::VP9 | CodecFormat::AV1 => {
+                Box::new(WebmRecorder::new(ctx.clone())?)
+            }
             #[cfg(feature = "hwcodec")]
             _ => Box::new(HwRecorder::new(ctx.clone())?),
             #[cfg(not(feature = "hwcodec"))]
@@ -161,6 +166,15 @@ impl Recorder {
                 }
                 vp9s.frames.iter().map(|f| self.write_video(f)).count();
             }
+            video_frame::Union::Av1s(av1s) => {
+                if self.ctx.format != CodecFormat::AV1 {
+                    self.change(RecorderContext {
+                        format: CodecFormat::AV1,
+                        ..self.ctx.clone()
+                    })?;
+                }
+                av1s.frames.iter().map(|f| self.write_video(f)).count();
+            }
             #[cfg(feature = "hwcodec")]
             video_frame::Union::H264s(h264s) => {
                 if self.ctx.format != CodecFormat::H264 {
@@ -227,10 +241,17 @@ impl RecorderApi for WebmRecorder {
             None,
             if ctx.format == CodecFormat::VP9 {
                 mux::VideoCodecId::VP9
-            } else {
+            } else if ctx.format == CodecFormat::VP8 {
                 mux::VideoCodecId::VP8
+            } else {
+                mux::VideoCodecId::AV1
             },
         );
+        if ctx.format == CodecFormat::AV1 {
+            // [129, 8, 12, 0] in 3.6.0, but zero works
+            let codec_private = vec![0, 0, 0, 0];
+            webm.set_codec_private(vt.track_number(), &codec_private);
+        }
         Ok(WebmRecorder {
             vt,
             webm: Some(webm),
diff --git a/src/server/video_service.rs b/src/server/video_service.rs
index 074b041c9..c2a9b6113 100644
--- a/src/server/video_service.rs
+++ b/src/server/video_service.rs
@@ -436,8 +436,7 @@ fn run(vs: VideoService) -> ResultType<()> {
     log::info!("init quality={:?}, abr enabled:{}", quality, abr);
     let codec_name = Encoder::negotiated_codec();
     let recorder = get_recorder(c.width, c.height, &codec_name);
-    let last_recording =
-        (recorder.lock().unwrap().is_some() || video_qos.record()) && codec_name != CodecName::AV1;
+    let last_recording = recorder.lock().unwrap().is_some() || video_qos.record();
     drop(video_qos);
     let encoder_cfg = get_encoder_config(&c, quality, last_recording);
 
@@ -479,8 +478,7 @@ fn run(vs: VideoService) -> ResultType<()> {
             allow_err!(encoder.set_quality(quality));
             video_qos.store_bitrate(encoder.bitrate());
         }
-        let recording = (recorder.lock().unwrap().is_some() || video_qos.record())
-            && codec_name != CodecName::AV1;
+        let recording = recorder.lock().unwrap().is_some() || video_qos.record();
         if recording != last_recording {
             bail!("SWITCH");
         }