diff --git a/src/deserializer/iter.rs b/src/deserializer/iter.rs index 558b64d..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)); } } } @@ -172,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 /// @@ -217,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 c6c604a..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), } } @@ -274,7 +280,7 @@ impl<'a> TypedStreamDeserializer<'a> { 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 @@ -306,11 +315,15 @@ impl<'a> TypedStreamDeserializer<'a> { } } } - 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)) } } } @@ -344,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); @@ -365,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)); @@ -418,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/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