Merge pull request #53 from francisdb/fix/num_dir_sectors

fix: num_dir_sectors

src/internal/directory.rs

@@ -413,13 +413,40 @@ impl<F: Write + Seek> Directory<F> {
         if self.dir_entries.len() % dir_entries_per_sector == 0 {
             let start_sector = self.dir_start_sector;
             self.allocator.extend_chain(start_sector, SectorInit::Dir)?;
+            self.update_num_dir_sectors()?;
         }
         // Add a new entry to the end of the directory and return it.
         let stream_id = self.dir_entries.len() as u32;
         self.dir_entries.push(unallocated_dir_entry);
         Ok(stream_id)
     }
 
+    /// Increase header num_dir_sectors if version V4
+    /// note: not updating this value breaks ole32 compatibility
+    fn update_num_dir_sectors(&mut self) -> io::Result<()> {
+        let start_sector = self.dir_start_sector;
+        if self.version() == Version::V4 {
+            let num_dir_sectors =
+                self.count_directory_sectors(start_sector)?;
+            self.seek_within_header(40)?
+                .write_u32::<LittleEndian>(num_dir_sectors)?;
+        }
+        Ok(())
+    }
+
+    fn count_directory_sectors(
+        &mut self,
+        start_sector: u32,
+    ) -> io::Result<u32> {
+        let mut num_dir_sectors = 1;
+        let mut next_sector = self.allocator.next(start_sector)?;
+        while next_sector != consts::END_OF_CHAIN {
+            num_dir_sectors += 1;
+            next_sector = self.allocator.next(next_sector)?;
+        }
+        Ok(num_dir_sectors)
+    }
+
     /// Deallocates the specified directory entry.
     fn free_dir_entry(&mut self, stream_id: u32) -> io::Result<()> {
         debug_assert_ne!(stream_id, consts::ROOT_STREAM_ID);
@@ -428,6 +455,7 @@ impl<F: Write + Seek> Directory<F> {
         *self.dir_entry_mut(stream_id) = dir_entry;
         // TODO: Truncate directory chain if last directory sector is now all
         //       unallocated.
+        //       In that case, also call update_num_dir_sectors()
         Ok(())
     }
 

src/lib.rs

@@ -493,7 +493,18 @@ impl<F: Read + Seek> CompoundFile<F> {
         let mut dir_entries = Vec::<DirEntry>::new();
         let mut seen_dir_sectors = FnvHashSet::default();
         let mut current_dir_sector = header.first_dir_sector;
+        let mut dir_sector_count = 1;
         while current_dir_sector != consts::END_OF_CHAIN {
+            if validation.is_strict()
+                && header.version == Version::V4
+                && dir_sector_count > header.num_dir_sectors
+            {
+                invalid_data!(
+                    "Directory chain includes at least {} sectors which is greater than header num_dir_sectors {}",
+                    dir_sector_count,
+                    header.num_dir_sectors
+                );
+            }
             if current_dir_sector > consts::MAX_REGULAR_SECTOR {
                 invalid_data!(
                     "Directory chain includes invalid sector index {}",
@@ -526,6 +537,7 @@ impl<F: Read + Seek> CompoundFile<F> {
                 }
             }
             current_dir_sector = allocator.next(current_dir_sector)?;
+            dir_sector_count += 1;
         }
 
         let mut directory = Directory::new(
@@ -985,10 +997,11 @@ impl<F: fmt::Debug> fmt::Debug for CompoundFile<F> {
 
 #[cfg(test)]
 mod tests {
-    use std::io::{self, Cursor};
+    use std::io::{self, Cursor, Seek, SeekFrom};
     use std::mem::size_of;
+    use std::path::Path;
 
-    use byteorder::{LittleEndian, WriteBytesExt};
+    use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 
     use crate::internal::{consts, DirEntry, Header, Version};
 
@@ -1048,6 +1061,28 @@ mod tests {
         // under Permissive validation.
         CompoundFile::open(Cursor::new(data)).expect("open");
     }
+
+    // Regression test for https://github.com/mdsteele/rust-cfb/issues/52.
+    #[test]
+    fn update_num_dir_sectors() {
+        // Create a CFB file with 2 sectors for the directory.
+        let cursor = Cursor::new(Vec::new());
+        let mut comp = CompoundFile::create(cursor).unwrap();
+        // root + 31 entries in the first sector
+        // 1 stream entry in the second sector
+        for i in 0..32 {
+            let path = format!("stream{}", i);
+            let path = Path::new(&path);
+            comp.create_stream(path).unwrap();
+        }
+        comp.flush().unwrap();
+
+        // read num_dir_sectors from the header
+        let mut cursor = comp.into_inner();
+        cursor.seek(SeekFrom::Start(40)).unwrap();
+        let num_dir_sectors = cursor.read_u32::<LittleEndian>().unwrap();
+        assert_eq!(num_dir_sectors, 2);
+    }
 }
 
 //===========================================================================//

tests/malformed.rs

@@ -456,3 +456,31 @@ fn open_difat_terminate_freesect() {
 fn open_strict_difat_terminate_freesect() {
     CompoundFile::open_strict(difat_terminate_in_freesect()).unwrap();
 }
+
+/// Regression test for https://github.com/mdsteele/rust-cfb/issues/52.
+#[test]
+#[should_panic(
+    expected = "Directory chain includes at least 2 sectors which is greater than header num_dir_sectors 1"
+)]
+fn invalid_num_dir_sectors_issue_52() {
+    // Create a CFB file with 2 sectors for the directory.
+    let cursor = Cursor::new(Vec::new());
+    let mut comp = CompoundFile::create(cursor).unwrap();
+    // root + 31 entries in the first sector
+    // 1 stream entry in the second sector
+    for i in 0..32 {
+        let path = format!("stream{}", i);
+        let path = Path::new(&path);
+        comp.create_stream(path).unwrap();
+    }
+    comp.flush().unwrap();
+
+    // update the header and set num_dir_sectors = 1 instead of 2
+    let mut cursor = comp.into_inner();
+    cursor.seek(SeekFrom::Start(40)).unwrap();
+    cursor.write_u32::<LittleEndian>(1).unwrap();
+    cursor.flush().unwrap();
+
+    // Read the file back in.
+    CompoundFile::open_strict(cursor).unwrap();
+}