From f69e178045bba01b13604b0093690f26bb40226c Mon Sep 17 00:00:00 2001
From: Zaggy1024 <Zaggy1024@gmail.com>
Date: Thu, 9 Feb 2023 19:36:52 -0600
Subject: [PATCH] Limit the allocated size of the sample indice table by the
 file size.

---
 mp4parse/src/unstable.rs | 58 +++++++++++++++++++++++++++++++++++-----
 mp4parse_capi/src/lib.rs |  9 ++++++-
 2 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/mp4parse/src/unstable.rs b/mp4parse/src/unstable.rs
index 6cca058f..0ea9f89b 100644
--- a/mp4parse/src/unstable.rs
+++ b/mp4parse/src/unstable.rs
@@ -2,7 +2,10 @@
 // License, v. 2.0. If a copy of the MPL was not distributed with this
 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
 use num_traits::{CheckedAdd, CheckedSub, PrimInt, Zero};
-use std::ops::{Add, Neg, Sub};
+use std::{
+    convert::TryInto,
+    ops::{Add, Neg, Sub},
+};
 
 use super::*;
 
@@ -139,6 +142,7 @@ pub struct Indice {
 pub fn create_sample_table(
     track: &Track,
     track_offset_time: CheckedInteger<i64>,
+    optional_available_size: Option<usize>,
 ) -> Option<TryVec<Indice>> {
     let timescale = match track.timescale {
         Some(ref t) => TrackTimeScale::<i64>(t.0 as i64, t.1),
@@ -153,16 +157,48 @@ pub fn create_sample_table(
     // According to spec, no sync table means every sample is sync sample.
     let has_sync_table = matches!(track.stss, Some(_));
 
-    let mut sample_size_iter = stsz.sample_sizes.iter();
-
     // Get 'stsc' iterator for (chunk_id, chunk_sample_count) and calculate the sample
     // offset address.
 
     // With large numbers of samples, the cost of many allocations dominates,
     // so it's worth iterating twice to allocate sample_table just once.
-    let total_sample_count = sample_to_chunk_iter(&stsc.samples, &stco.offsets)
-        .map(|(_, sample_counts)| sample_counts.to_usize())
-        .try_fold(0usize, usize::checked_add)?;
+    let total_sample_count = {
+        let mut sample_size_iter = stsz.sample_sizes.iter();
+        let mut sample_number: usize = 0;
+        for (chunk_id, sample_counts) in sample_to_chunk_iter(&stsc.samples, &stco.offsets) {
+            match optional_available_size {
+                Some(available_size) => {
+                    let mut end_offset = match stco.offsets.get(chunk_id as usize) {
+                        Some(v) => *v as usize,
+                        None => return None,
+                    };
+                    for _ in 0..sample_counts {
+                        match (stsz.sample_size, sample_size_iter.next()) {
+                            (_, Some(single_sample_size)) => {
+                                end_offset += *single_sample_size as usize
+                            }
+                            (sample_size, _) if sample_size > 0 => {
+                                end_offset += sample_size as usize
+                            }
+                            _ => return None,
+                        }
+                        if end_offset > available_size {
+                            break;
+                        }
+                        sample_number = match sample_number.checked_add(1) {
+                            Some(v) => v,
+                            None => return None,
+                        };
+                    }
+                }
+                None => sample_number += sample_counts as usize,
+            }
+        }
+        sample_number
+    };
+
+    let mut sample_size_iter = stsz.sample_sizes.iter();
+
     let mut sample_table = TryVec::with_capacity(total_sample_count).ok()?;
 
     for i in sample_to_chunk_iter(&stsc.samples, &stco.offsets) {
@@ -174,11 +210,19 @@ pub fn create_sample_table(
         };
         for _ in 0..sample_counts {
             let start_offset = cur_position;
-            let end_offset = match (stsz.sample_size, sample_size_iter.next()) {
+            let end_offset: CheckedInteger<u64> = match (stsz.sample_size, sample_size_iter.next())
+            {
                 (_, Some(t)) => (start_offset + *t)?,
                 (t, _) if t > 0 => (start_offset + t)?,
                 _ => 0.into(),
             };
+            match optional_available_size
+                .map(TryInto::try_into)
+                .and_then(Result::ok)
+            {
+                Some(available_size) if end_offset.0 > available_size => break,
+                _ => (),
+            }
             if end_offset == 0 {
                 return None;
             }
diff --git a/mp4parse_capi/src/lib.rs b/mp4parse_capi/src/lib.rs
index 2f4d8f1b..e4f77a39 100644
--- a/mp4parse_capi/src/lib.rs
+++ b/mp4parse_capi/src/lib.rs
@@ -1295,6 +1295,7 @@ pub unsafe extern "C" fn mp4parse_get_indice_table(
         &(*parser).context,
         &mut (*parser).sample_table,
         track_id,
+        None,
         &mut *indices,
     )
     .into()
@@ -1312,6 +1313,7 @@ pub unsafe extern "C" fn mp4parse_get_indice_table(
 pub unsafe extern "C" fn mp4parse_avif_get_indice_table(
     parser: *mut Mp4parseAvifParser,
     track_id: u32,
+    available_size: usize,
     indices: *mut Mp4parseByteData,
 ) -> Mp4parseStatus {
     if parser.is_null() {
@@ -1330,6 +1332,10 @@ pub unsafe extern "C" fn mp4parse_avif_get_indice_table(
             sequence,
             &mut (*parser).sample_table,
             track_id,
+            match available_size {
+                0 => None,
+                _ => Some(available_size),
+            },
             &mut *indices,
         )
         .into();
@@ -1342,6 +1348,7 @@ fn get_indice_table(
     context: &MediaContext,
     sample_table_cache: &mut TryHashMap<u32, TryVec<Indice>>,
     track_id: u32,
+    available_size: Option<usize>,
     indices: &mut Mp4parseByteData,
 ) -> Result<(), Mp4parseStatus> {
     let tracks = &context.tracks;
@@ -1380,7 +1387,7 @@ fn get_indice_table(
         _ => 0.into(),
     };
 
-    if let Some(v) = create_sample_table(track, offset_time) {
+    if let Some(v) = create_sample_table(track, offset_time, available_size) {
         indices.set_indices(&v);
         sample_table_cache.insert(track_id, v)?;
         return Ok(());