From 3298ee3f3cad74dc407fe3ee76b4180fbd3169a3 Mon Sep 17 00:00:00 2001
From: Scott Lamb <slamb@slamb.org>
Date: Wed, 11 Dec 2024 14:56:44 -0800
Subject: [PATCH] strip Annex B prefixes from `sprop-parameter-sets`

---
 CHANGELOG.md      |  5 +++++
 src/codec/h264.rs | 57 ++++++++++++++++++++++++++++++++++++++---------
 2 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c7444bd..42fa87d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,8 @@
+## unreleased
+
+*   improve interoperability with Tenda CP3 Pro cameras by stripping an
+    Annex B separator prefix (if any) from each NAL in the H.264 `sprop-parameter-sets`.
+
 ## `v0.4.10` (2024-08-19)
 
 *   update `base64` dep to 0.22
diff --git a/src/codec/h264.rs b/src/codec/h264.rs
index 38974a6..b7e56ce 100644
--- a/src/codec/h264.rs
+++ b/src/codec/h264.rs
@@ -725,31 +725,40 @@ impl InternalParameters {
             let nal = base64::engine::general_purpose::STANDARD
                 .decode(nal)
                 .map_err(|_| {
-                    "bad sprop-parameter-sets: NAL has invalid base64 encoding".to_string()
+                    format!("bad sprop-parameter-sets: invalid base64 encoding in NAL: {nal}")
                 })?;
-            if nal.is_empty() {
+
+            // The parameter set should just be the NAL with no Annex B prefix, but at least the
+            // Tenda CP3 Pro gets this wrong. Strip out the prefix.
+            let nal = strip_annexb_prefix(&nal).map(<[u8]>::to_vec).unwrap_or(nal);
+            let hex = crate::hex::LimitedHex::new(&nal, 256);
+            let Some(&header) = nal.first() else {
                 return Err("bad sprop-parameter-sets: empty NAL".into());
-            }
-            let header = h264_reader::nal::NalHeader::new(nal[0])
-                .map_err(|_| format!("bad sprop-parameter-sets: bad NAL header {:0x}", nal[0]))?;
+            };
+            let header = h264_reader::nal::NalHeader::new(header)
+                .map_err(|_| format!("bad sprop-parameter-sets: bad header in NAL {hex}"))?;
             match header.nal_unit_type() {
                 UnitType::SeqParameterSet => {
                     if sps_nal.is_some() {
-                        return Err("multiple SPSs".into());
+                        return Err("multiple SPSs are currently unsupported".into());
                     }
                     sps_nal = Some(nal);
                 }
                 UnitType::PicParameterSet => {
                     if pps_nal.is_some() {
-                        return Err("multiple PPSs".into());
+                        return Err("multiple PPSs are currently unsupported".into());
                     }
                     pps_nal = Some(nal);
                 }
-                _ => return Err("only SPS and PPS expected in parameter sets".into()),
+                _ => {
+                    return Err(format!(
+                        "bad sprop-parameter-sets: unexpected non-SPS/PPS NAL: {hex}"
+                    ));
+                }
             }
         }
-        let sps_nal = sps_nal.ok_or_else(|| "no sps".to_string())?;
-        let pps_nal = pps_nal.ok_or_else(|| "no pps".to_string())?;
+        let sps_nal = sps_nal.ok_or_else(|| "bad sprop-parameter-sets: no sps".to_string())?;
+        let pps_nal = pps_nal.ok_or_else(|| "bad sprop-parameter-sets: no pps".to_string())?;
         Self::parse_sps_and_pps(&sps_nal, &pps_nal, false)
     }
 
@@ -901,6 +910,22 @@ fn to_bytes(hdr: NalHeader, len: u32, pieces: &[Bytes]) -> Bytes {
     out.into()
 }
 
+/// Returns a subslice with an Annex B separator prefix removed.
+///
+/// Annex B separators are two or more zero bytes followed by a one.
+fn strip_annexb_prefix(nal: &[u8]) -> Option<&[u8]> {
+    let mut post_zeros = match nal {
+        // Annex B separator starts with 2 zeros.
+        [0, 0, rest @ ..] => rest,
+        _ => return None,
+    };
+    while let [0, rest @ ..] = post_zeros {
+        // ...skip over any additional zeros also.
+        post_zeros = rest;
+    }
+    post_zeros.strip_prefix(&b"\x01"[..])
+}
+
 /// A simple packetizer, currently only for testing/benchmarking. Unstable.
 ///
 /// Only uses plain NALs and FU-As, never STAP-A.
@@ -1565,6 +1590,18 @@ mod tests {
         .unwrap_err();
     }
 
+    #[test]
+    fn cp3pro_params() {
+        init_logging();
+        let params = super::InternalParameters::parse_format_specific_params(
+            "profile-level-id=4DE028;\
+             packetization-mode=1;\
+             sprop-parameter-sets=AAAAAWdNABaNjUBQF/yzcBAQFAAALuAABX5AEA==,AAAAAWjuOIA=",
+        )
+        .unwrap();
+        assert_eq!(&params.generic_parameters.rfc6381_codec, "avc1.4D0016");
+    }
+
     #[test]
     fn bad_format_specific_params() {
         init_logging();