PE: preparations for a writer outside of Goblin

src/pe/certificate_table.rs

@@ -3,12 +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::{ctx, Pread, Pwrite};
+use crate::pe::debug;
+use scroll::{ctx, Pread, Pwrite, SizeWith};
 
 use alloc::string::ToString;
 use alloc::vec::Vec;
 
-use super::utils::pad;
+use super::utils::{align_to, pad};
 
 #[repr(u16)]
 #[non_exhaustive]
@@ -51,6 +52,10 @@ pub enum AttributeCertificateType {
     Reserved1 = 0x0003,
     /// WIN_CERT_TYPE_TS_STACK_SIGNED
     TsStackSigned = 0x0004,
+    /// WIN_CERT_TYPE_EFI_PKCS115
+    EfiPkcs115 = 0xEF0,
+    /// WIN_CERT_TYPE_EFI_GUID
+    EfiGuid = 0x0EF1,
 }
 
 impl TryFrom<u16> for AttributeCertificateType {
@@ -77,16 +82,31 @@ impl TryFrom<u16> for AttributeCertificateType {
     }
 }
 
-#[derive(Clone, Pread)]
-struct AttributeCertificateHeader {
+/// WIN_CERTIFICATE header structure
+/// It's useful beyond only parsing PE certificates
+/// This can be used to parse EFI variable structures containing certificates for example.
+/// Example: https://dox.ipxe.org/structWIN__CERTIFICATE__UEFI__GUID.html
+#[derive(Debug, Clone, Pread, Pwrite, SizeWith)]
+pub struct AttributeCertificateHeader {
     /// dwLength
-    length: u32,
-    revision: u16,
-    certificate_type: u16,
+    pub length: u32,
+    /// wRevision
+    pub revision: u16,
+    /// wCertificateType
+    pub certificate_type: u16,
 }
 
-const CERTIFICATE_DATA_OFFSET: u32 = 8;
-#[derive(Debug)]
+/// An alternative name for the WIN_CERTIFICATE header structure.
+pub type WindowsCertificateHeader = AttributeCertificateHeader;
+
+/// Static size of the [`AttributeCertificateHeader`] structure
+/// Also known under the name WIN_CERTIFICATE header structure.
+pub const ATTRIBUTE_CERTIFICATE_HEADER_SIZEOF: usize =
+    core::mem::size_of::<AttributeCertificateHeader>();
+
+/// PE-specific structure to hold certificates to associate verifiable statements about this image.
+/// The header [`AttributeCertificateHeader`] is inlined in there.
+#[derive(Debug, Clone)]
 pub struct AttributeCertificate<'a> {
     pub length: u32,
     pub revision: AttributeCertificateRevision,
@@ -95,18 +115,56 @@ pub struct AttributeCertificate<'a> {
 }
 
 impl<'a> AttributeCertificate<'a> {
+    /// Takes the raw bytes constituting a certificate
+    /// and wrap it into an AttributeCertificate.
+    /// Caller is responsible for ensuring the consistency between
+    /// the certificate type and what is in the certificate (DER, etc.).
+    pub fn from_bytes(
+        certificate: &'a [u8],
+        revision: AttributeCertificateRevision,
+        certificate_type: AttributeCertificateType,
+    ) -> error::Result<Self> {
+        // SAFETY: `ATTRIBUTE_CERTIFICATE_HEADER_SIZEOF` should always fit in a
+        // `u32`
+        // as its value fits in a `u8`.
+        let length = (align_to(certificate.len(), 8usize) + ATTRIBUTE_CERTIFICATE_HEADER_SIZEOF)
+            .try_into()
+            .map_err(|_| {
+                error::Error::Malformed(
+                    "Attribute certificate length does not fit in a `u32`".to_string(),
+                )
+            })?;
+
+        debug_assert!(length as usize >= certificate.len(), "Attribute certificate length cannot be smaller than the actual certificate contents length (potentially unaligned)");
+
+        Ok(Self {
+            length,
+            revision,
+            certificate_type,
+            certificate,
+        })
+    }
+
     pub fn parse(
         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))
