diff --git a/examples/pe_add_section.rs b/examples/pe_add_section.rs new file mode 100644 index 000000000..2416c5b3d --- /dev/null +++ b/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 = 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( + §ion_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 + ); +} diff --git a/src/pe/certificate_table.rs b/src/pe/certificate_table.rs index be45ce48d..e71dd1eb7 100644 --- a/src/pe/certificate_table.rs +++ b/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 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, 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, diff --git a/src/pe/data_directories.rs b/src/pe/data_directories.rs index e65db5953..6878dc406 100644 --- a/src/pe/data_directories.rs +++ b/src/pe/data_directories.rs @@ -18,6 +18,18 @@ impl DataDirectory { pub fn parse(bytes: &[u8], offset: &mut usize) -> error::Result { 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 for DataDirectoryType { @@ -58,6 +71,7 @@ impl TryFrom 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 for DataDirectories { fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result { 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 { + 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 { 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))) } } diff --git a/src/pe/header.rs b/src/pe/header.rs index 06e23c0b0..c5cea0c46 100644 --- a/src/pe/header.rs +++ b/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 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)?; } diff --git a/src/pe/mod.rs b/src/pe/mod.rs index 2336fddc5..a7432b9ea 100644 --- a/src/pe/mod.rs +++ b/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; diff --git a/src/pe/optional_header.rs b/src/pe/optional_header.rs index 852f4dced..b043fc38e 100644 --- a/src/pe/optional_header.rs +++ b/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 for OptionalHeader { match self.standard_fields.magic { MAGIC_32 => { bytes.gwrite_with::(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::(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!(), diff --git a/src/pe/section_table.rs b/src/pe/section_table.rs index f49c2b689..b7229bd12 100644 --- a/src/pe/section_table.rs +++ b/src/pe/section_table.rs @@ -1,10 +1,14 @@ use crate::error::{self, Error}; use crate::pe::relocation; use alloc::borrow::Cow; +use alloc::borrow::ToOwned; 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 { @@ -23,6 +27,46 @@ pub struct SectionTable { pub const SIZEOF_SECTION_TABLE: usize = 8 * 5; +#[derive(Debug, PartialEq, Clone)] +pub struct Section<'a> { + pub(crate) table: SectionTable, + pub(crate) contents: Option>, + pub(crate) relocations: Vec, +} + +impl<'a> Section<'a> { + pub fn new( + name: &[u8; 8], + contents: Option>, + characteristics: u32, + ) -> error::Result { + let mut table = SectionTable::default(); + + table.name = *name; + table.characteristics = characteristics; + + // Filling this data requires a complete overview + // of the final PE which may involve rewriting + // the complete PE. + table.size_of_raw_data = 0; + 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; + + table.virtual_size = 0; + table.virtual_address = 0; + + Ok(Self { + table, + contents, + relocations: Vec::new(), + }) + } +} + // Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L70 // Decodes a string table entry in base 64 (//AAAAAA). Expects string without // prefixed slashes. @@ -55,6 +99,52 @@ fn base64_decode_string_entry(s: &str) -> Result { } impl SectionTable { + pub fn new( + pe: &PE, + name: &[u8; 8], + contents: &[u8], + characteristics: u32, + section_alignment: u32, + ) -> error::Result { + 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, @@ -81,7 +171,21 @@ impl SectionTable { Ok(table) } - pub fn data<'a, 'b: 'a>(&'a self, pe_bytes: &'b [u8]) -> error::Result>> { + /// Consumes this header into a fully-fledged and self-contained Section + /// The section may need to outlive the PE bytes. + pub fn into_section<'b, 'a: 'b>(mut self, pe_bytes: &'a [u8]) -> error::Result> { + let contents: Option> = self.data(pe_bytes)?.to_owned(); + let relocations = self.relocations(pe_bytes)?.collect::>(); + let table = core::mem::take(&mut self); + + Ok(Section { + table, + contents, + relocations, + }) + } + + pub fn data<'a, 'b: 'a>(&'a self, pe_bytes: &'b [u8]) -> error::Result>> { let section_start: usize = self.pointer_to_raw_data.try_into().map_err(|_| { Error::Malformed(format!("Virtual address cannot fit in platform `usize`")) })?; diff --git a/src/pe/utils.rs b/src/pe/utils.rs index 289ccc529..19a6f4807 100644 --- a/src/pe/utils.rs +++ b/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(value: N, align: N) -> N +where + N: core::ops::Add + + core::ops::Not + + core::ops::BitAnd + + core::ops::Sub + + core::cmp::PartialEq + + core::marker::Copy, + u8: Into, +{ + debug_assert!(align != 0u8.into(), "Align must be non-zero"); + return (value + align - 1u8.into()) & !(align - 1u8.into()); +} + #[inline] fn section_read_size(section: §ion_table::SectionTable, file_alignment: u32) -> usize { fn round_size(size: usize) -> usize { @@ -69,7 +86,7 @@ fn section_read_size(section: §ion_table::SectionTable, file_alignment: u32) } } -fn rva2offset(rva: usize, section: §ion_table::SectionTable) -> usize { +pub(crate) fn rva2offset(rva: usize, section: §ion_table::SectionTable) -> usize { (rva - section.virtual_address as usize) + aligned_pointer_to_raw_data(section.pointer_to_raw_data as usize) } diff --git a/src/pe/writer.rs b/src/pe/writer.rs new file mode 100644 index 000000000..1fb3c7cf1 --- /dev/null +++ b/src/pe/writer.rs @@ -0,0 +1,503 @@ +use alloc::vec::Vec; +use core::mem::size_of; + +/// Meta writer structures for PE +/// PE is a complicated format that requires meta knowledge about all its fields +/// and reorganization at write time as we cannot predict all fields based on local information. +/// This file contains global structure which possess the global information to make up +/// for the complexity of PE. +/// Heavily inspired of how LLVM objcopy works for COFF. +use log::debug; +use log::trace; +use scroll::Pread; +use scroll::Pwrite; + +use crate::pe::certificate_table::enumerate_certificates; +use crate::pe::data_directories::SIZEOF_DATA_DIRECTORY; +use crate::pe::header::SIZEOF_COFF_HEADER; +use crate::pe::optional_header::StandardFields32; +use crate::pe::optional_header::StandardFields64; +use crate::pe::optional_header::WindowsFields32; +use crate::pe::optional_header::WindowsFields64; + +use crate::pe::utils::align_to; + +use super::data_directories::DataDirectory; +use super::data_directories::DataDirectoryType; +use super::debug::ImageDebugDirectory; +use super::error; +use super::header::DosHeader; +use super::header::DosStub; +use super::optional_header::OptionalHeader; +use super::section_table::Section; +use super::section_table::SectionTable; +use super::section_table::IMAGE_SCN_CNT_INITIALIZED_DATA; +use super::section_table::IMAGE_SCN_MEM_EXECUTE; +use super::section_table::IMAGE_SCN_MEM_READ; +use super::section_table::IMAGE_SCN_MEM_WRITE; +use super::utils::is_in_range; +use super::utils::rva2offset; +use super::PE; + +// The maximum number of sections that a COFF object can have (inclusive) +// which is a strict limit for PE, taken from LLVM. +const MAX_NUMBER_OF_SECTIONS_PE: usize = 65279; + +pub struct PEWriter<'a, 'b> { + pe: PE<'a>, + file_size: u32, + file_alignment: u32, + section_alignment: u32, + size_of_initialized_data: u64, + pending_sections: Vec>, + ready_sections: Vec>, +} + +impl<'a, 'b> PEWriter<'a, 'b> { + /// Consume the PE and store on-the-side information to rewrite + /// this PE with new information, e.g. new sections. + /// Some data can be manipulated beforehand and will be correctly rewritten + /// but this is very driven by implementation details. + /// It is guaranteed to work for new sections and removed sections, not for much more. + pub fn new(pe: PE<'a>) -> error::Result { + let header = pe.header.optional_header.ok_or(error::Error::Malformed( + "Missing optional header, write is not supported in this usecase".into(), + ))?; + Ok(Self { + pe, + file_size: 0, + file_alignment: header.windows_fields.file_alignment, + section_alignment: header.windows_fields.section_alignment, + size_of_initialized_data: 0, + pending_sections: Vec::new(), + ready_sections: Vec::new(), + }) + } + + /// Enqueue a pending section to be laid out at write time. + /// Fields that are impossible to predict can be left out + /// and will be filled automatically. + pub fn insert_section(&mut self, mut section: Section<'b>) -> error::Result<()> { + // VA is needed only if characteristics is + // execute | read | write. + let need_virtual_address = (section.table.characteristics + & (IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE | IMAGE_SCN_MEM_EXECUTE)) + != 0; + + if let Some(contents) = §ion.contents { + debug_assert!(need_virtual_address, "contents present without any need for a virtual address; missing flag on characteristics?"); + section.table.virtual_size = contents.len().try_into()?; + let mut sections = self.pe.sections.clone(); + sections.sort_by_key(|sect| sect.virtual_address); + let last_section_offset = sections + .iter() + .chain(self.pending_sections.iter().map(|sect| §.table)) + .last() + .map(|last_section| last_section.virtual_address + last_section.virtual_size) + .ok_or(0u32) + .unwrap(); + + section.table.virtual_address = align_to(last_section_offset, self.section_alignment); + debug!( + "[section {:?}] virtual address assigned: {}", + section.table.name, section.table.virtual_address + ); + } + + self.pending_sections.push(section); + Ok(()) + } + + /// This will compute all the missing fields for a pending section + /// and put it inside the "ready" sections array for the writer + /// It relies on the global internal `self.file_size` and + /// `self.size_of_initialized_data` state to adjust the "on-disk" pointers. + fn layout_sections(&mut self) -> error::Result<()> { + fn layout_section( + file_size: &mut u32, + size_of_initialized_data: &mut u64, + header: &mut SectionTable, + data_length: usize, + n_relocations: usize, + file_alignment: u32, + ) -> error::Result<()> { + header.size_of_raw_data = align_to(data_length as u32, file_alignment); + if header.size_of_raw_data > 0 { + header.pointer_to_raw_data = *file_size; + } + + if n_relocations > 0 { + return Err(error::Error::Malformed( + "COFF are unsupported; PE should not have relocations!".into(), + )); + } + + *file_size += header.size_of_raw_data; + *file_size = align_to(*file_size, file_alignment); + + if header.characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA + == IMAGE_SCN_CNT_INITIALIZED_DATA + { + *size_of_initialized_data += header.size_of_raw_data as u64; + } + + Ok(()) + } + for section in &mut self.pe.sections { + layout_section( + &mut self.file_size, + &mut self.size_of_initialized_data, + section, + section.size_of_raw_data as usize, + section.number_of_relocations.into(), + self.file_alignment, + )?; + } + while !self.pending_sections.is_empty() { + let mut section = self.pending_sections.pop().unwrap(); + + layout_section( + &mut self.file_size, + &mut self.size_of_initialized_data, + &mut section.table, + section.contents.as_ref().map(|c| c.len()).unwrap_or(0), + section.relocations.len(), + self.file_alignment, + )?; + + self.ready_sections.push(section); + } + + // Sections were added in LIFO style. + // This means that the last element here is the first pending section. + // i.e. section with the lowest virtual address. + // To maintain the sorting invariant, we just need to reverse the list. + self.ready_sections.reverse(); + + Ok(()) + } + + fn layout_certificates(&mut self) -> error::Result { + let mut total_length = 0; + for certificate in &self.pe.certificates { + self.file_size += certificate.length; + total_length += certificate.length; + } + Ok(total_length) + } + + fn layout_data_directories_contents( + &mut self, + opt_header: &mut OptionalHeader, + ) -> error::Result<()> { + for (index, dir) in opt_header + .data_directories + .data_directories + .iter_mut() + .enumerate() + { + trace!("{}: {:?}", index, dir); + let dd_type: DataDirectoryType = index.try_into()?; + // skip certificate table, we don't use size here. + // skip the debug table, it must be ordered *after* the certificate table + // as per: + // > Another exception is that attribute certificate and debug information must be placed + // > at the very end of an image file, with the attribute certificate table immediately + // > preceding the debug section, because the loader does not map these into memory. The + // > rule about attribute certificate and debug information does not apply to object + // > files, however. + // https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#other-contents-of-the-file + if dd_type == DataDirectoryType::CertificateTable + || dd_type == DataDirectoryType::DebugTable + { + continue; + } + + if let Some((offset, dd)) = dir { + *offset = self.file_size as usize; + self.file_size += dd.size; + } + } + // 4 := certificate data directory + // it is special because it virtual size does not reflect the full size + // of attribute certificates available. + // virtual size is only the size of a single bundle of certificate. + if let Some((_cert_offset, cert_table)) = + opt_header.data_directories.data_directories[4].as_mut() + { + cert_table.virtual_address = self.file_size; + cert_table.size = self.layout_certificates()?; + } + // 6 := debug table + // it is special because if it exist, it must be *after* the certificate table. + // this is incorrect, the data directory offset must point at some section offset. + if let Some((debug_offset, debug_table)) = + opt_header.data_directories.data_directories[6].as_mut() + { + *debug_offset = self.file_size as usize; + self.file_size += debug_table.size; + } + Ok(()) + } + + fn finalize(&mut self) -> error::Result<()> { + // XXX(RaitoBezarius): some steps of finalization are "commented out" + // They would be necessary if you are planning to support those codepaths for COFF Object + // File write support, I do not want to support them, I will stop at supporting *PE + // executables*. + // 1. finalize symbol table + // FIXME: COFF are unsupported ; self.finalize_symbol_table()?; + // 2. finalize relocation targets + // FIXME: COFF are unsupported ; self.finalize_relocation_targets()?; + // 3. finalize symbol contents + // FIXME: COFF are unsupported ; self.finalize_symbol_contents()?; + // 4. compute the address of the new exe header + let mut size_of_headers: u32 = 0; + let pe_header_size: u32 = { + if self.pe.is_64 { + size_of::() as u32 + size_of::() as u32 + } else { + size_of::() as u32 + size_of::() as u32 + } + }; + self.pe.header.dos_header.pe_pointer = + (size_of::() + size_of::()) as u32; + debug_assert!( + self.pe.header.dos_header.pe_pointer >= 0x40, + "PE pointer < 0x40, this is not expected." + ); + // 5. compute the initial pe header size + let mut opt_header = self + .pe + .header + .optional_header + .ok_or(error::Error::Malformed( + "Missing optional header for a PE".into(), + ))?; + // Count data directories in the PE. + opt_header.windows_fields.number_of_rva_and_sizes = 16; // TODO(raito): opt_header.data_directories.dirs().count() as u32; is better but requires the write operation for DD to skip none dds. + size_of_headers += pe_header_size + + (SIZEOF_DATA_DIRECTORY as u32) * opt_header.windows_fields.number_of_rva_and_sizes; + // 6. compute the number of sections + self.pe.header.coff_header.number_of_sections = + (self.pe.sections.len() + self.pending_sections.len()) as u16; + size_of_headers += SIZEOF_COFF_HEADER as u32; + size_of_headers += (size_of::() as u32) + * (self.pe.header.coff_header.number_of_sections as u32); + size_of_headers = align_to(size_of_headers, self.file_alignment); + // 7. compute the optional header size + self.pe.header.coff_header.size_of_optional_header = u16::try_from(pe_header_size)? + + (SIZEOF_DATA_DIRECTORY as u16) + * u16::try_from(opt_header.windows_fields.number_of_rva_and_sizes)?; + // 8. set file size + self.file_size = size_of_headers; + self.size_of_initialized_data = 0; + // 9. layout all sections and data directories contents + self.layout_sections()?; + self.layout_data_directories_contents(&mut opt_header)?; + // 10. adjust PE specific headers w.r.t to sizes + opt_header.windows_fields.size_of_headers = size_of_headers; + opt_header.standard_fields.size_of_initialized_data = self.size_of_initialized_data; + + if let Some(last_section) = self + .pe + .sections + .iter() + .chain(self.ready_sections.iter().map(|s| &s.table)) + .last() + { + opt_header.windows_fields.size_of_image = align_to( + last_section.virtual_address + last_section.virtual_size, + self.section_alignment, + ); + } + + // Clear the checksum and do not compute it. + opt_header.windows_fields.check_sum = 0; + + // 11. FIXME: COFF are unsupported ; finalize string tables + + self.pe.header.optional_header = Some(opt_header); + self.file_size = align_to(self.file_size, self.file_alignment); + Ok(()) + } + + fn write_headers(&mut self, buf: &mut Vec) -> error::Result { + let offset = &mut 0; + // 1. write the header + debug!("writing this header: {:#?}", self.pe.header); + buf.gwrite(self.pe.header, offset)?; + // 2. write the section tables + for section in self + .pe + .sections + .iter() + .chain(self.ready_sections.iter().map(|s| &s.table)) + { + debug!( + "writing section table {} at {}", + section.name().unwrap_or("unknown name"), + offset + ); + buf.gwrite(section, offset)?; + } + + Ok(*offset) + } + + fn write_sections(&mut self, buf: &mut Vec) -> error::Result { + // For each section, seek at the pointer to raw data, write the contents. + // For executable sections, pad the remainder of the raw data size + // with 0xCC, because it's useful on x86 (debugger breakpoint). + let mut written = 0; + let ready_sections = core::mem::take(&mut self.ready_sections); + for section in self + .pe + .sections + .iter() + .cloned() + .map(|s| s.into_section(&self.pe.bytes).unwrap()) + .chain(ready_sections.into_iter()) + { + let offset = section.table.pointer_to_raw_data as usize; + if let Some(contents) = §ion.contents { + written += buf.pwrite(contents.as_ref(), offset)?; + debug!( + "wrote {} (true size: {}) contents at {}", + contents.len(), + section.table.size_of_raw_data, + offset + ); + if section.table.size_of_raw_data as usize > contents.len() { + written += buf.pwrite( + &vec![0xCC; (section.table.size_of_raw_data as usize) - contents.len()][..], + offset + contents.len(), + )?; + } + } + } + // FIXME: COFF are unsupported but you would need to write the relocations here and + // distinguish based on the size of the COFF object. + Ok(written) + } + + fn patch_debug_directory( + &mut self, + debug_directory: &DataDirectory, + w: &mut Vec, + ) -> error::Result { + if debug_directory.size == 0 { + return Ok(0); + } + + for section in &self.pe.sections { + let section_end = section.virtual_address + section.virtual_size; + + if is_in_range( + debug_directory.virtual_address as usize, + section.virtual_address as usize, + section_end as usize, + ) { + if debug_directory.virtual_address + debug_directory.size > section_end { + return Err(error::Error::Malformed( + "debug directory extends past end of section".into(), + )); + } + + // We compute the relative difference inside the section + let offset = debug_directory.virtual_address - section.virtual_address; + // We compute the pointer to raw data for the debug dir + // based on the on-disk offset section + relative diff + // as mapping is linear. + let mut target_offset = (section.pointer_to_raw_data + offset) as usize; + let end = target_offset + debug_directory.size as usize; + // Read until target_offset + debug_directory.size + while target_offset < end { + let mut debug_data: ImageDebugDirectory = + w.gread::(&mut target_offset)?; + if debug_data.pointer_to_raw_data != 0 { + debug_data.pointer_to_raw_data = + rva2offset(debug_data.address_of_raw_data as usize, section) + .try_into()?; + // We rewrite the previous pointer inside the memory buffer + // Right now, we are sitting potentially onto the next ImageDebugDirectory + // or the end. + // It is therefore enough to start from target_offset, go back to previous + // element and go to the relevant field immediately. + + w.pwrite( + debug_data.pointer_to_raw_data, + target_offset + 0x18 - size_of::(), + )?; + } + } + } + } + + Ok(0) + } + + pub fn write_into(&mut self) -> error::Result> { + let total_sections = self.pending_sections.len() + self.pe.sections.len(); + let is_too_large = total_sections >= MAX_NUMBER_OF_SECTIONS_PE; + + if is_too_large { + return Err(error::Error::Malformed( + format!("Trying to write {total_sections} sections, the limit is {MAX_NUMBER_OF_SECTIONS_PE} for a PE binary") + )); + } + + let mut written = 0; + + self.finalize()?; + debug!("finalized the new PE binary at {} bytes", self.file_size); + let mut buffer = vec![0; self.file_size as usize]; + + written += self.write_headers(&mut buffer)?; + debug!("wrote headers"); + written += self.write_sections(&mut buffer)?; + debug!("wrote sections"); + // FIXME: COFF are unsupported ; written += self.write_symbol_string_tables(&mut buffer)?; + if let Some((_, debug_dir)) = &self + .pe + .header + .optional_header + .and_then(|opt_header| opt_header.data_directories.data_directories[6]) + { + self.patch_debug_directory(debug_dir, &mut buffer)?; + debug!("patched debug directory"); + } + written += self + .pe + .write_data_directories( + &mut buffer[..], + &self.pe.header.optional_header.unwrap().data_directories, + scroll::LE, + )? + .0; + debug!("wrote data directories contents"); + + // Specification says that: + // if cert table and debug table exist, they must be at the very end in this order + // if cert table exist, it should be the last element + // if debug table exist, it should be the last element + // This is important because they are not mapped in memory. + // TODO: reintroduce it. + + // We cannot guarantee that written == self.file_size + // as PE cannot be perfectly efficient vs. how we do write them. + // For example, if you have 1 data directory and it is the last one, + // you will have to say that you have all data directories and will only write one data + // directory header, but your file size will reflect the potential size of all data + // directories contents. + // Of course, it is possible to improve many moving parts and make it quite efficient. + // PRs are welcome as correctness is already a good enough goal with PEs. + debug_assert!( + written <= self.file_size as usize, + "incorrect amount of bytes written, expected at most: {}, wrote: {}", + self.file_size, + written + ); + Ok(buffer) + } +}