From 3220e3415f7494df2cd8b41e90201a704fa3ed69 Mon Sep 17 00:00:00 2001
From: rbran <git@rubens.io>
Date: Fri, 10 Jan 2025 08:24:02 -0300
Subject: [PATCH] fix id0 dirtree not hadling EoF correctly

---
 Cargo.toml         |  2 +-
 src/id0/dirtree.rs | 20 +++++++++++++-------
 src/ida_reader.rs  | 21 +++++++++++++++------
 3 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 3d2d581..fe1cb87 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,7 +16,7 @@ num_enum = "0.7.3"
 
 [features]
 default = []
-permissive = []
+restrictive = []
 
 [[bin]]
 name = "idb-tools"
diff --git a/src/id0/dirtree.rs b/src/id0/dirtree.rs
index 287e0ed..174d88d 100644
--- a/src/id0/dirtree.rs
+++ b/src/id0/dirtree.rs
@@ -203,8 +203,7 @@ impl DirTreeEntryRaw {
         let mut entries = vec![];
         for is_value in core::iter::successors(Some(false), |x| Some(!(*x))) {
             // TODO unpack_dw/u8?
-            // TODO diferenciate an error from EoF
-            let Ok(entries_len) = data.unpack_dd() else {
+            let Some(entries_len) = data.unpack_dd_or_eof()? else {
                 break;
             };
             parse_entries(&mut *data, &mut entries, entries_len, is_value)?;
@@ -264,6 +263,7 @@ impl DirTreeEntryRaw {
         // this value had known values of 0 and 4, as long it's smaller then 0x80 there no
         // much of a problem, otherwise this could be a unpack_dw/unpack_dd
         let _unknown: u8 = bincode::deserialize_from(&mut *data)?;
+        #[cfg(feature = "restrictive")]
         ensure!(_unknown < 0x80);
         // TODO unpack_dw/u8?
         let entries_len = data.unpack_dd()?;
@@ -280,11 +280,17 @@ impl DirTreeEntryRaw {
         // NOTE in case the folder have 0 elements, there will be a 0 value, but don't take that for granted
         for is_value in core::iter::successors(Some(false), |x| Some(!(*x))) {
             // TODO unpack_dw/u8?
-            let num = match data.unpack_dd() {
-                Ok(num) => num,
-                // this is an empty folder, so the last value is optional
-                Err(_) if entries_len == 0 => break,
-                Err(e) => return Err(e),
+            let Some(num) = data.unpack_dd_or_eof()? else {
+                if entries_len == 0 {
+                    // this is an empty folder, so the last value is optional
+                    break;
+                } else {
+                    return Err(std::io::Error::new(
+                        std::io::ErrorKind::UnexpectedEof,
+                        "Unexpected EoF while reading dirtree entries",
+                    )
+                    .into());
+                }
             };
             let num = usize::try_from(num).map_err(|_| anyhow!("Invalid number of entries"))?;
             ensure!(
diff --git a/src/ida_reader.rs b/src/ida_reader.rs
index 0a605f5..52b110c 100644
--- a/src/ida_reader.rs
+++ b/src/ida_reader.rs
@@ -251,6 +251,17 @@ pub trait IdaGenericBufUnpack: IdaGenericUnpack + BufRead {
         Ok(self.fill_buf()?.first().copied())
     }
 
+    // InnerRef b47f2c2-3c08-4d40-b7ab-3c7736dce31d 0x46b690 unpack_dd
+    // NOTE the orignal implementation never fails, if input hit EoF it a partial result or 0
+    /// Reads 1 to 5 bytes.
+    fn unpack_dd_or_eof(&mut self) -> Result<Option<u32>> {
+        let Some(b1) = self.peek_u8()? else {
+            return Ok(None);
+        };
+        self.consume(1);
+        self.unpack_dd_from_byte(b1).map(Option::Some)
+    }
+
     // InnerRef fb47f2c2-3c08-4d40-b7ab-3c7736dce31d 0x48ce40
     fn read_ext_att(&mut self) -> Result<u64> {
         // InnerRef fb47f2c2-3c08-4d40-b7ab-3c7736dce31d 0x48cec0
@@ -408,13 +419,11 @@ pub trait IdaGenericUnpack: Read {
     // NOTE the orignal implementation never fails, if input hit EoF it a partial result or 0
     /// Reads 1 to 5 bytes.
     fn unpack_dd(&mut self) -> Result<u32> {
-        #[cfg(feature = "restrictive")]
         let b1 = self.read_u8()?;
-        #[cfg(not(feature = "restrictive"))]
-        let Some(b1) = self.read_u8_or_nothing()?
-        else {
-            return Ok(0);
-        };
+        self.unpack_dd_from_byte(b1)
+    }
+
+    fn unpack_dd_from_byte(&mut self, b1: u8) -> Result<u32> {
         match b1 {
             // 7 bit value
             // [0xxx xxxx]