diff --git a/Cargo.toml b/Cargo.toml index 6cbbd0f18..c976939e4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -54,3 +54,6 @@ archive = ["alloc"] [badges.travis-ci] branch = "master" repository = "m4b/goblin" + +[dev-dependencies] +stderrlog = "0.5.4" diff --git a/examples/dotnet_pe_analysis.rs b/examples/dotnet_pe_analysis.rs index 6a87dfb60..b368b9687 100644 --- a/examples/dotnet_pe_analysis.rs +++ b/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(); diff --git a/examples/rewrite_pe.rs b/examples/rewrite_pe.rs new file mode 100644 index 000000000..45705a370 --- /dev/null +++ b/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 = 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() + ); +} diff --git a/src/error.rs b/src/error.rs index e2dd517e1..0ec4e1e0c 100644 --- a/src/error.rs +++ b/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 for Error { } } +impl From for Error { + fn from(err: TryFromIntError) -> Error { + Error::Malformed(format!("Integer do not fit: {err}")) + } +} + impl From for Error { fn from(err: scroll::Error) -> Error { Error::Scroll(err) diff --git a/src/pe/certificate_table.rs b/src/pe/certificate_table.rs index 4f4430d99..be45ce48d 100644 --- a/src/pe/certificate_table.rs +++ b/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 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 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 { + 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>; + pub(crate) fn enumerate_certificates( bytes: &[u8], table_virtual_address: u32, diff --git a/src/pe/data_directories.rs b/src/pe/data_directories.rs index 265e4e27f..e65db5953 100644 --- a/src/pe/data_directories.rs +++ b/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 { - 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 for DataDirectoryType { + type Error = error::Error; + fn try_from(value: usize) -> Result { + 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; NUM_DATA_DIRECTORIES], + pub data_directories: [Option<(usize, DataDirectory)>; NUM_DATA_DIRECTORIES], +} + +impl ctx::TryIntoCtx for DataDirectories { + type Error = error::Error; + + 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; + } 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 { - let idx = 0; - &self.data_directories[idx] - } - pub fn get_import_table(&self) -> &Option { - let idx = 1; - &self.data_directories[idx] - } - pub fn get_resource_table(&self) -> &Option { - let idx = 2; - &self.data_directories[idx] - } - pub fn get_exception_table(&self) -> &Option { - let idx = 3; - &self.data_directories[idx] - } - pub fn get_certificate_table(&self) -> &Option { - let idx = 4; - &self.data_directories[idx] - } - pub fn get_base_relocation_table(&self) -> &Option { - let idx = 5; - &self.data_directories[idx] - } - pub fn get_debug_table(&self) -> &Option { - let idx = 6; - &self.data_directories[idx] - } - pub fn get_architecture(&self) -> &Option { - let idx = 7; - &self.data_directories[idx] - } - pub fn get_global_ptr(&self) -> &Option { - let idx = 8; - &self.data_directories[idx] - } - pub fn get_tls_table(&self) -> &Option { - let idx = 9; - &self.data_directories[idx] - } - pub fn get_load_config_table(&self) -> &Option { - let idx = 10; - &self.data_directories[idx] - } - pub fn get_bound_import_table(&self) -> &Option { - let idx = 11; - &self.data_directories[idx] - } - pub fn get_import_address_table(&self) -> &Option { - let idx = 12; - &self.data_directories[idx] - } - pub fn get_delay_import_descriptor(&self) -> &Option { - let idx = 13; - &self.data_directories[idx] - } - pub fn get_clr_runtime_header(&self) -> &Option { - 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 { + self.data_directories + .into_iter() + .enumerate() + // (Index, Option
) -> 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))) } } diff --git a/src/pe/header.rs b/src/pe/header.rs index c1ded9dd9..06e23c0b0 100644 --- a/src/pe/header.rs +++ b/src/pe/header.rs @@ -3,24 +3,60 @@ use crate::pe::{optional_header, section_table, symbol}; use crate::strtab; use alloc::vec::Vec; use log::debug; -use scroll::{IOread, IOwrite, Pread, Pwrite, SizeWith}; +use scroll::{ctx, IOread, IOwrite, Pread, Pwrite, SizeWith}; /// DOS header present in all PE binaries #[repr(C)] -#[derive(Debug, PartialEq, Copy, Clone, Default)] +#[derive(Debug, PartialEq, Copy, Clone, Default, Pwrite)] pub struct DosHeader { /// Magic number: 5a4d pub signature: u16, - /// Pointer to PE header, always at offset 0x3c + /// e_cblp + pub bytes_on_last_page: u16, + /// e_cp + pub pages_in_file: u16, + /// e_crlc + pub relocations: u16, + /// e_cparhdr + pub size_of_header_in_paragraphs: u16, + /// e_minalloc + pub minimum_extra_paragraphs_needed: u16, + /// e_maxalloc + pub maximum_extra_paragraphs_needed: u16, + /// e_ss + pub initial_relative_ss: u16, + /// e_sp + pub initial_sp: u16, + /// e_csum + pub checksum: u16, + /// e_ip + pub initial_ip: u16, + /// e_cs + pub initial_relative_cs: u16, + /// e_lfarlc + pub file_address_of_relocation_table: u16, + /// e_ovno + pub overlay_number: u16, + /// e_res[4] + pub reserved: [u16; 4], + /// e_oemid + pub oem_id: u16, + /// e_oeminfo + pub oem_info: u16, + /// e_res2[10] + pub reserved2: [u16; 10], + /// e_lfanew: pointer to PE header, always at offset 0x3c pub pe_pointer: u32, } pub const DOS_MAGIC: u16 = 0x5a4d; pub const PE_POINTER_OFFSET: u32 = 0x3c; +pub const DOS_STUB_OFFSET: u32 = PE_POINTER_OFFSET + (core::mem::size_of::() as u32); impl DosHeader { pub fn parse(bytes: &[u8]) -> error::Result { - let signature = bytes.pread_with(0, scroll::LE).map_err(|_| { + let mut offset = 0; + let signature = bytes.gread_with(&mut offset, scroll::LE).map_err(|_| { error::Error::Malformed(format!("cannot parse DOS signature (offset {:#x})", 0)) })?; if signature != DOS_MAGIC { @@ -29,6 +65,33 @@ impl DosHeader { signature ))); } + + let bytes_on_last_page = bytes.gread_with(&mut offset, scroll::LE)?; + let pages_in_file = bytes.gread_with(&mut offset, scroll::LE)?; + let relocations = bytes.gread_with(&mut offset, scroll::LE)?; + let size_of_header_in_paragraphs = bytes.gread_with(&mut offset, scroll::LE)?; + let minimum_extra_paragraphs_needed = bytes.gread_with(&mut offset, scroll::LE)?; + let maximum_extra_paragraphs_needed = bytes.gread_with(&mut offset, scroll::LE)?; + let initial_relative_ss = bytes.gread_with(&mut offset, scroll::LE)?; + let initial_sp = bytes.gread_with(&mut offset, scroll::LE)?; + let checksum = bytes.gread_with(&mut offset, scroll::LE)?; + let initial_ip = bytes.gread_with(&mut offset, scroll::LE)?; + let initial_relative_cs = bytes.gread_with(&mut offset, scroll::LE)?; + let file_address_of_relocation_table = bytes.gread_with(&mut offset, scroll::LE)?; + let overlay_number = bytes.gread_with(&mut offset, scroll::LE)?; + let reserved = [0x0; 4]; + offset += core::mem::size_of_val(&reserved); + let oem_id = bytes.gread_with(&mut offset, scroll::LE)?; + let oem_info = bytes.gread_with(&mut offset, scroll::LE)?; + let reserved2 = [0x0; 10]; + offset += core::mem::size_of_val(&reserved2); + + debug_assert!( + offset == PE_POINTER_OFFSET as usize, + "expected offset ({:#x}) after reading DOS header to be at 0x3C", + offset + ); + let pe_pointer = bytes .pread_with(PE_POINTER_OFFSET as usize, scroll::LE) .map_err(|_| { @@ -37,6 +100,7 @@ impl DosHeader { PE_POINTER_OFFSET )) })?; + let pe_signature: u32 = bytes .pread_with(pe_pointer as usize, scroll::LE) @@ -52,13 +116,48 @@ impl DosHeader { pe_signature ))); } + Ok(DosHeader { signature, + bytes_on_last_page, + pages_in_file, + relocations, + size_of_header_in_paragraphs, + minimum_extra_paragraphs_needed, + maximum_extra_paragraphs_needed, + initial_relative_ss, + initial_sp, + checksum, + initial_ip, + initial_relative_cs, + file_address_of_relocation_table, + overlay_number, + reserved, + oem_id, + oem_info, + reserved2, pe_pointer, }) } } +#[repr(C)] +#[derive(Debug, PartialEq, Copy, Clone, Pread, Pwrite)] +/// The DOS stub program which should be executed in DOS mode +pub struct DosStub(pub [u8; 0x40]); +impl Default for DosStub { + fn default() -> Self { + // "This program cannot be run in DOS mode" error program + Self([ + 0x0E, 0x1F, 0xBA, 0x0E, 0x00, 0xB4, 0x09, 0xCD, 0x21, 0xB8, 0x01, 0x4C, 0xCD, 0x21, + 0x54, 0x68, 0x69, 0x73, 0x20, 0x70, 0x72, 0x6F, 0x67, 0x72, 0x61, 0x6D, 0x20, 0x63, + 0x61, 0x6E, 0x6E, 0x6F, 0x74, 0x20, 0x62, 0x65, 0x20, 0x72, 0x75, 0x6E, 0x20, 0x69, + 0x6E, 0x20, 0x44, 0x4F, 0x53, 0x20, 0x6D, 0x6F, 0x64, 0x65, 0x2E, 0x0D, 0x0D, 0x0A, + 0x24, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + ]) + } +} + /// COFF Header #[repr(C)] #[derive(Debug, PartialEq, Copy, Clone, Default, Pread, Pwrite, IOread, IOwrite, SizeWith)] @@ -163,14 +262,24 @@ impl CoffHeader { } /// Return the COFF symbol table. - pub fn symbols<'a>(&self, bytes: &'a [u8]) -> error::Result> { + pub fn symbols<'a>(&self, bytes: &'a [u8]) -> error::Result>> { let offset = self.pointer_to_symbol_table as usize; let number = self.number_of_symbol_table as usize; - symbol::SymbolTable::parse(bytes, offset, number) + if offset == 0 { + Ok(None) + } else { + symbol::SymbolTable::parse(bytes, offset, number).map(Some) + } } /// Return the COFF string table. - pub fn strings<'a>(&self, bytes: &'a [u8]) -> error::Result> { + pub fn strings<'a>(&self, bytes: &'a [u8]) -> error::Result>> { + // > 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. + if self.pointer_to_symbol_table == 0 { + return Ok(None); + } + let mut offset = self.pointer_to_symbol_table as usize + symbol::SymbolTable::size(self.number_of_symbol_table as usize); @@ -180,13 +289,15 @@ impl CoffHeader { // The offset needs to be advanced in order to read the strings. offset += length_field_size; - Ok(strtab::Strtab::parse(bytes, offset, length, 0)?) + Ok(Some(strtab::Strtab::parse(bytes, offset, length, 0)?)) } } #[derive(Debug, PartialEq, Copy, Clone, Default)] pub struct Header { pub dos_header: DosHeader, + /// DOS program for legacy loaders + pub dos_stub: DosStub, /// PE Magic: PE\0\0, little endian pub signature: u32, pub coff_header: CoffHeader, @@ -196,6 +307,12 @@ pub struct Header { impl Header { pub fn parse(bytes: &[u8]) -> error::Result { let dos_header = DosHeader::parse(&bytes)?; + let dos_stub = bytes.pread(DOS_STUB_OFFSET as usize).map_err(|_| { + error::Error::Malformed(format!( + "cannot parse DOS stub (offset {:#x})", + DOS_STUB_OFFSET + )) + })?; let mut offset = dos_header.pe_pointer as usize; let signature = bytes.gread_with(&mut offset, scroll::LE).map_err(|_| { error::Error::Malformed(format!("cannot parse PE signature (offset {:#x})", offset)) @@ -208,6 +325,7 @@ impl Header { }; Ok(Header { dos_header, + dos_stub, signature, coff_header, optional_header, @@ -215,6 +333,22 @@ impl Header { } } +impl ctx::TryIntoCtx for Header { + type Error = error::Error; + + fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result { + let offset = &mut 0; + bytes.gwrite_with(self.dos_header, offset, ctx)?; + bytes.gwrite_with(self.dos_stub, offset, ctx)?; + bytes.gwrite_with(self.signature, offset, scroll::LE)?; + bytes.gwrite_with(self.coff_header, offset, ctx)?; + if let Some(opt_header) = self.optional_header { + bytes.gwrite_with(opt_header, offset, ctx)?; + } + Ok(*offset) + } +} + /// Convert machine to str representation pub fn machine_to_str(machine: u16) -> &'static str { match machine { diff --git a/src/pe/mod.rs b/src/pe/mod.rs index c903f3904..935136198 100644 --- a/src/pe/mod.rs +++ b/src/pe/mod.rs @@ -3,7 +3,12 @@ // TODO: panics with unwrap on None for apisetschema.dll, fhuxgraphics.dll and some others +use core::cmp::max; + +use alloc::borrow::Cow; +use alloc::string::String; use alloc::vec::Vec; +use log::warn; pub mod authenticode; pub mod certificate_table; @@ -23,8 +28,11 @@ pub mod utils; use crate::container; use crate::error; +use crate::pe::utils::pad; use crate::strtab; +use scroll::{ctx, Pwrite}; + use log::debug; #[derive(Debug)] @@ -140,7 +148,7 @@ impl<'a> PE<'a> { entry, image_base, is_64 ); let file_alignment = optional_header.windows_fields.file_alignment; - if let Some(export_table) = *optional_header.data_directories.get_export_table() { + if let Some(&export_table) = optional_header.data_directories.get_export_table() { if let Ok(ed) = export::ExportData::parse_with_opts( bytes, export_table, @@ -162,7 +170,7 @@ impl<'a> PE<'a> { } } debug!("exports: {:#?}", exports); - if let Some(import_table) = *optional_header.data_directories.get_import_table() { + if let Some(&import_table) = optional_header.data_directories.get_import_table() { let id = if is_64 { import::ImportData::parse_with_opts::( bytes, @@ -196,7 +204,7 @@ impl<'a> PE<'a> { import_data = Some(id); } debug!("imports: {:#?}", imports); - if let Some(debug_table) = *optional_header.data_directories.get_debug_table() { + if let Some(&debug_table) = optional_header.data_directories.get_debug_table() { debug_data = Some(debug::DebugData::parse_with_opts( bytes, debug_table, @@ -209,8 +217,8 @@ impl<'a> PE<'a> { if header.coff_header.machine == header::COFF_MACHINE_X86_64 { // currently only x86_64 is supported debug!("exception data: {:#?}", exception_data); - if let Some(exception_table) = - *optional_header.data_directories.get_exception_table() + if let Some(&exception_table) = + optional_header.data_directories.get_exception_table() { exception_data = Some(exception::ExceptionData::parse_with_opts( bytes, @@ -224,8 +232,8 @@ impl<'a> PE<'a> { // Parse attribute certificates unless opted out of let certtable = if opts.parse_attribute_certificates { - if let Some(certificate_table) = - *optional_header.data_directories.get_certificate_table() + if let Some(&certificate_table) = + optional_header.data_directories.get_certificate_table() { certificates = certificate_table::enumerate_certificates( bytes, @@ -270,6 +278,192 @@ impl<'a> PE<'a> { certificates, }) } + + pub fn write_sections( + &self, + bytes: &mut [u8], + offset: &mut usize, + file_alignment: Option, + ctx: scroll::Endian, + ) -> Result { + // sections table and data + debug_assert!( + self.sections + .iter() + .flat_map(|section_a| { + self.sections + .iter() + .map(move |section_b| (section_a, section_b)) + }) + // given sections = (s_1, …, s_n) + // for all (s_i, s_j), i != j, verify that s_i does not overlap with s_j and vice versa. + .all(|(section_i, section_j)| section_i == section_j + || !section_i.overlaps_with(section_j)), + "Overlapping sections were found, this is not supported." + ); + + for section in &self.sections { + let section_data = section.data(&self.bytes)?.ok_or_else(|| { + error::Error::Malformed(format!( + "Section data `{}` is malformed", + section.name().unwrap_or("unknown name") + )) + })?; + let file_section_offset = + usize::try_from(section.pointer_to_raw_data).map_err(|_| { + error::Error::Malformed(format!( + "Section `{}`'s pointer to raw data does not fit in platform `usize`", + section.name().unwrap_or("unknown name") + )) + })?; + let vsize: usize = section.virtual_size.try_into()?; + let ondisk_size: usize = section.size_of_raw_data.try_into()?; + let section_name = String::from(section.name().unwrap_or("unknown name")); + + let mut file_offset = file_section_offset; + // `file_section_offset` is a on-disk offset which can be anywhere in the file. + // Write section data first to avoid the final consumption. + match section_data { + Cow::Borrowed(borrowed) => bytes.gwrite(borrowed, &mut file_offset)?, + Cow::Owned(owned) => bytes.gwrite(owned.as_slice(), &mut file_offset)?, + }; + + // Section tables follows the header. + bytes.gwrite_with(section, offset, ctx)?; + + // for size size_of_raw_data + // if < virtual_size, pad with 0 + // Pad with zeros if necessary + if file_offset < vsize { + bytes.gwrite(vec![0u8; vsize - file_offset].as_slice(), &mut file_offset)?; + } + + // Align on a boundary as per file alignement field. + if let Some(pad) = pad(file_offset - file_section_offset, file_alignment) { + debug!( + "aligning `{}` {:#x} -> {:#x} bytes'", + section_name, + file_offset - file_section_offset, + file_offset - file_section_offset + pad.len() + ); + bytes.gwrite(pad.as_slice(), &mut file_offset)?; + } + + let written_data_size = file_offset - file_section_offset; + if ondisk_size != written_data_size { + warn!("Original PE is inefficient or bug (on-disk data size in PE: {:#x}), we wrote {:#x} bytes", + ondisk_size, + written_data_size); + } + } + + Ok(*offset) + } + + pub fn write_certificates( + &self, + bytes: &mut [u8], + ctx: scroll::Endian, + ) -> Result { + let opt_header = self + .header + .optional_header + .ok_or(error::Error::Malformed(format!( + "This PE binary has no optional header; it is required to write certificates" + )))?; + let mut max_offset = 0; + + if let Some(certificate_directory) = opt_header.data_directories.get_certificate_table() { + let mut certificate_start = certificate_directory.virtual_address.try_into()?; + for certificate in &self.certificates { + bytes.gwrite_with(certificate, &mut certificate_start, ctx)?; + max_offset = max(certificate_start, max_offset); + } + } + + Ok(max_offset) + } +} + +impl<'a> ctx::TryIntoCtx for PE<'a> { + type Error = error::Error; + + fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result { + let mut offset = 0; + // We need to maintain a `max_offset` because + // we could be writing sections in the wrong order (i.e. not an increasing order for the + // pointer on raw disk) + // and there could be holes between sections. + // If we don't re-layout sections, we cannot fix that ourselves. + // Same can be said about the certificate table, there could be a hole between sections + // and the certificate data. + // To avoid those troubles, we maintain the max over all offsets we see so far. + let mut max_offset = 0; + let file_alignment: Option = match self.header.optional_header { + Some(opt_header) => { + debug_assert!( + opt_header.windows_fields.file_alignment.count_ones() == 1, + "file alignment should be a power of 2" + ); + Some(opt_header.windows_fields.file_alignment.try_into()?) + } + _ => None, + }; + bytes.gwrite_with(self.header, &mut offset, ctx)?; + max_offset = max(offset, max_offset); + self.write_sections(bytes, &mut offset, file_alignment, ctx)?; + // We want the section offset for which we have the highest pointer on disk. + // The next offset is reserved for debug tables (outside of sections) and/or certificate + // tables. + max_offset = max( + self.sections + .iter() + .max_by_key(|section| section.pointer_to_raw_data as usize) + .map(|section| (section.pointer_to_raw_data + section.size_of_raw_data) as usize) + .unwrap_or(offset), + max_offset, + ); + + // COFF Symbol Table + // Auxiliary Symbol Records + // COFF String Table + assert!( + self.header.coff_header.pointer_to_symbol_table == 0, + "Symbol tables in PE are deprecated and not supported to write" + ); + + // The following data directories are + // taken care inside a section: + // - 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) <- this one is special, it can be outside of a + // section. + // - 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) + // + // Nonetheless, we need to write the attribute certificate table one. + max_offset = max(max_offset, self.write_certificates(bytes, ctx)?); + + // TODO: we would like to support debug table outside of a section. + // i.e. debug tables that are never mapped in memory + // See https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#debug-directory-image-only + // > The debug directory can be in a discardable .debug section (if one exists), or it can be included in any other section in the image file, or not be in a section at all. + // In case it's not in a section at all, we need to find a way + // to rewrite it again. + // and we need to respect the ordering between attribute certificates + // and debug table. + + Ok(max_offset) + } } /// An analyzed COFF object @@ -279,10 +473,12 @@ pub struct Coff<'a> { pub header: header::CoffHeader, /// A list of the sections in this COFF binary pub sections: Vec, - /// The COFF symbol table. - pub symbols: symbol::SymbolTable<'a>, - /// The string table. - pub strings: strtab::Strtab<'a>, + /// The COFF symbol table, they are not guaranteed to exist. + /// For an image, this is expected to be None as COFF debugging information + /// has been deprecated. + pub symbols: Option>, + /// The string table, they don't exist if COFF symbol table does not exist. + pub strings: Option>, } impl<'a> Coff<'a> { @@ -419,7 +615,7 @@ mod tests { #[test] fn string_table_excludes_length() { let coff = Coff::parse(&&COFF_FILE_SINGLE_STRING_IN_STRING_TABLE[..]).unwrap(); - let string_table = coff.strings.to_vec().unwrap(); + let string_table = coff.strings.unwrap().to_vec().unwrap(); assert!(string_table == vec!["ExitProcess"]); } @@ -427,9 +623,10 @@ mod tests { #[test] fn symbol_name_excludes_length() { let coff = Coff::parse(&COFF_FILE_SINGLE_STRING_IN_STRING_TABLE).unwrap(); - let strings = coff.strings; + let strings = coff.strings.unwrap(); let symbols = coff .symbols + .unwrap() .iter() .filter(|(_, name, _)| name.is_none()) .map(|(_, _, sym)| sym.name(&strings).unwrap().to_owned()) diff --git a/src/pe/optional_header.rs b/src/pe/optional_header.rs index b1d319276..852f4dced 100644 --- a/src/pe/optional_header.rs +++ b/src/pe/optional_header.rs @@ -71,6 +71,22 @@ impl From for StandardFields { } } +impl From for StandardFields32 { + fn from(fields: StandardFields) -> Self { + StandardFields32 { + magic: fields.magic, + major_linker_version: fields.major_linker_version, + minor_linker_version: fields.minor_linker_version, + size_of_code: fields.size_of_code as u32, + size_of_initialized_data: fields.size_of_initialized_data as u32, + size_of_uninitialized_data: fields.size_of_uninitialized_data as u32, + address_of_entry_point: fields.address_of_entry_point as u32, + base_of_code: fields.base_of_code as u32, + base_of_data: fields.base_of_data, + } + } +} + impl From for StandardFields { fn from(fields: StandardFields64) -> Self { StandardFields { @@ -87,6 +103,21 @@ impl From for StandardFields { } } +impl From for StandardFields64 { + fn from(fields: StandardFields) -> Self { + StandardFields64 { + magic: fields.magic, + major_linker_version: fields.major_linker_version, + minor_linker_version: fields.minor_linker_version, + size_of_code: fields.size_of_code as u32, + size_of_initialized_data: fields.size_of_initialized_data as u32, + size_of_uninitialized_data: fields.size_of_uninitialized_data as u32, + address_of_entry_point: fields.address_of_entry_point as u32, + base_of_code: fields.base_of_code as u32, + } + } +} + /// Standard fields magic number for 32-bit binary pub const MAGIC_32: u16 = 0x10b; /// Standard fields magic number for 64-bit binary @@ -208,6 +239,36 @@ impl From for WindowsFields { } } +impl TryFrom for WindowsFields32 { + type Error = crate::error::Error; + + fn try_from(value: WindowsFields64) -> Result { + Ok(WindowsFields32 { + image_base: value.image_base.try_into()?, + section_alignment: value.section_alignment, + file_alignment: value.file_alignment, + major_operating_system_version: value.major_operating_system_version, + minor_operating_system_version: value.minor_operating_system_version, + major_image_version: value.major_image_version, + minor_image_version: value.minor_image_version, + major_subsystem_version: value.major_subsystem_version, + minor_subsystem_version: value.minor_subsystem_version, + win32_version_value: value.win32_version_value, + size_of_image: value.size_of_image, + size_of_headers: value.size_of_headers, + check_sum: value.check_sum, + subsystem: value.subsystem, + dll_characteristics: value.dll_characteristics, + size_of_stack_reserve: value.size_of_stack_reserve.try_into()?, + size_of_stack_commit: value.size_of_stack_commit.try_into()?, + size_of_heap_reserve: value.size_of_heap_reserve.try_into()?, + size_of_heap_commit: value.size_of_heap_commit.try_into()?, + loader_flags: value.loader_flags, + number_of_rva_and_sizes: value.number_of_rva_and_sizes, + }) + } +} + // impl From for WindowsFields { // fn from(windows: WindowsFields32) -> Self { // WindowsFields { @@ -289,6 +350,28 @@ impl<'a> ctx::TryFromCtx<'a, Endian> for OptionalHeader { } } +impl ctx::TryIntoCtx for OptionalHeader { + type Error = error::Error; + + fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result { + let offset = &mut 0; + match self.standard_fields.magic { + MAGIC_32 => { + bytes.gwrite_with::(self.standard_fields.into(), offset, ctx)?; + bytes.gwrite_with(WindowsFields32::try_from(self.windows_fields)?, offset, ctx)?; + bytes.gwrite_with(self.data_directories, offset, ctx)?; + } + MAGIC_64 => { + bytes.gwrite_with::(self.standard_fields.into(), offset, ctx)?; + bytes.gwrite_with(self.windows_fields, offset, ctx)?; + bytes.gwrite_with(self.data_directories, offset, ctx)?; + } + _ => panic!(), + } + Ok(*offset) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/pe/section_table.rs b/src/pe/section_table.rs index 4491827f1..f49c2b689 100644 --- a/src/pe/section_table.rs +++ b/src/pe/section_table.rs @@ -1,6 +1,8 @@ use crate::error::{self, Error}; use crate::pe::relocation; +use alloc::borrow::Cow; use alloc::string::{String, ToString}; +use alloc::vec::Vec; use scroll::{ctx, Pread, Pwrite}; #[repr(C)] @@ -79,6 +81,32 @@ impl SectionTable { Ok(table) } + 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`")) + })?; + + // assert!(self.virtual_size <= self.size_of_raw_data); + // if vsize > size_of_raw_data, the section is zero padded. + let section_end: usize = section_start + + usize::try_from(self.size_of_raw_data).map_err(|_| { + Error::Malformed(format!("Virtual size cannot fit in platform `usize`")) + })?; + + let original_bytes = pe_bytes.get(section_start..section_end).map(Cow::Borrowed); + + if original_bytes.is_some() && self.virtual_size > self.size_of_raw_data { + let mut bytes: Vec = Vec::new(); + bytes.resize(self.size_of_raw_data.try_into()?, 0); + bytes.copy_from_slice(&original_bytes.unwrap()); + bytes.resize(self.virtual_size.try_into()?, 0); + + Ok(Some(Cow::Owned(bytes))) + } else { + Ok(original_bytes) + } + } + pub fn name_offset(&self) -> error::Result> { // Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L1054 if self.name[0] == b'/' { @@ -163,6 +191,15 @@ impl SectionTable { let number = self.number_of_relocations as usize; relocation::Relocations::parse(bytes, offset, number) } + + /// Tests if `another_section` on-disk ranges will collide. + pub fn overlaps_with(&self, another_section: &SectionTable) -> bool { + let self_end = self.pointer_to_raw_data + self.size_of_raw_data; + let another_end = another_section.pointer_to_raw_data + another_section.size_of_raw_data; + + !((self_end <= another_section.pointer_to_raw_data) + || (another_end <= self.pointer_to_raw_data)) + } } impl ctx::SizeWith for SectionTable { @@ -171,7 +208,7 @@ impl ctx::SizeWith for SectionTable { } } -impl ctx::TryIntoCtx for SectionTable { +impl ctx::TryIntoCtx for &SectionTable { type Error = error::Error; fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result { let offset = &mut 0; @@ -189,7 +226,7 @@ impl ctx::TryIntoCtx for SectionTable { } } -impl ctx::IntoCtx for SectionTable { +impl ctx::IntoCtx for &SectionTable { fn into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) { bytes.pwrite_with(self, 0, ctx).unwrap(); } diff --git a/src/pe/symbol.rs b/src/pe/symbol.rs index b2ced808a..e5aba25be 100644 --- a/src/pe/symbol.rs +++ b/src/pe/symbol.rs @@ -412,6 +412,7 @@ pub struct AuxSectionDefinition { } /// A COFF symbol table. +// TODO: #[derive(Pwrite)] produce unparseable tokens pub struct SymbolTable<'a> { symbols: &'a [u8], } @@ -483,6 +484,14 @@ impl<'a> SymbolTable<'a> { } } +impl<'a> ctx::TryIntoCtx for SymbolTable<'a> { + type Error = error::Error; + + fn try_into_ctx(self, bytes: &mut [u8], _ctx: scroll::Endian) -> Result { + bytes.pwrite(self.symbols, 0).map_err(|err| err.into()) + } +} + impl<'a> Debug for SymbolTable<'a> { fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { fmt.debug_struct("SymbolTable") diff --git a/src/pe/utils.rs b/src/pe/utils.rs index 07a320d57..289ccc529 100644 --- a/src/pe/utils.rs +++ b/src/pe/utils.rs @@ -1,5 +1,6 @@ use crate::error; use alloc::string::ToString; +use alloc::vec::Vec; use scroll::Pread; use super::options; @@ -178,3 +179,18 @@ where let result: T = bytes.pread_with(offset, scroll::LE)?; Ok(result) } + +pub(crate) fn pad(length: usize, alignment: Option) -> Option> { + match alignment { + Some(alignment) => { + let overhang = length % alignment; + if overhang != 0 { + let repeat = alignment - overhang; + Some(vec![0u8; repeat]) + } else { + None + } + } + None => None, + } +} diff --git a/src/strtab.rs b/src/strtab.rs index dc7b8080f..487d8d427 100644 --- a/src/strtab.rs +++ b/src/strtab.rs @@ -34,6 +34,11 @@ impl<'a> Strtab<'a> { Self::from_slice_unparsed(bytes, 0, bytes.len(), delim) } + /// Returns the length of this `Strtab` in bytes + pub fn len(&self) -> usize { + self.bytes.len() + } + /// Creates a `Strtab` directly without bounds check and without parsing it. /// /// This is potentially unsafe and should only be used if `feature = "alloc"` is disabled.