-            .map_err(|_err| {
-                error::Error::Malformed(
-                    "Attribute certificate size do not fit in usize".to_string(),
-                )
-            })?;
+        let cert_size = usize::try_from(
+            header
+                .length
+                .saturating_sub(ATTRIBUTE_CERTIFICATE_HEADER_SIZEOF as u32),
+        )
+        .map_err(|_err| {
+            error::Error::Malformed("Attribute certificate size do not fit in usize".to_string())
+        })?;
+
+        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 {
@@ -136,13 +194,32 @@ impl<'a> ctx::TryIntoCtx<scroll::Endian> for &AttributeCertificate<'a> {
     /// 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;
+        debug_assert!(
+            (self.length - ATTRIBUTE_CERTIFICATE_HEADER_SIZEOF as u32) % 8 == 0,
+            "Attribute certificate's length field is unaligned"
+        );
+        debug_assert!(
+            bytes.len() >= self.length as usize,
+            "Insufficient buffer to write an aligned certificate"
+        );
         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));
+        // Extend by zero the buffer until it is aligned on a quadword (16 bytes), according to
+        // spec:
+        // > If the bCertificate content does not end on a quadword boundary, the attribute
+        // > certificate entry is padded with zeros, from the end of bCertificate to the next
+        // > quadword boundary.
+        let maybe_certificate_padding = pad(self.certificate.len(), Some(8usize));
         bytes.gwrite(self.certificate, offset)?;
         if let Some(cert_padding) = maybe_certificate_padding {
+            debug!(
+                "Extending the buffer ({}) at offset {} with {} extra bytes for quadword alignment",
+                bytes.len(),
+                *offset,
+                cert_padding.len()
+            );
+
             bytes.gwrite(&cert_padding[..], offset)?;
         }
 

src/pe/data_directories.rs

@@ -18,6 +18,25 @@ impl DataDirectory {
     pub fn parse(bytes: &[u8], offset: &mut usize) -> error::Result<Self> {
         Ok(bytes.gread_with(offset, scroll::LE)?)
     }
+
+    /// Given a view of the PE binary, represented by `bytes` and the on-disk offset `disk_offset`
+    /// this will return a view of the data directory's contents.
+    /// If the range are out of bands, this will fail with a [`error::Error::Malformed`] error.
+    pub fn data<'a>(&self, bytes: &'a [u8], disk_offset: usize) -> error::Result<&'a [u8]> {
+        let disk_end = disk_offset.saturating_add(self.size.try_into().map_err(|_| {
+            error::Error::Malformed(format!("Data directory size cannot fit in platform `usize"))
+        })?);
+
+        bytes
+            .get(disk_offset..disk_end)
+            .ok_or(error::Error::Malformed(format!(
+                "Requesting bytes from data directory at {} (end: {}) of size {}, buffer is {}",
+                disk_offset,
+                disk_end,
+                self.size,
+                bytes.len()
+            )))
+    }
 }
 
 #[derive(Debug, PartialEq, Copy, Clone)]
