pe: offer basic section manipulations

examples/pe_add_section.rs

@@ -0,0 +1,50 @@
+use std::borrow::Cow;
+
+use goblin::pe::{
+    section_table::{Section, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ},
+    writer::PEWriter,
+    PE,
+};
+
+fn main() {
+    stderrlog::new().verbosity(3).init().unwrap();
+    let args: Vec<String> = std::env::args().collect();
+
+    let file = std::fs::read(&args[1]).unwrap();
+    let file = &file[..];
+    let pe = PE::parse(file).unwrap();
+    println!("{}", file.len());
+    let mut pe_writer = PEWriter::new(pe).expect("Failed to create a wrapper");
+
+    let section_name: [u8; 8] = *b".added\0\0";
+    pe_writer
+        .insert_section(
+            Section::new(
+                &section_name,
+                Some(Cow::Borrowed(&[0x0])),
+                IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ,
+            )
+            .expect("Failed to create a section"),
+        )
+        .unwrap();
+
+    let new_pe = pe_writer.write_into().unwrap();
+    std::fs::write(&args[2], &new_pe[..]).unwrap();
+
+    let old_pe_reread = PE::parse(&new_pe).unwrap();
+    let section = old_pe_reread
+        .sections
+        .into_iter()
+        .find(|s| s.name().expect("Failed to obtain name of a section") == ".added")
+        .expect("Failed to find written section");
+
+    let data = section
+        .data(&new_pe)
+        .expect("Failed to read section data")
+        .expect("Failed to obtain section data");
+    assert!(
+        data[0] == 0x0,
+        "Invalid data written in the section, got: {:?}",
+        data
+    );
+}

src/pe/certificate_table.rs

@@ -3,6 +3,7 @@
 /// https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#the-attribute-certificate-table-image-only
 /// https://learn.microsoft.com/en-us/windows/win32/api/wintrust/ns-wintrust-win_certificate
 use crate::error;
+use crate::pe::debug;
 use scroll::{ctx, Pread, Pwrite};
 
 use alloc::string::ToString;
@@ -77,7 +78,7 @@ impl TryFrom<u16> for AttributeCertificateType {
     }
 }
 
-#[derive(Clone, Pread)]
+#[derive(Debug, Clone, Pread)]
 struct AttributeCertificateHeader {
     /// dwLength
     length: u32,
@@ -99,6 +100,7 @@ impl<'a> AttributeCertificate<'a> {
         bytes: &'a [u8],
         current_offset: &mut usize,
     ) -> Result<AttributeCertificate<'a>, error::Error> {
+        debug!("reading certificate header at {current_offset}");
         // `current_offset` is moved sizeof(AttributeCertificateHeader) = 8 bytes further.
         let header: AttributeCertificateHeader = bytes.gread_with(current_offset, scroll::LE)?;
         let cert_size = usize::try_from(header.length.saturating_sub(CERTIFICATE_DATA_OFFSET))
@@ -108,6 +110,11 @@ impl<'a> AttributeCertificate<'a> {
                 )
             })?;
 
+        debug!(
+            "parsing certificate header {:#?}, predicted certificate size: {}",
+            header, cert_size
+        );
+
         if let Some(bytes) = bytes.get(*current_offset..(*current_offset + cert_size)) {
             let attr = Self {
                 length: header.length,

src/pe/data_directories.rs

@@ -18,6 +18,18 @@ impl DataDirectory {
     pub fn parse(bytes: &[u8], offset: &mut usize) -> error::Result<Self> {
         Ok(bytes.gread_with(offset, scroll::LE)?)
     }
+
+    pub fn data<'a>(&self, bytes: &'a [u8], disk_offset: usize) -> error::Result<&'a [u8]> {
+        bytes
+            .get(disk_offset..disk_offset + self.size as usize)
+            .ok_or(error::Error::Malformed(format!(
+                "Requesting bytes from data directory at {} (end: {}) of size {}, buffer is {}",
+                disk_offset,
+                disk_offset + self.size as usize,
+                self.size,
+                bytes.len()
+            )))
+    }
 }
 
 #[derive(Debug, PartialEq, Copy, Clone)]
