diff --git a/src/deserializer/iter.rs b/src/deserializer/iter.rs index 628aaa6..f51c4c0 100644 --- a/src/deserializer/iter.rs +++ b/src/deserializer/iter.rs @@ -80,6 +80,8 @@ impl<'a, 'b> PropertyIterator<'a, 'b> { impl<'a, 'b: 'a> PropertyIterator<'a, 'b> { /// Collects only primitive data values from a `typedstream` using a depth-first-search over the deserialized object graph. /// + /// Note: There is a max depth of 100 and a max item limit of 1,000,000. + /// /// # Examples /// /// ```no_run @@ -101,23 +103,53 @@ impl<'a, 'b: 'a> PropertyIterator<'a, 'b> { /// ``` #[must_use] pub fn primitives(self) -> Vec<&'b OutputData<'a>> { + self.primitives_with_limits(100, 1000000) + } + + /// Collects primitive data values with safety limits to prevent infinite loops. + /// + /// # Arguments + /// + /// * `max_depth` - Maximum depth to traverse (prevents infinite recursion on cycles) + /// * `max_items` - Maximum total items to process (prevents runaway expansion) + #[must_use] + pub fn primitives_with_limits( + self, + max_depth: usize, + max_items: usize, + ) -> Vec<&'b OutputData<'a>> { let mut primitives = Vec::new(); - // Use an explicit stack for depth-first traversal - let mut stack: Vec> = self.collect(); - while let Some(prop) = stack.pop() { + let mut processed_items = 0; + + // Use an explicit stack for depth-first traversal with depth tracking + let initial_props: Vec> = self.collect(); + let mut stack: Vec<(Property<'a, 'b>, usize)> = + initial_props.into_iter().map(|p| (p, 0)).collect(); + + while let Some((prop, depth)) = stack.pop() { + // Safety checks to prevent infinite expansion + if processed_items >= max_items { + break; + } + if depth >= max_depth { + continue; + } + + processed_items += 1; + match prop { Property::Primitive(p) => primitives.push(p), Property::Group(mut group) => { // push children in reverse to preserve order while let Some(child) = group.pop() { - stack.push(child); + stack.push((child, depth + 1)); } } Property::Object { data, .. } => { // data is a PropertyIterator; collect its items let mut nested: Vec<_> = data.collect(); while let Some(child) = nested.pop() { - stack.push(child); + stack.push((child, depth + 1)); } } } @@ -142,26 +174,25 @@ impl<'a, 'b: 'a> Iterator for PropertyIterator<'a, 'b> { class: cls, data: _, }) = self.object_table.get(*idx) + && let Some(Archived::Class(cls)) = self.object_table.get(*cls) { - if let Some(Archived::Class(cls)) = self.object_table.get(*cls) { - let class_name = self - .type_table - .get(cls.name_index) - .and_then(|types| types.first()) - .and_then(|t| match t { - Type::String(name) => Some(*name), - _ => None, - }) - .unwrap_or("Unknown Class"); - // recurse into that object’s own data - let sub_iter = - PropertyIterator::new(self.object_table, self.type_table, *idx)?; - resolved.push(Property::Object { - class: cls, - name: class_name, - data: sub_iter, - }); - } + let class_name = self + .type_table + .get(cls.name_index) + .and_then(|types| types.first()) + .and_then(|t| match t { + Type::String(name) => Some(*name), + _ => None, + }) + .unwrap_or("Unknown Class"); + // recurse into that object’s own data + let sub_iter = + PropertyIterator::new(self.object_table, self.type_table, *idx)?; + resolved.push(Property::Object { + class: cls, + name: class_name, + data: sub_iter, + }); } } prim => resolved.push(Property::Primitive(prim)), @@ -173,8 +204,10 @@ impl<'a, 'b: 'a> Iterator for PropertyIterator<'a, 'b> { /// Print a resolved [`PropertyIterator`] in a human-readable tree format for debugging. /// -/// This function recursively prints all properties with proper indentation to show the nested structure -/// of the deserialized object graph. +/// This function iteratively prints all properties with proper indentation to show the nested structure +/// of the deserialized object graph. Uses an explicit stack to avoid stack overflow for large structures. +/// +/// Note: There is a max depth of 100 and a max item limit of 1,000,000. /// /// # Arguments /// @@ -218,40 +251,80 @@ impl<'a, 'b: 'a> Iterator for PropertyIterator<'a, 'b> { /// Primitive: SignedInteger(0) /// ``` pub fn print_resolved(iter: PropertyIterator<'_, '_>, indent: usize) { - for prop in iter { - print_property(prop, indent); - } + print_resolved_with_limits(iter, indent, 100, 1000000); } -/// Print a single `Property` with indentation, recursing for nested data. +/// Print a resolved [`PropertyIterator`] with depth and item limits to prevent infinite expansion. /// /// # Arguments /// -/// * `prop` - The property to print. -/// * `indent` - Number of spaces to indent each level. -/// -/// This function is intended for debugging purposes. -pub(crate) fn print_property<'a, 'b: 'a>(prop: Property<'a, 'b>, indent: usize) { - match prop { - Property::Object { - class: _, - name, - data, - } => { - // Print the object itself - println!("{:indent$}Object: {:?}", "", name, indent = indent); - // Recurse into its children with increased indent - print_resolved(data, indent + 2); +/// * `iter` - The property iterator to print +/// * `indent` - Number of spaces to indent each level +/// * `max_depth` - Maximum depth to traverse (prevents infinite recursion on cycles) +/// * `max_items` - Maximum total items to print (prevents runaway output) +fn print_resolved_with_limits( + iter: PropertyIterator<'_, '_>, + indent: usize, + max_depth: usize, + max_items: usize, +) { + // Use an explicit stack to avoid recursion and potential stack overflow + let mut stack: Vec<(Property<'_, '_>, usize)> = Vec::new(); + let mut items_printed = 0; + + // Push all properties from the iterator onto the stack with their indent level + for prop in iter { + stack.push((prop, indent)); + } + + // Process the stack + while let Some((prop, current_indent)) = stack.pop() { + // Safety checks to prevent infinite expansion + if items_printed >= max_items { + println!( + "{:indent$}... (truncated after {max_items} items)", + "", + indent = current_indent + ); + break; } - Property::Group(slice) => { - println!("{:indent$}Group:", "", indent = indent); - // Drill into every slot in the group - for slot in slice { - print_property(slot, indent + 2); - } + + let depth = (current_indent - indent) / 2; + if depth >= max_depth { + println!( + "{:indent$}... (max depth {max_depth} reached)", + "", + indent = current_indent + ); + continue; } - Property::Primitive(p) => { - println!("{:indent$}Primitive: {:?}", "", p, indent = indent); + + items_printed += 1; + + match prop { + Property::Object { + class: _, + name, + data, + } => { + // Print the object itself + println!("{:indent$}Object: {:?}", "", name, indent = current_indent); + // Push its children onto the stack with increased indent (in reverse order) + let children: Vec<_> = data.collect(); + for child in children.into_iter().rev() { + stack.push((child, current_indent + 2)); + } + } + Property::Group(group) => { + println!("{:indent$}Group:", "", indent = current_indent); + // Push every slot in the group onto the stack with increased indent (in reverse order) + for slot in group.into_iter().rev() { + stack.push((slot, current_indent + 2)); + } + } + Property::Primitive(p) => { + println!("{:indent$}Primitive: {:?}", "", p, indent = current_indent); + } } } } diff --git a/src/deserializer/typedstream.rs b/src/deserializer/typedstream.rs index cc55ba7..d29d974 100644 --- a/src/deserializer/typedstream.rs +++ b/src/deserializer/typedstream.rs @@ -54,12 +54,18 @@ impl<'a> TypedStreamDeserializer<'a> { /// ``` #[must_use] pub fn new(data: &'a [u8]) -> Self { + // Estimate initial capacities based on data size to reduce reallocations + let estimated_size = data.len(); + let type_capacity = (estimated_size / 64).clamp(16, 256); + let object_capacity = (estimated_size / 32).clamp(32, 512); + let embedded_capacity = (estimated_size / 128).clamp(8, 64); + Self { data, position: 0, - type_table: Vec::with_capacity(16), - object_table: Vec::with_capacity(32), - seen_embedded_types: Vec::with_capacity(8), + type_table: Vec::with_capacity(type_capacity), + object_table: Vec::with_capacity(object_capacity), + seen_embedded_types: Vec::with_capacity(embedded_capacity), } } @@ -232,10 +238,10 @@ impl<'a> TypedStreamDeserializer<'a> { // The class we just appended (*idx*) is the **parent** of the // class we appended in the previous iteration (*prev_new*) - if let Some(child_idx) = prev_new { - if let Archived::Class(ref mut child_cls) = self.object_table[child_idx] { - child_cls.parent_index = Some(idx); - } + if let Some(child_idx) = prev_new + && let Archived::Class(ref mut child_cls) = self.object_table[child_idx] + { + child_cls.parent_index = Some(idx); } // remember the first class we ever pushed @@ -264,17 +270,17 @@ impl<'a> TypedStreamDeserializer<'a> { // Patch the outer-most newly created class so that it points to the // already-existing parent (or to `None` if EMPTY terminated the list). - if let Some(outer_idx) = prev_new { - if let Archived::Class(ref mut outer_cls) = self.object_table[outer_idx] { - outer_cls.parent_index = final_parent; - } + if let Some(outer_idx) = prev_new + && let Archived::Class(ref mut outer_cls) = self.object_table[outer_idx] + { + outer_cls.parent_index = final_parent; } // Return the index of the bottom-most child we created first. Ok(Some(first_idx)) } - fn read_object(&mut self) -> Result { + fn read_object(&mut self) -> Result> { match *read_byte_at(self.data, self.position)? { START => { let placeholder_index = self.object_table.len(); @@ -284,9 +290,12 @@ impl<'a> TypedStreamDeserializer<'a> { self.position += 1; if let Some(cls) = self.read_class()? { + // Estimate initial capacity for object data to reduce reallocations + let estimated_data_capacity = + ((self.data.len() - self.position) / 64).clamp(8, 64); self.object_table[placeholder_index] = Archived::Object { class: cls, - data: Vec::with_capacity(8), + data: Vec::with_capacity(estimated_data_capacity), }; while self.position < self.data.len() && *read_byte_at(self.data, self.position)? != END @@ -294,24 +303,27 @@ impl<'a> TypedStreamDeserializer<'a> { // Read the next type, which should be an object if let Some(next_index) = self.read_type(false)? { // Recursively read the types for this object - if let Some(data) = self.read_types(next_index)? { - if let Some(Archived::Object { + if let Some(data) = self.read_types(next_index)? + && let Some(Archived::Object { class: _, data: data_vec, }) = self.object_table.get_mut(placeholder_index) - { - // Add the data to the object - data_vec.push(data); - } + { + // Add the data to the object + data_vec.push(data); } } } } - Ok(placeholder_index) + Ok(Some(placeholder_index)) + } + EMPTY => { + self.position += 1; + Ok(None) } ptr => { let pointer = read_pointer(&ptr)?; - Ok(pointer.value as usize) + Ok(Some(pointer.value as usize)) } } } @@ -345,6 +357,7 @@ impl<'a> TypedStreamDeserializer<'a> { fn read_types(&mut self, types_index: usize) -> Result>>> { // Start reading types from the specified index in the type table + let len = self.type_table[types_index].len(); let mut out_v = Vec::with_capacity(len); @@ -366,7 +379,11 @@ impl<'a> TypedStreamDeserializer<'a> { Type::Object => { let obj_idx = self.read_object()?; self.position += 1; - out_v.push(OutputData::Object(obj_idx)); + if let Some(obj_idx) = obj_idx { + out_v.push(OutputData::Object(obj_idx)); + } else { + out_v.push(OutputData::Null); + } } Type::String(s) => { out_v.push(OutputData::String(s)); @@ -419,21 +436,20 @@ impl<'a> TypedStreamDeserializer<'a> { let pointer = read_pointer(&ptr)?; let ref_tag = pointer.value as usize; - if ref_tag as usize >= self.type_table.len() { + // Optimize bounds checking + if ref_tag >= self.type_table.len() { return Ok(None); } if is_embedded_type { // We only want to include the first embedded reference tag, not subsequent references to the same embed - if !self.seen_embedded_types.contains(&ref_tag) - && self.type_table.get(ref_tag as usize).is_some() - { + if !self.seen_embedded_types.contains(&ref_tag) { self.object_table.push(Archived::Type(ref_tag)); self.seen_embedded_types.push(ref_tag); } } - Ok(Some(ref_tag as usize)) + Ok(Some(ref_tag)) } } } diff --git a/src/error.rs b/src/error.rs index 4bbdcb7..b84e930 100644 --- a/src/error.rs +++ b/src/error.rs @@ -56,7 +56,7 @@ impl Display for TypedStreamError { } TypedStreamError::InvalidHeader => write!(f, "Invalid header in typedstream!"), TypedStreamError::InvalidPointer(pointer) => { - write!(f, "Invalid pointer: {pointer}") + write!(f, "Invalid pointer: {pointer:x}") } TypedStreamError::InvalidArray(offset) => { write!(f, "Invalid array at index: {offset:x}") diff --git a/src/lib.rs b/src/lib.rs index 167593f..b98bb51 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6866,4 +6866,28 @@ mod test_typedstream_deserializer { assert_eq!(typedstream.type_table, expected_types); assert_eq!(typedstream.object_table, expected_objects); } + + #[test] + fn test_parse_large_with_null() { + let typedstream_path = current_dir() + .unwrap() + .as_path() + .join("src/test_data/HugeWithRefs"); + let mut file = File::open(typedstream_path).unwrap(); + let mut bytes = vec![]; + file.read_to_end(&mut bytes).unwrap(); + + let mut typedstream = TypedStreamDeserializer::new(&bytes); + + // Unwrapping here means we resolved the object + let root = typedstream.oxidize().unwrap(); + println!("\nResults: {root:?}"); + // Unwrapping here means we resolved the properties + let root_obj = typedstream.resolve_properties(root).unwrap(); + println!("\nResults: {root_obj:?}"); + + // This will only complete if the circular references are handled + let primitives = root_obj.primitives(); + println!("\nPrimitive Values: {primitives:?}"); + } } diff --git a/src/models/output_data.rs b/src/models/output_data.rs index 90a46de..565dcb9 100644 --- a/src/models/output_data.rs +++ b/src/models/output_data.rs @@ -19,6 +19,8 @@ pub enum OutputData<'a> { Array(&'a [u8]), /// Reference to another object by index in the [`object_table`](crate::deserializer::typedstream::TypedStreamDeserializer::object_table). Object(usize), + /// Represents a null value. + Null, } impl<'a> OutputData<'a> { @@ -187,6 +189,7 @@ impl std::fmt::Display for OutputData<'_> { OutputData::Byte(b) => write!(f, "0x{b:02x}"), OutputData::Array(arr) => write!(f, "[{arr:02x?}]"), OutputData::Object(idx) => write!(f, "Object({idx})"), + OutputData::Null => write!(f, "Null"), } } } diff --git a/src/models/types.rs b/src/models/types.rs index fe8259c..622ecaf 100644 --- a/src/models/types.rs +++ b/src/models/types.rs @@ -93,7 +93,7 @@ impl<'a> Type<'a> { Self::String(str) } - pub(crate) fn get_array_length(types: &[u8]) -> Option> { + pub(crate) fn get_array_length(types: &'_ [u8]) -> Option>> { if types.first() == Some(&0x5b) { let len = types[1..] @@ -109,7 +109,7 @@ impl<'a> Type<'a> { None } - pub(crate) fn read_new_type(data: &[u8]) -> Result>> { + pub(crate) fn read_new_type(data: &'_ [u8]) -> Result>>> { // Get the type of the object let type_length = read_unsigned_int(data)?; diff --git a/src/test_data/HugeWithRefs b/src/test_data/HugeWithRefs new file mode 100644 index 0000000..a4718d9 Binary files /dev/null and b/src/test_data/HugeWithRefs differ