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)
 }