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/1092>

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 e084e9d17..cc9f5210b 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 53360a72b..50f2e40b1 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 ab2eb4230..6f261e14c 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))
     }