From fd320224d85c7cadd2c76ce8bd8a7e71a2223366 Mon Sep 17 00:00:00 2001
From: Philip Craig <philipjcraig@gmail.com>
Date: Tue, 11 Jan 2022 12:14:08 +1000
Subject: [PATCH] read/pe: add `ImageDataDirectory::file_range`

---
 src/read/pe/data_directory.rs | 19 ++++++++++++++
 src/read/pe/section.rs        | 48 +++++++++++++++++++++++++----------
 2 files changed, 53 insertions(+), 14 deletions(-)

diff --git a/src/read/pe/data_directory.rs b/src/read/pe/data_directory.rs
index 561fda26..3b6d58d2 100644
--- a/src/read/pe/data_directory.rs
+++ b/src/read/pe/data_directory.rs
@@ -128,6 +128,25 @@ impl pe::ImageDataDirectory {
         (self.virtual_address.get(LE), self.size.get(LE))
     }
 
+    /// Return the file offset and size of this directory entry.
+    ///
+    /// This function has some limitations:
+    /// - It requires that the data is contained in a single section.
+    /// - It uses the size field of the directory entry, which is
+    /// not desirable for all data directories.
+    /// - It uses the `virtual_address` of the directory entry as an address,
+    /// which is not valid for `IMAGE_DIRECTORY_ENTRY_SECURITY`.
+    pub fn file_range<'data>(&self, sections: &SectionTable<'data>) -> Result<(u32, u32)> {
+        let (offset, section_size) = sections
+            .pe_file_range_at(self.virtual_address.get(LE))
+            .read_error("Invalid data dir virtual address")?;
+        let size = self.size.get(LE);
+        if size > section_size {
+            return None.read_error("Invalid data dir size");
+        }
+        Ok((offset, size))
+    }
+
     /// Get the data referenced by this directory entry.
     ///
     /// This function has some limitations:
diff --git a/src/read/pe/section.rs b/src/read/pe/section.rs
index 70ac6264..54e2904d 100644
--- a/src/read/pe/section.rs
+++ b/src/read/pe/section.rs
@@ -293,9 +293,20 @@ where
 }
 
 impl<'data> SectionTable<'data> {
-    /// Return the data at the given virtual address in a PE file.
+    /// Return the file offset of the given virtual address, and the size up
+    /// to the end of the section containing it.
+    ///
+    /// Returns `None` if no section contains the address.
+    pub fn pe_file_range_at(&self, va: u32) -> Option<(u32, u32)> {
+        self.iter().find_map(|section| section.pe_file_range_at(va))
+    }
+
+    /// Return the data starting at the given virtual address, up to the end of the
+    /// section containing it.
     ///
     /// Ignores sections with invalid data.
+    ///
+    /// Returns `None` if no section contains the address.
     pub fn pe_data_at<R: ReadRef<'data>>(&self, data: R, va: u32) -> Option<&'data [u8]> {
         self.iter().find_map(|section| section.pe_data_at(data, va))
     }
@@ -326,6 +337,22 @@ impl pe::ImageSectionHeader {
         (offset, size)
     }
 
+    /// Return the file offset of the given virtual address, and the remaining size up
+    /// to the end of the section.
+    ///
+    /// Returns `None` if the section does not contain the address.
+    pub fn pe_file_range_at(&self, va: u32) -> Option<(u32, u32)> {
+        let section_va = self.virtual_address.get(LE);
+        let offset = va.checked_sub(section_va)?;
+        let (section_offset, section_size) = self.pe_file_range();
+        // Address must be within section (and not at its end).
+        if offset < section_size {
+            Some((section_offset.checked_add(offset)?, section_size - offset))
+        } else {
+            None
+        }
+    }
+
     /// Return the virtual address and size of the section.
     pub fn pe_address_range(&self) -> (u32, u32) {
         (self.virtual_address.get(LE), self.virtual_size.get(LE))
@@ -340,22 +367,15 @@ impl pe::ImageSectionHeader {
             .read_error("Invalid PE section offset or size")
     }
 
-    /// Return the data at the given virtual address if this section contains it.
+    /// Return the data starting at the given virtual address, up to the end of the
+    /// section.
     ///
     /// Ignores sections with invalid data.
+    ///
+    /// Returns `None` if the section does not contain the address.
     pub fn pe_data_at<'data, R: ReadRef<'data>>(&self, data: R, va: u32) -> Option<&'data [u8]> {
-        let section_va = self.virtual_address.get(LE);
-        let offset = va.checked_sub(section_va)?;
-        let (section_offset, section_size) = self.pe_file_range();
-        // Address must be within section (and not at its end).
-        if offset < section_size {
-            let section_data = data
-                .read_bytes_at(section_offset.into(), section_size.into())
-                .ok()?;
-            section_data.get(offset as usize..)
-        } else {
-            None
-        }
+        let (offset, size) = self.pe_file_range_at(va)?;
+        data.read_bytes_at(offset.into(), size.into()).ok()
     }
 
     /// Return the section data if it contains the given virtual address.