rtpav1pay: Fix Leb128Bytes size parsing

There are multiple ways of encoding the value, and don't assume
that bitstream used the way used in this plugin. Instead, count
the number of used bytes.

Fixes: https://gitlab.freedesktop.org/gstreamer/gst-plugins-rs/-/issues/312
Part-of: <https://gitlab.freedesktop.org/gstreamer/gst-plugins-rs/-/merge_requests/1090>

3 files changed, 22 insertions, 17 deletions

diff --git a/net/rtp/src/av1/common/integers.rs b/net/rtp/src/av1/common/integers.rs
index aebee7a2..02b5831e 100644
--- a/net/rtp/src/av1/common/integers.rs
+++ b/net/rtp/src/av1/common/integers.rs
@@ -12,23 +12,25 @@
 use bitstream_io::{BitRead, BitReader, BitWrite, BitWriter, Endianness};
 use std::io::{self, Read, Seek, Write};
 
-pub fn parse_leb128<R, E>(reader: &mut BitReader<R, E>) -> io::Result<u32>
+pub fn parse_leb128<R, E>(reader: &mut BitReader<R, E>) -> io::Result<(u32, u32)>
 where
     R: Read + Seek,
     E: Endianness,
 {
     let mut value = 0;
+    let mut num_bytes = 0;
 
     for i in 0..8 {
         let byte = reader.read::<u32>(8)?;
         value |= (byte & 0x7f) << (i * 7);
+        num_bytes += 1;
         if byte & 0x80 == 0 {
             break;
         }
     }
 
     reader.byte_align();
-    Ok(value)
+    Ok((value, num_bytes))
 }
 
 pub fn write_leb128<W, E>(writer: &mut BitWriter<W, E>, mut value: u32) -> io::Result<()>
@@ -82,7 +84,10 @@ mod tests {
             println!("testing: value={value}");
 
             let mut reader = BitReader::endian(Cursor::new(&encoding), BigEndian);
-            assert_eq!(value, parse_leb128(&mut reader).unwrap());
+            assert_eq!(
+                (value, encoding.len() as u32),
+                parse_leb128(&mut reader).unwrap()
+            );
             assert_eq!(
                 encoding.len() as u64 * 8,
                 reader.position_in_bits().unwrap()
@@ -96,7 +101,10 @@ mod tests {
             data.set_position(0);
 
             let mut reader = BitReader::endian(data, BigEndian);
-            assert_eq!(value, parse_leb128(&mut reader).unwrap());
+            assert_eq!(
+                (value, encoding.len() as u32),
+                parse_leb128(&mut reader).unwrap()
+            );
         }
     }
 }
diff --git a/net/rtp/src/av1/common/obu.rs b/net/rtp/src/av1/common/obu.rs
index c223072d..54e31414 100644
--- a/net/rtp/src/av1/common/obu.rs
+++ b/net/rtp/src/av1/common/obu.rs
@@ -7,7 +7,7 @@
 //
 // SPDX-License-Identifier: MPL-2.0
 
-use crate::av1::common::{leb128_size, parse_leb128};
+use crate::av1::common::parse_leb128;
 use bitstream_io::{BitRead, BitReader, Endianness};
 use std::io::{self, Read, Seek};
 
@@ -165,8 +165,7 @@ impl SizedObu {
 
         reader.byte_align();
 
-        let size = parse_leb128(reader)?;
-        let leb_size = leb128_size(size) as u32;
+        let (size, leb_size) = parse_leb128(reader)?;
 
         Ok(Self {
             obu_type,
diff --git a/net/rtp/src/av1/depay/imp.rs b/net/rtp/src/av1/depay/imp.rs
index 59cb2021..44d1c8be 100644
--- a/net/rtp/src/av1/depay/imp.rs
+++ b/net/rtp/src/av1/depay/imp.rs
@@ -393,24 +393,21 @@ impl RTPAv1Depay {
         aggr_header: &AggregationHeader,
         index: u32,
     ) -> Result<(u32, bool), gst::FlowError> {
-        let element_size: u32;
         let is_last_obu: bool;
 
-        if let Some(count) = aggr_header.obu_count {
+        let element_size = if let Some(count) = aggr_header.obu_count {
             is_last_obu = index + 1 == count as u32;
-            element_size = if is_last_obu {
+            if is_last_obu {
                 rtp.payload_size() - (reader.position() as u32)
             } else {
                 let mut bitreader = BitReader::endian(reader, ENDIANNESS);
-                parse_leb128(&mut bitreader).map_err(err_flow!(self, leb_read))?
+                let (size, _) = parse_leb128(&mut bitreader).map_err(err_flow!(self, leb_read))?;
+                size
             }
         } else {
-            element_size = parse_leb128(&mut BitReader::endian(&mut *reader, ENDIANNESS))
+            let (size, _) = parse_leb128(&mut BitReader::endian(&mut *reader, ENDIANNESS))
                 .map_err(err_flow!(self, leb_read))?;
-            is_last_obu = match rtp
-                .payload_size()
-                .cmp(&(reader.position() as u32 + element_size))
-            {
+            is_last_obu = match rtp.payload_size().cmp(&(reader.position() as u32 + size)) {
                 Ordering::Greater => false,
                 Ordering::Equal => true,
                 Ordering::Less => {
@@ -422,7 +419,8 @@ impl RTPAv1Depay {
                     return Err(gst::FlowError::Error);
                 }
             };
-        }
+            size
+        };
 
         Ok((element_size, is_last_obu))
     }