pe: add write support for COFF object files

src/pe/header.rs

@@ -2,7 +2,7 @@ use crate::error;
 use crate::pe::{optional_header, section_table, symbol};
 use crate::strtab;
 use log::debug;
-use scroll::Pread;
+use scroll::{Pread, Pwrite, IOread, IOwrite, SizeWith};
 
 /// DOS header present in all PE binaries
 #[repr(C)]
@@ -29,7 +29,7 @@ impl DosHeader {
 
 /// COFF Header
 #[repr(C)]
-#[derive(Debug, PartialEq, Copy, Clone, Default)]
+#[derive(Debug, PartialEq, Copy, Clone, Default, Pread, Pwrite, IOread, IOwrite, SizeWith)]
 pub struct CoffHeader {
     /// The machine type
     pub machine: u16,
@@ -41,30 +41,16 @@ pub struct CoffHeader {
     pub characteristics: u16,
 }
 
-pub const SIZEOF_COFF_HEADER: usize = 24;
+pub const SIZEOF_COFF_HEADER: usize = 20;
 /// PE\0\0, little endian
-pub const COFF_MAGIC: u32 = 0x0000_4550;
+pub const PE_MAGIC: u32 = 0x0000_4550;
+pub const SIZEOF_PE_MAGIC: usize = 4;
 pub const COFF_MACHINE_X86: u16 = 0x14c;
 pub const COFF_MACHINE_X86_64: u16 = 0x8664;
 
 impl CoffHeader {
     pub fn parse(bytes: &[u8], offset: &mut usize) -> error::Result<Self> {
-        let mut coff = CoffHeader::default();
-        coff.machine = bytes.gread_with(offset, scroll::LE)
-            .map_err(|_| error::Error::Malformed(format!("cannot parse COFF machine (offset {:#x})", offset)))?;
-        coff.number_of_sections = bytes.gread_with(offset, scroll::LE)
-            .map_err(|_| error::Error::Malformed(format!("cannot parse COFF number of sections (offset {:#x})", offset)))?;
-        coff.time_date_stamp = bytes.gread_with(offset, scroll::LE)
-            .map_err(|_| error::Error::Malformed(format!("cannot parse COFF time date stamp (offset {:#x})", offset)))?;
-        coff.pointer_to_symbol_table = bytes.gread_with(offset, scroll::LE)
-            .map_err(|_| error::Error::Malformed(format!("cannot parse COFF pointer to symbol table (offset {:#x})", offset)))?;
-        coff.number_of_symbol_table = bytes.gread_with(offset, scroll::LE)
-            .map_err(|_| error::Error::Malformed(format!("cannot parse COFF number of symbol (offset {:#x})", offset)))?;
-        coff.size_of_optional_header = bytes.gread_with(offset, scroll::LE)
-            .map_err(|_| error::Error::Malformed(format!("cannot parse COFF size of optional header (offset {:#x})", offset)))?;
-        coff.characteristics = bytes.gread_with(offset, scroll::LE)
-            .map_err(|_| error::Error::Malformed(format!("cannot parse COFF characteristics (offset {:#x})", offset)))?;
-        Ok(coff)
+        Ok(bytes.gread_with(offset, scroll::LE)?)
     }
 
     /// Parse the COFF section headers.
@@ -138,7 +124,7 @@ impl Header {
 
 #[cfg(test)]
 mod tests {
-    use super::{DOS_MAGIC, COFF_MAGIC, COFF_MACHINE_X86, Header};
+    use super::{DOS_MAGIC, PE_MAGIC, COFF_MACHINE_X86, Header};
 
     const CRSS_HEADER: [u8; 688] =
         [0x4d, 0x5a, 0x90, 0x00, 0x03, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0xff, 0xff, 0x00, 0x00,
@@ -189,7 +175,7 @@ mod tests {
     fn crss_header () {
         let header = Header::parse(&&CRSS_HEADER[..]).unwrap();
         assert!(header.dos_header.signature == DOS_MAGIC);
-        assert!(header.signature == COFF_MAGIC);
+        assert!(header.signature == PE_MAGIC);
         assert!(header.coff_header.machine == COFF_MACHINE_X86);
         println!("header: {:?}", &header);
     }

src/pe/mod.rs

@@ -64,7 +64,7 @@ impl<'a> PE<'a> {
     pub fn parse(bytes: &'a [u8]) -> error::Result<Self> {
         let header = header::Header::parse(bytes)?;
         debug!("{:#?}", header);
-        let offset = &mut (header.dos_header.pe_pointer as usize + header::SIZEOF_COFF_HEADER + header.coff_header.size_of_optional_header as usize);
+        let offset = &mut (header.dos_header.pe_pointer as usize + header::SIZEOF_PE_MAGIC + header::SIZEOF_COFF_HEADER + header.coff_header.size_of_optional_header as usize);
         let sections = header.coff_header.sections(bytes, offset)?;
         let is_lib = characteristic::is_dll(header.coff_header.characteristics);
         let mut entry = 0;
@@ -160,6 +160,8 @@ impl<'a> Coff<'a> {
         let offset = &mut 0;
         let header = header::CoffHeader::parse(bytes, offset)?;
         debug!("{:#?}", header);
+        // TODO: maybe parse optional header, but it isn't present for Windows.
+        *offset += header.size_of_optional_header as usize;
         let sections = header.sections(bytes, offset)?;
         let symbols = header.symbols(bytes)?;
         let strings = header.strings(bytes)?;

src/pe/relocation.rs

@@ -1,5 +1,5 @@
 use crate::error;
-use scroll::{Pread, Pwrite};
+use scroll::{IOread, IOwrite, Pread, Pwrite, SizeWith};
 
 /// Size of a single COFF relocation.
 pub const COFF_RELOCATION_SIZE: usize = 10;
@@ -78,7 +78,7 @@ pub const IMAGE_REL_AMD64_SSPAN32: u16 = 0x0010;
 
 /// A COFF relocation.
 #[repr(C)]
-#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite)]
+#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite, IOread, IOwrite, SizeWith)]
 pub struct Relocation {
     /// The address of the item to which relocation is applied.
     ///

src/pe/section_table.rs

@@ -1,7 +1,8 @@
 use crate::alloc::string::{String, ToString};
-use scroll::Pread;
+use scroll::{ctx, Pread, Pwrite};
 use crate::error::{self, Error};
 use crate::pe::relocation;
+use std::io::Write;
 
 #[repr(C)]
 #[derive(Debug, PartialEq, Clone, Default)]
@@ -64,20 +65,56 @@ impl SectionTable {
         table.number_of_linenumbers = bytes.gread_with(offset, scroll::LE)?;
         table.characteristics = bytes.gread_with(offset, scroll::LE)?;
 
+        if let Some(idx) = table.name_offset()? {
+            table.real_name = Some(bytes.pread::<&str>(string_table_offset + idx)?.to_string());
+        }
+        Ok(table)
+    }
+
+    pub fn name_offset(&self) -> error::Result<Option<usize>> {
         // Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L1054
-        if name[0] == b'/' {
-            let idx: usize = if name[1] == b'/' {
-                let b64idx = name.pread::<&str>(2)?;
+        if self.name[0] == b'/' {
+            let idx: usize = if self.name[1] == b'/' {
+                let b64idx = self.name.pread::<&str>(2)?;
                 base64_decode_string_entry(b64idx).map_err(|_|
                     Error::Malformed(format!("Invalid indirect section name //{}: base64 decoding failed", b64idx)))?
             } else {
-                let name = name.pread::<&str>(1)?;
+                let name = self.name.pread::<&str>(1)?;
                 name.parse().map_err(|err|
                     Error::Malformed(format!("Invalid indirect section name /{}: {}", name, err)))?
             };
-            table.real_name = Some(bytes.pread::<&str>(string_table_offset + idx)?.to_string());
+            Ok(Some(idx))
+        } else {
+            Ok(None)
+        }
+    }
+
+    pub fn set_name_offset(&mut self, mut idx: usize) -> error::Result<()> {
+        if idx <= 9_999_999 { // 10^7 - 1
+            self.name = [0; 8];
+            self.name[0] = b'/';
+            write!(&mut self.name[1..], "{}", idx).unwrap();
+            Ok(())
+        } else if idx <= 0xfff_fff_fff { // 64^6 - 1
+            self.name[0] = b'/';
+            self.name[1] = b'/';
+            for i in 0..6 {
+                let rem = (idx % 64) as u8;
+                idx = idx / 64;
+                let c = match rem {
+                    0..=25 => b'A' + rem,
+                    26..=51 => b'a' + rem - 26,
+                    52..=61 => b'0' + rem - 52,
+                    62 => b'+',
+                    63 => b'/',
+                    _ => unreachable!(),
+                };
+                self.name[7 - i] = c;
+            }
+            Ok(())
+        } else {
+            Err(Error::Malformed(format!("Invalid section name offset: {}", idx)))
         }
-        Ok(table)
     }
 
     pub fn name(&self) -> error::Result<&str> {
@@ -94,6 +131,38 @@ impl SectionTable {
     }
 }
 
+impl ctx::SizeWith<scroll::Endian> for SectionTable {
+    type Units = usize;
+    fn size_with(_ctx: &scroll::Endian) -> usize {
+        SIZEOF_SECTION_TABLE
+    }
+}
+
+impl ctx::TryIntoCtx<scroll::Endian> for SectionTable {
+    type Error = error::Error;
+    type Size = usize;
+    fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result<Self::Size, Self::Error> {
+        let offset = &mut 0;
+        bytes.gwrite(&self.name[..], offset)?;
+        bytes.gwrite_with(self.virtual_size, offset, ctx)?;
+        bytes.gwrite_with(self.virtual_address, offset, ctx)?;
+        bytes.gwrite_with(self.size_of_raw_data, offset, ctx)?;
+        bytes.gwrite_with(self.pointer_to_raw_data, offset, ctx)?;
+        bytes.gwrite_with(self.pointer_to_relocations, offset, ctx)?;
+        bytes.gwrite_with(self.pointer_to_linenumbers, offset, ctx)?;
+        bytes.gwrite_with(self.number_of_relocations, offset, ctx)?;
+        bytes.gwrite_with(self.number_of_linenumbers, offset, ctx)?;
+        bytes.gwrite_with(self.characteristics, offset, ctx)?;
+        Ok(SIZEOF_SECTION_TABLE)
+    }
+}
+
+impl ctx::IntoCtx<scroll::Endian> for SectionTable {
+    fn into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) {
+        bytes.pwrite_with(self, 0, ctx).unwrap();
+    }
+}
+
 /// The section should not be padded to the next boundary. This flag is obsolete and is replaced
 /// by `IMAGE_SCN_ALIGN_1BYTES`. This is valid only for object files.
 pub const IMAGE_SCN_TYPE_NO_PAD: u32 = 0x0000_0008;
@@ -150,3 +219,25 @@ pub const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;
 pub const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;
 /// The section can be written to.
 pub const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn set_name_offset() {
+        let mut section = SectionTable::default();
+        for (offset, name) in vec![
+            (0, b"/0\0\0\0\0\0\0"),
+            (1, b"/1\0\0\0\0\0\0"),
+            (9_999_999, b"/9999999"),
+            (10_000_000, b"//AAmJaA"),
+            (0xfff_fff_fff, b"////////"),
+        ] {
+            section.set_name_offset(offset).unwrap();
+            assert_eq!(&section.name, name);
+            assert_eq!(section.name_offset().unwrap(), Some(offset));
+        }
+        assert!(section.set_name_offset(0x1_000_000_000).is_err());
+    }
+}

src/pe/symbol.rs

@@ -1,7 +1,7 @@
 use crate::error;
 use crate::strtab;
 use core::fmt::{self, Debug};
-use scroll::{ctx, Pread, Pwrite};
+use scroll::{ctx, IOread, IOwrite, Pread, Pwrite, SizeWith};
 
 /// Size of a single symbol in the COFF Symbol Table.
 pub const COFF_SYMBOL_SIZE: usize = 18;
@@ -16,6 +16,7 @@ pub const IMAGE_SYM_UNDEFINED: i16 = 0;
 pub const IMAGE_SYM_ABSOLUTE: i16 = -1;
 /// The symbol provides general type or debugging information but does not
 /// correspond to a section.
+pub const IMAGE_SYM_DEBUG: i16 = -2;
 
 // Base types for `Symbol::typ`.
 
@@ -162,7 +163,7 @@ pub const IMAGE_SYM_CLASS_CLR_TOKEN: u8 = 107;
 ///
 /// [`ExceptionData::get_unwind_info`]: struct.ExceptionData.html#method.get_unwind_info
 #[repr(C)]
-#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite)]
+#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite, IOread, IOwrite, SizeWith)]
 pub struct Symbol {
     /// The name of the symbol.
     ///
@@ -238,6 +239,12 @@ impl Symbol {
         }
     }
 
+    /// Set the strtab offset of the symbol name.
+    pub fn set_name_offset(&mut self, offset: u32) {
+        self.name[..4].copy_from_slice(&[0; 4]);
+        self.name.pwrite_with(offset, 4, scroll::LE).unwrap();
+    }
+
     /// Return the base type of the symbol.
     ///
     /// This type uses the `IMAGE_SYM_TYPE_*` definitions.
@@ -285,7 +292,7 @@ impl Symbol {
 
 /// Auxiliary symbol record for function definitions.
 #[repr(C)]
-#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite)]
+#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite, IOread, IOwrite, SizeWith)]
 pub struct AuxFunctionDefinition {
     /// The symbol-table index of the corresponding `.bf` (begin function) symbol record.
     pub tag_index: u32,
@@ -307,7 +314,7 @@ pub struct AuxFunctionDefinition {
 
 /// Auxiliary symbol record for symbols with storage class `IMAGE_SYM_CLASS_FUNCTION`.
 #[repr(C)]
-#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite)]
+#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite, IOread, IOwrite, SizeWith)]
 pub struct AuxBeginAndEndFunction {
     /// Unused padding.
     pub unused1: [u8; 4],
@@ -336,7 +343,7 @@ pub const IMAGE_WEAK_EXTERN_SEARCH_ALIAS: u32 = 3;
 
 /// Auxiliary symbol record for weak external symbols.
 #[repr(C)]
-#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite)]
+#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite, IOread, IOwrite, SizeWith)]
 pub struct AuxWeakExternal {
     /// The symbol-table index of the symbol to be linked if an external definition is not found.
     pub tag_index: u32,
@@ -376,7 +383,7 @@ pub const IMAGE_COMDAT_SELECT_LARGEST: u8 = 6;
 
 /// Auxiliary symbol record for section definitions.
 #[repr(C)]
-#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite)]
+#[derive(Debug, Copy, Clone, PartialEq, Default, Pread, Pwrite, IOread, IOwrite, SizeWith)]
 pub struct AuxSectionDefinition {
     /// The size of section data; the same as `size_of_raw_data` in the section header.
     pub length: u32,