Validate counts of parsed objects before allocating memory for them.

This prevents invalid or malicious files from causing the parser
to allocate more memory than could possibly be necessary to hold
the contents of the file.

src/model.rs

@@ -5,6 +5,8 @@ use nom::{
     IResult,
 };
 
+use crate::parser::validate_count;
+
 /// A renderable voxel model.
 #[derive(Debug, PartialEq, Eq)]
 pub struct Model {
@@ -71,5 +73,6 @@ fn parse_voxel(input: &[u8]) -> IResult<&[u8], Voxel> {
 
 pub fn parse_voxels(i: &[u8]) -> IResult<&[u8], Vec<Voxel>> {
     let (i, n) = le_u32(i)?;
-    count(parse_voxel, n as usize)(i)
+    let n = validate_count(i, n, 4)?;
+    count(parse_voxel, n)(i)
 }

src/parser.rs

@@ -5,12 +5,13 @@ use crate::{
 use nom::{
     bytes::complete::{tag, take},
     combinator::{flat_map, map_res},
+    error::make_error,
     multi::{fold_many_m_n, many0},
     number::complete::le_u32,
     sequence::pair,
     IResult,
 };
-use std::{str, str::Utf8Error};
+use std::{mem::size_of, str, str::Utf8Error};
 
 #[cfg(feature = "ahash")]
 use ahash::AHashMap as HashMap;
@@ -395,13 +396,14 @@ pub fn parse_material(i: &[u8]) -> IResult<&[u8], Material> {
 
 pub(crate) fn parse_dict(i: &[u8]) -> IResult<&[u8], Dict> {
     let (i, n) = le_u32(i)?;
+    let n = validate_count(i, n, size_of::<u32>() * 2)?;
 
-    let init = move || Dict::with_capacity(n as usize);
+    let init = move || Dict::with_capacity(n);
     let fold = |mut map: Dict, (key, value)| {
         map.insert(key, value);
         map
     };
-    fold_many_m_n(n as usize, n as usize, parse_dict_entry, init, fold)(i)
+    fold_many_m_n(n, n, parse_dict_entry, init, fold)(i)
 }
 
 fn parse_dict_entry(i: &[u8]) -> IResult<&[u8], (String, String)> {
@@ -413,6 +415,35 @@ fn parse_string(i: &[u8]) -> IResult<&[u8], String> {
     map_res(bytes, to_str)(i)
 }
 
+/// Validate that a given count of items is possible to achieve given the size of the
+/// input, then convert it to [`usize`].
+///
+/// This ensures that parsing an invalid/malicious file cannot cause excessive memory
+/// allocation in the form of data structures' capacity. It should be used whenever
+/// [`nom::multi::count`] or a `with_capacity()` function is used with a count taken from
+/// the file.
+///
+/// `minimum_object_size` must not be smaller than the minimum possible size of a parsed
+/// value.
+pub(crate) fn validate_count(
+    i: &[u8],
+    count: u32,
+    minimum_object_size: usize,
+) -> Result<usize, nom::Err<nom::error::Error<&[u8]>>> {
+    let Ok(count) = usize::try_from(count) else {
+        return Err(nom::Err::Failure(make_error(i, nom::error::ErrorKind::TooLarge)));
+    };
+
+    if count > i.len() / minimum_object_size {
+        Err(nom::Err::Failure(make_error(
+            i,
+            nom::error::ErrorKind::TooLarge,
+        )))
+    } else {
+        Ok(count)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/scene.rs

@@ -1,4 +1,6 @@
-use crate::{Color, Dict, Rotation};
+use std::mem::size_of;
+
+use crate::{parser::validate_count, Color, Dict, Rotation};
 use nom::{
     multi::count,
     number::complete::{le_i32, le_u32},
@@ -152,7 +154,8 @@ pub fn parse_scene_transform(i: &[u8]) -> IResult<&[u8], SceneTransform> {
     let (i, _ignored) = le_i32(i)?;
     let (i, layer_id) = le_u32(i)?;
     let (i, frame_count) = le_u32(i)?;
-    let (i, frames) = count(parse_dict, frame_count as usize)(i)?;
+    let frame_count = validate_count(i, frame_count, size_of::<u32>())?;
+    let (i, frames) = count(parse_dict, frame_count)(i)?;
     Ok((
         i,
         SceneTransform {
@@ -167,14 +170,16 @@ pub fn parse_scene_transform(i: &[u8]) -> IResult<&[u8], SceneTransform> {
 pub fn parse_scene_group(i: &[u8]) -> IResult<&[u8], SceneGroup> {
     let (i, header) = parse_node_header(i)?;
     let (i, child_count) = le_u32(i)?;
-    let (i, children) = count(le_u32, child_count as usize)(i)?;
+    let child_count = validate_count(i, child_count, size_of::<u32>())?;
+    let (i, children) = count(le_u32, child_count)(i)?;
     Ok((i, SceneGroup { header, children }))
 }
 
 pub fn parse_scene_shape(i: &[u8]) -> IResult<&[u8], SceneShape> {
     let (i, header) = parse_node_header(i)?;
     let (i, model_count) = le_u32(i)?;
-    let (i, models) = count(parse_scene_shape_model, model_count as usize)(i)?;
+    let model_count = validate_count(i, model_count, size_of::<u32>() * 2)?;
+    let (i, models) = count(parse_scene_shape_model, model_count)(i)?;
     Ok((i, SceneShape { header, models }))
 }
 
@@ -236,7 +241,7 @@ impl Frame {
             if let IResult::<&str, u8>::Ok((_, byte_rotation)) =
                 nom::character::complete::u8(value.as_str())
             {
-                return Some(Rotation::from_byte(byte_rotation))
+                return Some(Rotation::from_byte(byte_rotation));
             } else {
                 debug!("'_r' attribute for Frame could not be parsed! {}", value);
             }