pe: write support for PE binaries (#361)

* strtab: offer `len` method

* pe(header): introduce full DOS header and DOS stub

Required to rematerialize a whole PE, otherwise, it is a corrupted PE.

* pe(section_table): add `data` function to retrieve the section data safely

It now returns a Cow as we need to transfer ownership whenever we have
virtual size > size of raw data, which forces us to pad with zeroes
according to the PE specification.

* pe(utils): add a zero-extension padding helper

It can be used to zero-extend any buffer to make its length
aligned on some boundary, e.g. quadwords, for attribute certificates.

It is designed to be "zero cost" if you have no alignment to do, which
will result in no-allocation!

* pe: implement write support for AttributeCertificate

* pe: implement write support for DataDirectories

* pe: implement write support for Header

It enables a lossless conversion from 64 bits to 32 bits headers
except on platforms… that are strange.

* pe: implement write support for SymbolTable

* pe: implement write support for data directories

This tucks `DataDirectory` with a Deref-style and expose `offset` for internal consumers (this crate)
and offers write support for data directories.

* pe: implement write support for SectionTable

All that is needed is to import the allocation `Vec`.

* pe: implement write support

* examples(pe): add a identity rewrite PE binary

This is a trivial rewriter that performs the identity transformation.

* pe: debug assert that no overlapping sections are written

To further improve the debuggability, a non-overlapping sanity check is added as a
debug_assert!.

* pe(coff): COFF symbols and COFF strings are deprecated

Previously, I (we?) thought that COFF symbols/strings were always there.

In fact, this is not the case and they are deprecated
according to https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#coff-file-header-object-and-image

> The file offset of the COFF symbol table, or zero
> if no COFF symbol table is present. This value should be zero for an image
> because COFF debugging information is deprecated.

* pe(write): factor into `write_certificates`

We had the write certificate logic in the `TryIntoCtx`, but it makes
sense to have it separated to let consumers call it manually.

Cargo.toml

@@ -54,3 +54,6 @@ archive = ["alloc"]
 [badges.travis-ci]
 branch = "master"
 repository = "m4b/goblin"
+
+[dev-dependencies]
+stderrlog = "0.5.4"

examples/dotnet_pe_analysis.rs

@@ -67,7 +67,8 @@ fn main() {
         .expect("No CLI header");
     let sections = &pe.sections;
 
-    let cli_header_value: CliHeader = get_data(file, sections, cli_header, file_alignment).unwrap();
+    let cli_header_value: CliHeader =
+        get_data(file, sections, *cli_header, file_alignment).unwrap();
     println!("{:#?}", cli_header_value);
     let metadata_root: MetadataRoot =
         get_data(file, sections, cli_header_value.metadata, file_alignment).unwrap();

examples/rewrite_pe.rs

@@ -0,0 +1,42 @@
+use goblin::pe::PE;
+use scroll::Pwrite;
+
+fn main() {
+    stderrlog::new().verbosity(1).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!("read {}", &args[1]);
+
+    println!(
+        "file alignment: {:?}",
+        pe.header
+            .optional_header
+            .unwrap()
+            .windows_fields
+            .file_alignment
+    );
+
+    let mut new_pe = vec![0u8; file.len() + 8192];
+    let new_len = new_pe.pwrite(pe, 0).unwrap();
+    let pe = PE::parse(file).unwrap();
+
+    let out = &new_pe[..new_len];
+    std::fs::write(&args[2], &out).unwrap();
+    println!("written as {}", &args[2]);
+    println!(
+        "original PE size: {} bytes, new PE size: {} bytes, delta (new - original): {} bytes",
+        file.len(),
+        out.len(),
+        out.len() as isize - file.len() as isize
+    );
+
+    let new_pe = PE::parse(&new_pe).unwrap();
+    println!(
+        "original signatures: {}, new signatures: {}",
+        pe.certificates.len(),
+        new_pe.certificates.len()
+    );
+}

src/error.rs

@@ -3,6 +3,7 @@
 
 use alloc::string::String;
 use core::fmt;
+use core::num::TryFromIntError;
 use core::result;
 #[cfg(feature = "std")]
 use std::{error, io};
@@ -42,6 +43,12 @@ impl From<io::Error> for Error {
     }
 }
 
+impl From<TryFromIntError> for Error {
+    fn from(err: TryFromIntError) -> Error {
+        Error::Malformed(format!("Integer do not fit: {err}"))
+    }
+}
+
 impl From<scroll::Error> for Error {
     fn from(err: scroll::Error) -> Error {
         Error::Scroll(err)

src/pe/certificate_table.rs

@@ -3,11 +3,13 @@
 /// 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 scroll::Pread;
+use scroll::{ctx, Pread, Pwrite};
 
 use alloc::string::ToString;
 use alloc::vec::Vec;
 
+use super::utils::pad;
+
 #[repr(u16)]
 #[non_exhaustive]
 #[derive(Debug, PartialEq, Copy, Clone)]
@@ -39,7 +41,7 @@ impl TryFrom<u16> for AttributeCertificateRevision {
 }
 
 #[repr(u16)]
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Copy, Clone)]
 pub enum AttributeCertificateType {
     /// WIN_CERT_TYPE_X509
     X509 = 0x0001,
@@ -128,7 +130,28 @@ impl<'a> AttributeCertificate<'a> {
     }
 }
 
+impl<'a> ctx::TryIntoCtx<scroll::Endian> for &AttributeCertificate<'a> {
+    type Error = error::Error;
+
+    /// Writes an aligned attribute certificate in the buffer.
+    fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result<usize, Self::Error> {
+        let offset = &mut 0;
+        bytes.gwrite_with(self.length, offset, ctx)?;
+        bytes.gwrite_with(self.revision as u16, offset, ctx)?;
+        bytes.gwrite_with(self.certificate_type as u16, offset, ctx)?;
+        // Extend by zero the buffer until it is aligned on a quadword (16 bytes).
+        let maybe_certificate_padding = pad(self.certificate.len(), Some(16usize));
+        bytes.gwrite(self.certificate, offset)?;
+        if let Some(cert_padding) = maybe_certificate_padding {
+            bytes.gwrite(&cert_padding[..], offset)?;
+        }
+
+        Ok(*offset)
+    }
+}
+
 pub type CertificateDirectoryTable<'a> = Vec<AttributeCertificate<'a>>;
+
 pub(crate) fn enumerate_certificates(
     bytes: &[u8],
     table_virtual_address: u32,

src/pe/data_directories.rs

@@ -1,5 +1,8 @@
 use crate::error;
-use scroll::{Pread, Pwrite, SizeWith};
+use scroll::{
+    ctx::{self},
+    Pread, Pwrite, SizeWith,
+};
 
 #[repr(C)]
 #[derive(Debug, PartialEq, Copy, Clone, Default, Pread, Pwrite, SizeWith)]
@@ -13,14 +16,86 @@ const NUM_DATA_DIRECTORIES: usize = 16;
 
 impl DataDirectory {
     pub fn parse(bytes: &[u8], offset: &mut usize) -> error::Result<Self> {
-        let dd = bytes.gread_with(offset, scroll::LE)?;
-        Ok(dd)
+        Ok(bytes.gread_with(offset, scroll::LE)?)
+    }
+}
+
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum DataDirectoryType {
+    ExportTable,
+    ImportTable,
+    ResourceTable,
+    ExceptionTable,
+    CertificateTable,
+    BaseRelocationTable,
+    DebugTable,
+    Architecture,
+    GlobalPtr,
+    TlsTable,
+    LoadConfigTable,
+    BoundImportTable,
+    ImportAddressTable,
+    DelayImportDescriptor,
+    ClrRuntimeHeader,
+}
+
+impl TryFrom<usize> for DataDirectoryType {
+    type Error = error::Error;
+    fn try_from(value: usize) -> Result<Self, Self::Error> {
+        Ok(match value {
+            0 => Self::ExportTable,
+            1 => Self::ImportTable,
+            2 => Self::ResourceTable,
+            3 => Self::ExceptionTable,
+            4 => Self::CertificateTable,
+            5 => Self::BaseRelocationTable,
+            6 => Self::DebugTable,
+            7 => Self::Architecture,
+            8 => Self::GlobalPtr,
+            9 => Self::TlsTable,
+            10 => Self::LoadConfigTable,
+            11 => Self::BoundImportTable,
+            12 => Self::ImportAddressTable,
+            13 => Self::DelayImportDescriptor,
+            14 => Self::ClrRuntimeHeader,
+            _ => {
+                return Err(error::Error::Malformed(
+                    "Wrong data directory index number".into(),
+                ))
+            }
+        })
     }
 }
 
 #[derive(Debug, PartialEq, Copy, Clone, Default)]
 pub struct DataDirectories {
-    pub data_directories: [Option<DataDirectory>; NUM_DATA_DIRECTORIES],
+    pub data_directories: [Option<(usize, DataDirectory)>; NUM_DATA_DIRECTORIES],
+}
+
+impl ctx::TryIntoCtx<scroll::Endian> for DataDirectories {
+    type Error = error::Error;
+
+    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;
+            } else {
+                bytes.gwrite(&[0; SIZEOF_DATA_DIRECTORY][..], offset)?;
+            }
+        }
+        Ok(NUM_DATA_DIRECTORIES * SIZEOF_DATA_DIRECTORY)
+    }
+}
+
+macro_rules! build_dd_getter {
+    ($dd_name:tt, $index:tt) => {
+        pub fn $dd_name(&self) -> Option<&DataDirectory> {
+            let idx = $index;
+            self.data_directories[idx].as_ref().map(|(_, dd)| dd)
+        }
+    };
 }
 
 impl DataDirectories {
@@ -37,70 +112,42 @@ impl DataDirectories {
             let dd = if dd.virtual_address == 0 && dd.size == 0 {
                 None
             } else {
-                Some(dd)
+                Some((*offset, dd))
             };
             *dir = dd;
         }
         Ok(DataDirectories { data_directories })
     }