@@ -37,6 +56,7 @@ pub enum DataDirectoryType {
     ImportAddressTable,
     DelayImportDescriptor,
     ClrRuntimeHeader,
+    Reserved,
 }
 
 impl TryFrom<usize> for DataDirectoryType {
@@ -58,6 +78,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 +99,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 +156,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 +176,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

@@ -76,6 +76,11 @@ pub struct PE<'a> {
 }
 
 impl<'a> PE<'a> {
+    /// Returns a view in the raw bytes of this PE binary
+    pub fn bytes(&self) -> &'a [u8] {
+        self.bytes
+    }
+
     /// Reads a PE binary from the underlying `bytes`
     pub fn parse(bytes: &'a [u8]) -> error::Result<Self> {
         Self::parse_with_opts(bytes, &options::ParseOptions::default())

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!(),

src/pe/section_table.rs

@@ -5,6 +5,9 @@ use alloc::string::{String, ToString};
 use alloc::vec::Vec;
 use scroll::{ctx, Pread, Pwrite};
 
+use super::utils::align_to;
+use super::PE;
+
 #[repr(C)]
 #[derive(Debug, PartialEq, Clone, Default)]
 pub struct SectionTable {
@@ -55,6 +58,60 @@ fn base64_decode_string_entry(s: &str) -> Result<usize, ()> {
 }
 
 impl SectionTable {
+    /// Given a view of the PE binary and minimal section information,
+    /// this will return a [`SectionTable`] filled with a virtual address
+    /// automatically based on the last section offset present in the PE.
+    ///
+    /// Most of the fields will be zeroed when they require a full vision of the PE
+    /// to be derived.
+    ///
+    /// Caller is responsible to fill on-disk file offsets and various pointers.
+    pub fn new(
+        pe: &PE,
+        name: &[u8; 8],
+        contents: &[u8],
+        characteristics: u32,
+        section_alignment: u32,
+    ) -> error::Result<Self> {
+        let mut table = SectionTable::default();
+        // VA is needed only if characteristics is
+        // execute | read | write.
+        let need_virtual_address = true;
+
+        table.name = *name;
+        table.size_of_raw_data = contents.len().try_into()?;
+        table.characteristics = characteristics;
+
+        // Filling this data requires a complete overview
+        // of the final PE which may involve rewriting
+        // the complete PE.
+        table.pointer_to_raw_data = 0;
+        table.pointer_to_relocations = 0;
+
+        table.pointer_to_linenumbers = 0;
+        table.number_of_linenumbers = 0;
+        table.pointer_to_relocations = 0;
+
+        if need_virtual_address {
+            table.virtual_size = contents.len().try_into()?;
+            let mut sections = pe.sections.clone();
+            sections.sort_by_key(|sect| sect.virtual_address);
+            // Base VA = 0 ?
+            let last_section_offset = sections
+                .last()
+                .map(|last_section| last_section.virtual_address + last_section.virtual_size)
+                .ok_or(0u32)
+                .unwrap();
+
+            table.virtual_address = align_to(last_section_offset, section_alignment);
+        } else {
+            table.virtual_size = 0;
+            table.virtual_address = 0;
+        }
+
+        Ok(table)
+    }
+
     pub fn parse(
         bytes: &[u8],
         offset: &mut usize,

src/pe/utils.rs

@@ -22,6 +22,23 @@ fn aligned_pointer_to_raw_data(pointer_to_raw_data: usize) -> usize {
     pointer_to_raw_data & !PHYSICAL_ALIGN
 }
 
+// Performs arbitrary alignment of values
+// based on homogeneous numerical types.
+#[inline]
+pub fn align_to<N>(value: N, align: N) -> N
+where
+    N: core::ops::Add<Output = N>
+        + core::ops::Not<Output = N>
+        + core::ops::BitAnd<Output = N>
+        + core::ops::Sub<Output = N>
+        + core::cmp::PartialEq
+        + core::marker::Copy,
+    u8: Into<N>,
+{
+    debug_assert!(align != 0u8.into(), "Align must be non-zero");
+    return (value + align - 1u8.into()) & !(align - 1u8.into());
+}
+
 #[inline]
 fn section_read_size(section: &section_table::SectionTable, file_alignment: u32) -> usize {
     fn round_size(size: usize) -> usize {
@@ -69,7 +86,9 @@ fn section_read_size(section: &section_table::SectionTable, file_alignment: u32)
     }
 }
 
-fn rva2offset(rva: usize, section: &section_table::SectionTable) -> usize {
+/// Transforms a RVA, i.e. a relative virtual address, into an
+/// on-disk file offset, given the section table information.
+pub fn rva2offset(rva: usize, section: &section_table::SectionTable) -> usize {
     (rva - section.virtual_address as usize)
         + aligned_pointer_to_raw_data(section.pointer_to_raw_data as usize)
 }