@@ -37,6 +49,7 @@ pub enum DataDirectoryType {
     ImportAddressTable,
     DelayImportDescriptor,
     ClrRuntimeHeader,
+    Reserved,
 }
 
 impl TryFrom<usize> for DataDirectoryType {
@@ -58,6 +71,7 @@ impl TryFrom<usize> for DataDirectoryType {
             12 => Self::ImportAddressTable,
             13 => Self::DelayImportDescriptor,
             14 => Self::ClrRuntimeHeader,
+            15 => Self::Reserved,
             _ => {
                 return Err(error::Error::Malformed(
                     "Wrong data directory index number".into(),
@@ -78,9 +92,8 @@ impl ctx::TryIntoCtx<scroll::Endian> for DataDirectories {
     fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result<usize, Self::Error> {
         let offset = &mut 0;
         for opt_dd in self.data_directories {
-            if let Some((dd_offset, dd)) = opt_dd {
-                bytes.pwrite_with(dd, dd_offset, ctx)?;
-                *offset += dd_offset;
+            if let Some((_, dd)) = opt_dd {
+                bytes.gwrite_with(dd, offset, ctx)?;
             } else {
                 bytes.gwrite(&[0; SIZEOF_DATA_DIRECTORY][..], offset)?;
             }
@@ -136,6 +149,14 @@ impl DataDirectories {
     build_dd_getter!(get_clr_runtime_header, 14);
 
     pub fn dirs(&self) -> impl Iterator<Item = (DataDirectoryType, DataDirectory)> {
+        self.dirs_with_offset().map(|(a, _b, c)| (a, c))
+    }
+
+    /// Returns all data directories
+    /// with their types, offsets and contents.
+    pub fn dirs_with_offset(
+        &self,
+    ) -> impl Iterator<Item = (DataDirectoryType, usize, DataDirectory)> {
         self.data_directories
             .into_iter()
             .enumerate()
@@ -148,6 +169,6 @@ impl DataDirectories {
                 // takes into account the N possible data directories.
                 // Therefore, the unwrap can never fail as long as Rust guarantees
                 // on types are honored.
-                o.map(|(_, v)| (i.try_into().unwrap(), v)))
+                o.map(|(offset, v)| (i.try_into().unwrap(), offset, v)))
     }
 }

src/pe/header.rs

@@ -253,6 +253,7 @@ impl CoffHeader {
         let string_table_offset = self.pointer_to_symbol_table as usize
             + symbol::SymbolTable::size(self.number_of_symbol_table as usize);
         for i in 0..nsections {
+            debug!("parsing section at offset {offset}");
             let section =
                 section_table::SectionTable::parse(bytes, offset, string_table_offset as usize)?;
             debug!("({}) {:#?}", i, section);
@@ -342,6 +343,7 @@ impl ctx::TryIntoCtx<scroll::Endian> for Header {
         bytes.gwrite_with(self.dos_stub, offset, ctx)?;
         bytes.gwrite_with(self.signature, offset, scroll::LE)?;
         bytes.gwrite_with(self.coff_header, offset, ctx)?;
+        debug!("Non-optional header written, current offset: {}", offset);
         if let Some(opt_header) = self.optional_header {
             bytes.gwrite_with(opt_header, offset, ctx)?;
         }

src/pe/mod.rs

@@ -25,6 +25,7 @@ pub mod relocation;
 pub mod section_table;
 pub mod symbol;
 pub mod utils;
+pub mod writer;
 
 use crate::container;
 use crate::error;
@@ -35,6 +36,9 @@ use scroll::{ctx, Pwrite};
 
 use log::debug;
 
+use self::data_directories::DataDirectories;
+use self::data_directories::DataDirectoryType;
+
 #[derive(Debug)]
 /// An analyzed PE32/PE32+ binary
 pub struct PE<'a> {
@@ -278,6 +282,69 @@ impl<'a> PE<'a> {
         })
     }
 
+    /// This will write the data directories contents,
+    /// e.g. edata, idata, bound import table, import address table, delay import tables
+    /// .src, .pdata, .reloc, .debug, load config table, .tls, architecture table, global ptr
+    /// and clr runtime header (.cormeta) and finally the certificate table.
+    /// This will return a pair of (written data, maximum encountered offset)
+    /// This has to be done in this way because consumers are not necessarily in the same thing.
+    pub(crate) fn write_data_directories(
+        &self,
+        bytes: &mut [u8],
+        data_directories: &DataDirectories,
+        ctx: scroll::Endian,
+    ) -> Result<(usize, usize), error::Error> {
+        let mut max_offset = 0;
+        let mut written = 0;
+        // Takes care of:
+        // - export table (.edata)
+        // - import table (.idata)
+        // - bound import table
+        // - import address table
+        // - delay import tables
+        // - resource table (.rsrc)
+        // - exception table (.pdata)
+        // - base relocation table (.reloc)
+        // - debug table (.debug)
+        // - load config table
+        // - tls table (.tls)
+        // - architecture (reserved, 0 for now)
+        // - global ptr is a "empty" data directory (header-only)
+        // - clr runtime header (.cormeta is object-only)
+        for (dt_type, offset, dd) in data_directories.dirs_with_offset() {
+            // - attribute certificate table is a bit special
+            // its size is not the real size of all certificates
+            // you can check the parse logic to understand how that works
+            if dt_type == DataDirectoryType::CertificateTable {
+                let mut certificate_start = dd.virtual_address.try_into()?;
+                for certificate in &self.certificates {
+                    debug!("certificate size: {}", certificate.length);
+                    debug!("writing certificate at offset {}", certificate_start);
+                    written += bytes.gwrite_with(certificate, &mut certificate_start, ctx)?;
+                    max_offset = max(max_offset, certificate_start);
+                }
+            } else {
+                // FIXME: wow, this is very ugly. This is copying the whole PE.
+                // This should be replaced by a CoW pointer.
+                // Ideally, we should be able to explain to Rust that we
+                // are extracting from the `dd.data` location
+                // to write in the "dd.data" new location.
+                // This should be achieved with a smart usage of std::mem::take.
+                let mut view = Vec::with_capacity(bytes.len());
+                view.resize(bytes.len(), 0);
+                view.clone_from_slice(&bytes);
+                let dd_written = bytes.pwrite(dd.data(&view, offset)?, offset)?;
+                written += dd_written;
+                max_offset = max(max_offset, offset);
+                debug!(
+                    "writing {:?} at {} for {} bytes",
+                    dt_type, offset, dd_written
+                );
+            }
+        }
+        Ok((written, max_offset))
+    }
+
     pub fn write_sections(
         &self,
         bytes: &mut [u8],
@@ -502,6 +569,8 @@ impl<'a> Coff<'a> {
 
 #[cfg(test)]
 mod tests {
+    use scroll::Pwrite;
+
     use super::Coff;
     use super::PE;
 

src/pe/optional_header.rs

@@ -2,6 +2,7 @@ use crate::container;
 use crate::error;
 
 use crate::pe::data_directories;
+use crate::pe::debug;
 
 use scroll::{ctx, Endian, LE};
 use scroll::{Pread, Pwrite, SizeWith};
@@ -358,12 +359,16 @@ impl ctx::TryIntoCtx<scroll::Endian> for OptionalHeader {
         match self.standard_fields.magic {
             MAGIC_32 => {
                 bytes.gwrite_with::<StandardFields32>(self.standard_fields.into(), offset, ctx)?;
+                debug!("Wrote standard fields 32 bits (offset: {})", offset);
                 bytes.gwrite_with(WindowsFields32::try_from(self.windows_fields)?, offset, ctx)?;
+                debug!("Wrote windows fields 32 bits (offset: {})", offset);
                 bytes.gwrite_with(self.data_directories, offset, ctx)?;
             }
             MAGIC_64 => {
                 bytes.gwrite_with::<StandardFields64>(self.standard_fields.into(), offset, ctx)?;
+                debug!("Wrote standard fields 64 bits (offset: {})", offset);
                 bytes.gwrite_with(self.windows_fields, offset, ctx)?;
+                debug!("Wrote windows fields 64 bits (offset: {})", offset);
                 bytes.gwrite_with(self.data_directories, offset, ctx)?;
             }
             _ => panic!(),