-    pub fn get_export_table(&self) -> &Option<DataDirectory> {
-        let idx = 0;
-        &self.data_directories[idx]
-    }
-    pub fn get_import_table(&self) -> &Option<DataDirectory> {
-        let idx = 1;
-        &self.data_directories[idx]
-    }
-    pub fn get_resource_table(&self) -> &Option<DataDirectory> {
-        let idx = 2;
-        &self.data_directories[idx]
-    }
-    pub fn get_exception_table(&self) -> &Option<DataDirectory> {
-        let idx = 3;
-        &self.data_directories[idx]
-    }
-    pub fn get_certificate_table(&self) -> &Option<DataDirectory> {
-        let idx = 4;
-        &self.data_directories[idx]
-    }
-    pub fn get_base_relocation_table(&self) -> &Option<DataDirectory> {
-        let idx = 5;
-        &self.data_directories[idx]
-    }
-    pub fn get_debug_table(&self) -> &Option<DataDirectory> {
-        let idx = 6;
-        &self.data_directories[idx]
-    }
-    pub fn get_architecture(&self) -> &Option<DataDirectory> {
-        let idx = 7;
-        &self.data_directories[idx]
-    }
-    pub fn get_global_ptr(&self) -> &Option<DataDirectory> {
-        let idx = 8;
-        &self.data_directories[idx]
-    }
-    pub fn get_tls_table(&self) -> &Option<DataDirectory> {
-        let idx = 9;
-        &self.data_directories[idx]
-    }
-    pub fn get_load_config_table(&self) -> &Option<DataDirectory> {
-        let idx = 10;
-        &self.data_directories[idx]
-    }
-    pub fn get_bound_import_table(&self) -> &Option<DataDirectory> {
-        let idx = 11;
-        &self.data_directories[idx]
-    }
-    pub fn get_import_address_table(&self) -> &Option<DataDirectory> {
-        let idx = 12;
-        &self.data_directories[idx]
-    }
-    pub fn get_delay_import_descriptor(&self) -> &Option<DataDirectory> {
-        let idx = 13;
-        &self.data_directories[idx]
-    }
-    pub fn get_clr_runtime_header(&self) -> &Option<DataDirectory> {
-        let idx = 14;
-        &self.data_directories[idx]
+
+    build_dd_getter!(get_export_table, 0);
+    build_dd_getter!(get_import_table, 1);
+    build_dd_getter!(get_resource_table, 2);
+    build_dd_getter!(get_exception_table, 3);
+    build_dd_getter!(get_certificate_table, 4);
+    build_dd_getter!(get_base_relocation_table, 5);
+    build_dd_getter!(get_debug_table, 6);
+    build_dd_getter!(get_architecture, 7);
+    build_dd_getter!(get_global_ptr, 8);
+    build_dd_getter!(get_tls_table, 9);
+    build_dd_getter!(get_load_config_table, 10);
+    build_dd_getter!(get_bound_import_table, 11);
+    build_dd_getter!(get_import_address_table, 12);
+    build_dd_getter!(get_delay_import_descriptor, 13);
+    build_dd_getter!(get_clr_runtime_header, 14);
+
+    pub fn dirs(&self) -> impl Iterator<Item = (DataDirectoryType, DataDirectory)> {
+        self.data_directories
+            .into_iter()
+            .enumerate()
+            // (Index, Option<DD>) -> Option<(Index, DD)> -> (DDT, DD)
+            .filter_map(|(i, o)|
+                // We should not have invalid indexes.
+                // Indeed: `data_directories: &[_; N]` where N is the number
+                // of data directories.
+                // The `TryFrom` trait for integers to DataDirectoryType
+                // 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)))
     }
 }