1 //! Helper for writing PE files. 2 use std::mem; 3 use std::string::String; 4 use std::vec::Vec; 5 6 use crate::endian::{LittleEndian as LE, *}; 7 use crate::pe; 8 use crate::write::util; 9 use crate::write::{Error, Result, WritableBuffer}; 10 11 /// A helper for writing PE files. 12 /// 13 /// Writing uses a two phase approach. The first phase reserves file ranges and virtual 14 /// address ranges for everything in the order that they will be written. 15 /// 16 /// The second phase writes everything out in order. Thus the caller must ensure writing 17 /// is in the same order that file ranges were reserved. 18 #[allow(missing_debug_implementations)] 19 pub struct Writer<'a> { 20 is_64: bool, 21 section_alignment: u32, 22 file_alignment: u32, 23 24 buffer: &'a mut dyn WritableBuffer, 25 len: u32, 26 virtual_len: u32, 27 headers_len: u32, 28 29 code_address: u32, 30 data_address: u32, 31 code_len: u32, 32 data_len: u32, 33 bss_len: u32, 34 35 nt_headers_offset: u32, 36 data_directories: Vec<DataDirectory>, 37 section_header_num: u16, 38 sections: Vec<Section>, 39 40 symbol_offset: u32, 41 symbol_num: u32, 42 43 reloc_blocks: Vec<RelocBlock>, 44 relocs: Vec<U16<LE>>, 45 reloc_offset: u32, 46 } 47 48 impl<'a> Writer<'a> { 49 /// Create a new `Writer`. new( is_64: bool, section_alignment: u32, file_alignment: u32, buffer: &'a mut dyn WritableBuffer, ) -> Self50 pub fn new( 51 is_64: bool, 52 section_alignment: u32, 53 file_alignment: u32, 54 buffer: &'a mut dyn WritableBuffer, 55 ) -> Self { 56 Writer { 57 is_64, 58 section_alignment, 59 file_alignment, 60 61 buffer, 62 len: 0, 63 virtual_len: 0, 64 headers_len: 0, 65 66 code_address: 0, 67 data_address: 0, 68 code_len: 0, 69 data_len: 0, 70 bss_len: 0, 71 72 nt_headers_offset: 0, 73 data_directories: Vec::new(), 74 section_header_num: 0, 75 sections: Vec::new(), 76 77 symbol_offset: 0, 78 symbol_num: 0, 79 80 reloc_blocks: Vec::new(), 81 relocs: Vec::new(), 82 reloc_offset: 0, 83 } 84 } 85 86 /// Return the current virtual address size that has been reserved. 87 /// 88 /// This is only valid after section headers have been reserved. virtual_len(&self) -> u3289 pub fn virtual_len(&self) -> u32 { 90 self.virtual_len 91 } 92 93 /// Reserve a virtual address range with the given size. 94 /// 95 /// The reserved length will be increased to match the section alignment. 96 /// 97 /// Returns the aligned offset of the start of the range. reserve_virtual(&mut self, len: u32) -> u3298 pub fn reserve_virtual(&mut self, len: u32) -> u32 { 99 let offset = self.virtual_len; 100 self.virtual_len += len; 101 self.virtual_len = util::align_u32(self.virtual_len, self.section_alignment); 102 offset 103 } 104 105 /// Reserve up to the given virtual address. 106 /// 107 /// The reserved length will be increased to match the section alignment. reserve_virtual_until(&mut self, address: u32)108 pub fn reserve_virtual_until(&mut self, address: u32) { 109 debug_assert!(self.virtual_len <= address); 110 self.virtual_len = util::align_u32(address, self.section_alignment); 111 } 112 113 /// Return the current file length that has been reserved. reserved_len(&self) -> u32114 pub fn reserved_len(&self) -> u32 { 115 self.len 116 } 117 118 /// Return the current file length that has been written. 119 #[allow(clippy::len_without_is_empty)] len(&self) -> usize120 pub fn len(&self) -> usize { 121 self.buffer.len() 122 } 123 124 /// Reserve a file range with the given size and starting alignment. 125 /// 126 /// Returns the aligned offset of the start of the range. reserve(&mut self, len: u32, align_start: u32) -> u32127 pub fn reserve(&mut self, len: u32, align_start: u32) -> u32 { 128 if len == 0 { 129 return self.len; 130 } 131 self.reserve_align(align_start); 132 let offset = self.len; 133 self.len += len; 134 offset 135 } 136 137 /// Reserve a file range with the given size and using the file alignment. 138 /// 139 /// Returns the aligned offset of the start of the range. reserve_file(&mut self, len: u32) -> u32140 pub fn reserve_file(&mut self, len: u32) -> u32 { 141 self.reserve(len, self.file_alignment) 142 } 143 144 /// Write data. write(&mut self, data: &[u8])145 pub fn write(&mut self, data: &[u8]) { 146 self.buffer.write_bytes(data); 147 } 148 149 /// Reserve alignment padding bytes. reserve_align(&mut self, align_start: u32)150 pub fn reserve_align(&mut self, align_start: u32) { 151 self.len = util::align_u32(self.len, align_start); 152 } 153 154 /// Write alignment padding bytes. write_align(&mut self, align_start: u32)155 pub fn write_align(&mut self, align_start: u32) { 156 util::write_align(self.buffer, align_start as usize); 157 } 158 159 /// Reserve the file range up to the given file offset. reserve_until(&mut self, offset: u32)160 pub fn reserve_until(&mut self, offset: u32) { 161 debug_assert!(self.len <= offset); 162 self.len = offset; 163 } 164 165 /// Write padding up to the given file offset. pad_until(&mut self, offset: u32)166 pub fn pad_until(&mut self, offset: u32) { 167 debug_assert!(self.buffer.len() <= offset as usize); 168 self.buffer.resize(offset as usize); 169 } 170 171 /// Reserve the range for the DOS header. 172 /// 173 /// This must be at the start of the file. 174 /// 175 /// When writing, you may use `write_custom_dos_header` or `write_empty_dos_header`. reserve_dos_header(&mut self)176 pub fn reserve_dos_header(&mut self) { 177 debug_assert_eq!(self.len, 0); 178 self.reserve(mem::size_of::<pe::ImageDosHeader>() as u32, 1); 179 } 180 181 /// Write a custom DOS header. 182 /// 183 /// This must be at the start of the file. write_custom_dos_header(&mut self, dos_header: &pe::ImageDosHeader) -> Result<()>184 pub fn write_custom_dos_header(&mut self, dos_header: &pe::ImageDosHeader) -> Result<()> { 185 debug_assert_eq!(self.buffer.len(), 0); 186 187 // Start writing. 188 self.buffer 189 .reserve(self.len as usize) 190 .map_err(|_| Error(String::from("Cannot allocate buffer")))?; 191 192 self.buffer.write(dos_header); 193 Ok(()) 194 } 195 196 /// Write the DOS header for a file without a stub. 197 /// 198 /// This must be at the start of the file. 199 /// 200 /// Uses default values for all fields. write_empty_dos_header(&mut self) -> Result<()>201 pub fn write_empty_dos_header(&mut self) -> Result<()> { 202 self.write_custom_dos_header(&pe::ImageDosHeader { 203 e_magic: U16::new(LE, pe::IMAGE_DOS_SIGNATURE), 204 e_cblp: U16::new(LE, 0), 205 e_cp: U16::new(LE, 0), 206 e_crlc: U16::new(LE, 0), 207 e_cparhdr: U16::new(LE, 0), 208 e_minalloc: U16::new(LE, 0), 209 e_maxalloc: U16::new(LE, 0), 210 e_ss: U16::new(LE, 0), 211 e_sp: U16::new(LE, 0), 212 e_csum: U16::new(LE, 0), 213 e_ip: U16::new(LE, 0), 214 e_cs: U16::new(LE, 0), 215 e_lfarlc: U16::new(LE, 0), 216 e_ovno: U16::new(LE, 0), 217 e_res: [U16::new(LE, 0); 4], 218 e_oemid: U16::new(LE, 0), 219 e_oeminfo: U16::new(LE, 0), 220 e_res2: [U16::new(LE, 0); 10], 221 e_lfanew: U32::new(LE, self.nt_headers_offset), 222 }) 223 } 224 225 /// Reserve a fixed DOS header and stub. 226 /// 227 /// Use `reserve_dos_header` and `reserve` if you need a custom stub. reserve_dos_header_and_stub(&mut self)228 pub fn reserve_dos_header_and_stub(&mut self) { 229 self.reserve_dos_header(); 230 self.reserve(64, 1); 231 } 232 233 /// Write a fixed DOS header and stub. 234 /// 235 /// Use `write_custom_dos_header` and `write` if you need a custom stub. write_dos_header_and_stub(&mut self) -> Result<()>236 pub fn write_dos_header_and_stub(&mut self) -> Result<()> { 237 self.write_custom_dos_header(&pe::ImageDosHeader { 238 e_magic: U16::new(LE, pe::IMAGE_DOS_SIGNATURE), 239 e_cblp: U16::new(LE, 0x90), 240 e_cp: U16::new(LE, 3), 241 e_crlc: U16::new(LE, 0), 242 e_cparhdr: U16::new(LE, 4), 243 e_minalloc: U16::new(LE, 0), 244 e_maxalloc: U16::new(LE, 0xffff), 245 e_ss: U16::new(LE, 0), 246 e_sp: U16::new(LE, 0xb8), 247 e_csum: U16::new(LE, 0), 248 e_ip: U16::new(LE, 0), 249 e_cs: U16::new(LE, 0), 250 e_lfarlc: U16::new(LE, 0x40), 251 e_ovno: U16::new(LE, 0), 252 e_res: [U16::new(LE, 0); 4], 253 e_oemid: U16::new(LE, 0), 254 e_oeminfo: U16::new(LE, 0), 255 e_res2: [U16::new(LE, 0); 10], 256 e_lfanew: U32::new(LE, self.nt_headers_offset), 257 })?; 258 259 #[rustfmt::skip] 260 self.buffer.write_bytes(&[ 261 0x0e, 0x1f, 0xba, 0x0e, 0x00, 0xb4, 0x09, 0xcd, 262 0x21, 0xb8, 0x01, 0x4c, 0xcd, 0x21, 0x54, 0x68, 263 0x69, 0x73, 0x20, 0x70, 0x72, 0x6f, 0x67, 0x72, 264 0x61, 0x6d, 0x20, 0x63, 0x61, 0x6e, 0x6e, 0x6f, 265 0x74, 0x20, 0x62, 0x65, 0x20, 0x72, 0x75, 0x6e, 266 0x20, 0x69, 0x6e, 0x20, 0x44, 0x4f, 0x53, 0x20, 267 0x6d, 0x6f, 0x64, 0x65, 0x2e, 0x0d, 0x0d, 0x0a, 268 0x24, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 269 ]); 270 271 Ok(()) 272 } 273 nt_headers_size(&self) -> u32274 fn nt_headers_size(&self) -> u32 { 275 if self.is_64 { 276 mem::size_of::<pe::ImageNtHeaders64>() as u32 277 } else { 278 mem::size_of::<pe::ImageNtHeaders32>() as u32 279 } 280 } 281 optional_header_size(&self) -> u32282 fn optional_header_size(&self) -> u32 { 283 let size = if self.is_64 { 284 mem::size_of::<pe::ImageOptionalHeader64>() as u32 285 } else { 286 mem::size_of::<pe::ImageOptionalHeader32>() as u32 287 }; 288 size + self.data_directories.len() as u32 * mem::size_of::<pe::ImageDataDirectory>() as u32 289 } 290 291 /// Return the offset of the NT headers, if reserved. nt_headers_offset(&self) -> u32292 pub fn nt_headers_offset(&self) -> u32 { 293 self.nt_headers_offset 294 } 295 296 /// Reserve the range for the NT headers. reserve_nt_headers(&mut self, data_directory_num: usize)297 pub fn reserve_nt_headers(&mut self, data_directory_num: usize) { 298 debug_assert_eq!(self.nt_headers_offset, 0); 299 self.nt_headers_offset = self.reserve(self.nt_headers_size(), 8); 300 self.data_directories = vec![DataDirectory::default(); data_directory_num]; 301 self.reserve( 302 data_directory_num as u32 * mem::size_of::<pe::ImageDataDirectory>() as u32, 303 1, 304 ); 305 } 306 307 /// Set the virtual address and size of a data directory. set_data_directory(&mut self, index: usize, virtual_address: u32, size: u32)308 pub fn set_data_directory(&mut self, index: usize, virtual_address: u32, size: u32) { 309 self.data_directories[index] = DataDirectory { 310 virtual_address, 311 size, 312 } 313 } 314 315 /// Write the NT headers. write_nt_headers(&mut self, nt_headers: NtHeaders)316 pub fn write_nt_headers(&mut self, nt_headers: NtHeaders) { 317 self.pad_until(self.nt_headers_offset); 318 self.buffer.write(&U32::new(LE, pe::IMAGE_NT_SIGNATURE)); 319 let file_header = pe::ImageFileHeader { 320 machine: U16::new(LE, nt_headers.machine), 321 number_of_sections: U16::new(LE, self.section_header_num), 322 time_date_stamp: U32::new(LE, nt_headers.time_date_stamp), 323 pointer_to_symbol_table: U32::new(LE, self.symbol_offset), 324 number_of_symbols: U32::new(LE, self.symbol_num), 325 size_of_optional_header: U16::new(LE, self.optional_header_size() as u16), 326 characteristics: U16::new(LE, nt_headers.characteristics), 327 }; 328 self.buffer.write(&file_header); 329 if self.is_64 { 330 let optional_header = pe::ImageOptionalHeader64 { 331 magic: U16::new(LE, pe::IMAGE_NT_OPTIONAL_HDR64_MAGIC), 332 major_linker_version: nt_headers.major_linker_version, 333 minor_linker_version: nt_headers.minor_linker_version, 334 size_of_code: U32::new(LE, self.code_len), 335 size_of_initialized_data: U32::new(LE, self.data_len), 336 size_of_uninitialized_data: U32::new(LE, self.bss_len), 337 address_of_entry_point: U32::new(LE, nt_headers.address_of_entry_point), 338 base_of_code: U32::new(LE, self.code_address), 339 image_base: U64::new(LE, nt_headers.image_base), 340 section_alignment: U32::new(LE, self.section_alignment), 341 file_alignment: U32::new(LE, self.file_alignment), 342 major_operating_system_version: U16::new( 343 LE, 344 nt_headers.major_operating_system_version, 345 ), 346 minor_operating_system_version: U16::new( 347 LE, 348 nt_headers.minor_operating_system_version, 349 ), 350 major_image_version: U16::new(LE, nt_headers.major_image_version), 351 minor_image_version: U16::new(LE, nt_headers.minor_image_version), 352 major_subsystem_version: U16::new(LE, nt_headers.major_subsystem_version), 353 minor_subsystem_version: U16::new(LE, nt_headers.minor_subsystem_version), 354 win32_version_value: U32::new(LE, 0), 355 size_of_image: U32::new(LE, self.virtual_len), 356 size_of_headers: U32::new(LE, self.headers_len), 357 check_sum: U32::new(LE, 0), 358 subsystem: U16::new(LE, nt_headers.subsystem), 359 dll_characteristics: U16::new(LE, nt_headers.dll_characteristics), 360 size_of_stack_reserve: U64::new(LE, nt_headers.size_of_stack_reserve), 361 size_of_stack_commit: U64::new(LE, nt_headers.size_of_stack_commit), 362 size_of_heap_reserve: U64::new(LE, nt_headers.size_of_heap_reserve), 363 size_of_heap_commit: U64::new(LE, nt_headers.size_of_heap_commit), 364 loader_flags: U32::new(LE, 0), 365 number_of_rva_and_sizes: U32::new(LE, self.data_directories.len() as u32), 366 }; 367 self.buffer.write(&optional_header); 368 } else { 369 let optional_header = pe::ImageOptionalHeader32 { 370 magic: U16::new(LE, pe::IMAGE_NT_OPTIONAL_HDR32_MAGIC), 371 major_linker_version: nt_headers.major_linker_version, 372 minor_linker_version: nt_headers.minor_linker_version, 373 size_of_code: U32::new(LE, self.code_len), 374 size_of_initialized_data: U32::new(LE, self.data_len), 375 size_of_uninitialized_data: U32::new(LE, self.bss_len), 376 address_of_entry_point: U32::new(LE, nt_headers.address_of_entry_point), 377 base_of_code: U32::new(LE, self.code_address), 378 base_of_data: U32::new(LE, self.data_address), 379 image_base: U32::new(LE, nt_headers.image_base as u32), 380 section_alignment: U32::new(LE, self.section_alignment), 381 file_alignment: U32::new(LE, self.file_alignment), 382 major_operating_system_version: U16::new( 383 LE, 384 nt_headers.major_operating_system_version, 385 ), 386 minor_operating_system_version: U16::new( 387 LE, 388 nt_headers.minor_operating_system_version, 389 ), 390 major_image_version: U16::new(LE, nt_headers.major_image_version), 391 minor_image_version: U16::new(LE, nt_headers.minor_image_version), 392 major_subsystem_version: U16::new(LE, nt_headers.major_subsystem_version), 393 minor_subsystem_version: U16::new(LE, nt_headers.minor_subsystem_version), 394 win32_version_value: U32::new(LE, 0), 395 size_of_image: U32::new(LE, self.virtual_len), 396 size_of_headers: U32::new(LE, self.headers_len), 397 check_sum: U32::new(LE, 0), 398 subsystem: U16::new(LE, nt_headers.subsystem), 399 dll_characteristics: U16::new(LE, nt_headers.dll_characteristics), 400 size_of_stack_reserve: U32::new(LE, nt_headers.size_of_stack_reserve as u32), 401 size_of_stack_commit: U32::new(LE, nt_headers.size_of_stack_commit as u32), 402 size_of_heap_reserve: U32::new(LE, nt_headers.size_of_heap_reserve as u32), 403 size_of_heap_commit: U32::new(LE, nt_headers.size_of_heap_commit as u32), 404 loader_flags: U32::new(LE, 0), 405 number_of_rva_and_sizes: U32::new(LE, self.data_directories.len() as u32), 406 }; 407 self.buffer.write(&optional_header); 408 } 409 410 for dir in &self.data_directories { 411 self.buffer.write(&pe::ImageDataDirectory { 412 virtual_address: U32::new(LE, dir.virtual_address), 413 size: U32::new(LE, dir.size), 414 }) 415 } 416 } 417 418 /// Reserve the section headers. 419 /// 420 /// The number of reserved section headers must be the same as the number of sections that 421 /// are later reserved. 422 // TODO: change this to a maximum number of sections? reserve_section_headers(&mut self, section_header_num: u16)423 pub fn reserve_section_headers(&mut self, section_header_num: u16) { 424 debug_assert_eq!(self.section_header_num, 0); 425 self.section_header_num = section_header_num; 426 self.reserve( 427 u32::from(section_header_num) * mem::size_of::<pe::ImageSectionHeader>() as u32, 428 1, 429 ); 430 // Padding before sections must be included in headers_len. 431 self.reserve_align(self.file_alignment); 432 self.headers_len = self.len; 433 self.reserve_virtual(self.len); 434 } 435 436 /// Write the section headers. 437 /// 438 /// This uses information that was recorded when the sections were reserved. write_section_headers(&mut self)439 pub fn write_section_headers(&mut self) { 440 debug_assert_eq!(self.section_header_num as usize, self.sections.len()); 441 for section in &self.sections { 442 let section_header = pe::ImageSectionHeader { 443 name: section.name, 444 virtual_size: U32::new(LE, section.range.virtual_size), 445 virtual_address: U32::new(LE, section.range.virtual_address), 446 size_of_raw_data: U32::new(LE, section.range.file_size), 447 pointer_to_raw_data: U32::new(LE, section.range.file_offset), 448 pointer_to_relocations: U32::new(LE, 0), 449 pointer_to_linenumbers: U32::new(LE, 0), 450 number_of_relocations: U16::new(LE, 0), 451 number_of_linenumbers: U16::new(LE, 0), 452 characteristics: U32::new(LE, section.characteristics), 453 }; 454 self.buffer.write(§ion_header); 455 } 456 } 457 458 /// Reserve a section. 459 /// 460 /// Returns the file range and virtual address range that are reserved 461 /// for the section. reserve_section( &mut self, name: [u8; 8], characteristics: u32, virtual_size: u32, data_size: u32, ) -> SectionRange462 pub fn reserve_section( 463 &mut self, 464 name: [u8; 8], 465 characteristics: u32, 466 virtual_size: u32, 467 data_size: u32, 468 ) -> SectionRange { 469 let virtual_address = self.reserve_virtual(virtual_size); 470 471 // Padding after section must be included in section file size. 472 let file_size = util::align_u32(data_size, self.file_alignment); 473 let file_offset = if file_size != 0 { 474 self.reserve(file_size, self.file_alignment) 475 } else { 476 0 477 }; 478 479 // Sizes in optional header use the virtual size with the file alignment. 480 let aligned_virtual_size = util::align_u32(virtual_size, self.file_alignment); 481 if characteristics & pe::IMAGE_SCN_CNT_CODE != 0 { 482 if self.code_address == 0 { 483 self.code_address = virtual_address; 484 } 485 self.code_len += aligned_virtual_size; 486 } else if characteristics & pe::IMAGE_SCN_CNT_INITIALIZED_DATA != 0 { 487 if self.data_address == 0 { 488 self.data_address = virtual_address; 489 } 490 self.data_len += aligned_virtual_size; 491 } else if characteristics & pe::IMAGE_SCN_CNT_UNINITIALIZED_DATA != 0 { 492 if self.data_address == 0 { 493 self.data_address = virtual_address; 494 } 495 self.bss_len += aligned_virtual_size; 496 } 497 498 let range = SectionRange { 499 virtual_address, 500 virtual_size, 501 file_offset, 502 file_size, 503 }; 504 self.sections.push(Section { 505 name, 506 characteristics, 507 range, 508 }); 509 range 510 } 511 512 /// Write the data for a section. write_section(&mut self, offset: u32, data: &[u8])513 pub fn write_section(&mut self, offset: u32, data: &[u8]) { 514 if data.is_empty() { 515 return; 516 } 517 self.pad_until(offset); 518 self.write(data); 519 self.write_align(self.file_alignment); 520 } 521 522 /// Reserve a `.text` section. 523 /// 524 /// Contains executable code. reserve_text_section(&mut self, size: u32) -> SectionRange525 pub fn reserve_text_section(&mut self, size: u32) -> SectionRange { 526 self.reserve_section( 527 *b".text\0\0\0", 528 pe::IMAGE_SCN_CNT_CODE | pe::IMAGE_SCN_MEM_EXECUTE | pe::IMAGE_SCN_MEM_READ, 529 size, 530 size, 531 ) 532 } 533 534 /// Reserve a `.data` section. 535 /// 536 /// Contains initialized data. 537 /// 538 /// May also contain uninitialized data if `virtual_size` is greater than `data_size`. reserve_data_section(&mut self, virtual_size: u32, data_size: u32) -> SectionRange539 pub fn reserve_data_section(&mut self, virtual_size: u32, data_size: u32) -> SectionRange { 540 self.reserve_section( 541 *b".data\0\0\0", 542 pe::IMAGE_SCN_CNT_INITIALIZED_DATA | pe::IMAGE_SCN_MEM_READ | pe::IMAGE_SCN_MEM_WRITE, 543 virtual_size, 544 data_size, 545 ) 546 } 547 548 /// Reserve a `.rdata` section. 549 /// 550 /// Contains read-only initialized data. reserve_rdata_section(&mut self, size: u32) -> SectionRange551 pub fn reserve_rdata_section(&mut self, size: u32) -> SectionRange { 552 self.reserve_section( 553 *b".rdata\0\0", 554 pe::IMAGE_SCN_CNT_INITIALIZED_DATA | pe::IMAGE_SCN_MEM_READ, 555 size, 556 size, 557 ) 558 } 559 560 /// Reserve a `.bss` section. 561 /// 562 /// Contains uninitialized data. reserve_bss_section(&mut self, size: u32) -> SectionRange563 pub fn reserve_bss_section(&mut self, size: u32) -> SectionRange { 564 self.reserve_section( 565 *b".bss\0\0\0\0", 566 pe::IMAGE_SCN_CNT_UNINITIALIZED_DATA | pe::IMAGE_SCN_MEM_READ | pe::IMAGE_SCN_MEM_WRITE, 567 size, 568 0, 569 ) 570 } 571 572 /// Reserve an `.idata` section. 573 /// 574 /// Contains import tables. Note that it is permissible to store import tables in a different 575 /// section. 576 /// 577 /// This also sets the `pe::IMAGE_DIRECTORY_ENTRY_IMPORT` data directory. reserve_idata_section(&mut self, size: u32) -> SectionRange578 pub fn reserve_idata_section(&mut self, size: u32) -> SectionRange { 579 let range = self.reserve_section( 580 *b".idata\0\0", 581 pe::IMAGE_SCN_CNT_INITIALIZED_DATA | pe::IMAGE_SCN_MEM_READ | pe::IMAGE_SCN_MEM_WRITE, 582 size, 583 size, 584 ); 585 let dir = &mut self.data_directories[pe::IMAGE_DIRECTORY_ENTRY_IMPORT]; 586 debug_assert_eq!(dir.virtual_address, 0); 587 *dir = DataDirectory { 588 virtual_address: range.virtual_address, 589 size, 590 }; 591 range 592 } 593 594 /// Reserve an `.edata` section. 595 /// 596 /// Contains export tables. 597 /// 598 /// This also sets the `pe::IMAGE_DIRECTORY_ENTRY_EXPORT` data directory. reserve_edata_section(&mut self, size: u32) -> SectionRange599 pub fn reserve_edata_section(&mut self, size: u32) -> SectionRange { 600 let range = self.reserve_section( 601 *b".edata\0\0", 602 pe::IMAGE_SCN_CNT_INITIALIZED_DATA | pe::IMAGE_SCN_MEM_READ, 603 size, 604 size, 605 ); 606 let dir = &mut self.data_directories[pe::IMAGE_DIRECTORY_ENTRY_EXPORT]; 607 debug_assert_eq!(dir.virtual_address, 0); 608 *dir = DataDirectory { 609 virtual_address: range.virtual_address, 610 size, 611 }; 612 range 613 } 614 615 /// Reserve a `.pdata` section. 616 /// 617 /// Contains exception information. 618 /// 619 /// This also sets the `pe::IMAGE_DIRECTORY_ENTRY_EXCEPTION` data directory. reserve_pdata_section(&mut self, size: u32) -> SectionRange620 pub fn reserve_pdata_section(&mut self, size: u32) -> SectionRange { 621 let range = self.reserve_section( 622 *b".pdata\0\0", 623 pe::IMAGE_SCN_CNT_INITIALIZED_DATA | pe::IMAGE_SCN_MEM_READ, 624 size, 625 size, 626 ); 627 let dir = &mut self.data_directories[pe::IMAGE_DIRECTORY_ENTRY_EXCEPTION]; 628 debug_assert_eq!(dir.virtual_address, 0); 629 *dir = DataDirectory { 630 virtual_address: range.virtual_address, 631 size, 632 }; 633 range 634 } 635 636 /// Reserve a `.xdata` section. 637 /// 638 /// Contains exception information. reserve_xdata_section(&mut self, size: u32) -> SectionRange639 pub fn reserve_xdata_section(&mut self, size: u32) -> SectionRange { 640 self.reserve_section( 641 *b".xdata\0\0", 642 pe::IMAGE_SCN_CNT_INITIALIZED_DATA | pe::IMAGE_SCN_MEM_READ, 643 size, 644 size, 645 ) 646 } 647 648 /// Reserve a `.rsrc` section. 649 /// 650 /// Contains the resource directory. 651 /// 652 /// This also sets the `pe::IMAGE_DIRECTORY_ENTRY_RESOURCE` data directory. reserve_rsrc_section(&mut self, size: u32) -> SectionRange653 pub fn reserve_rsrc_section(&mut self, size: u32) -> SectionRange { 654 let range = self.reserve_section( 655 *b".rsrc\0\0\0", 656 pe::IMAGE_SCN_CNT_INITIALIZED_DATA | pe::IMAGE_SCN_MEM_READ, 657 size, 658 size, 659 ); 660 let dir = &mut self.data_directories[pe::IMAGE_DIRECTORY_ENTRY_RESOURCE]; 661 debug_assert_eq!(dir.virtual_address, 0); 662 *dir = DataDirectory { 663 virtual_address: range.virtual_address, 664 size, 665 }; 666 range 667 } 668 669 /// Add a base relocation. 670 /// 671 /// `typ` must be one of the `IMAGE_REL_BASED_*` constants. add_reloc(&mut self, mut virtual_address: u32, typ: u16)672 pub fn add_reloc(&mut self, mut virtual_address: u32, typ: u16) { 673 let reloc = U16::new(LE, typ << 12 | (virtual_address & 0xfff) as u16); 674 virtual_address &= !0xfff; 675 if let Some(block) = self.reloc_blocks.last_mut() { 676 if block.virtual_address == virtual_address { 677 self.relocs.push(reloc); 678 block.count += 1; 679 return; 680 } 681 // Blocks must have an even number of relocations. 682 if block.count & 1 != 0 { 683 self.relocs.push(U16::new(LE, 0)); 684 block.count += 1; 685 } 686 debug_assert!(block.virtual_address < virtual_address); 687 } 688 self.relocs.push(reloc); 689 self.reloc_blocks.push(RelocBlock { 690 virtual_address, 691 count: 1, 692 }); 693 } 694 695 /// Return true if a base relocation has been added. has_relocs(&mut self) -> bool696 pub fn has_relocs(&mut self) -> bool { 697 !self.relocs.is_empty() 698 } 699 700 /// Reserve a `.reloc` section. 701 /// 702 /// This contains the base relocations that were added with `add_reloc`. 703 /// 704 /// This also sets the `pe::IMAGE_DIRECTORY_ENTRY_BASERELOC` data directory. reserve_reloc_section(&mut self) -> SectionRange705 pub fn reserve_reloc_section(&mut self) -> SectionRange { 706 if let Some(block) = self.reloc_blocks.last_mut() { 707 // Blocks must have an even number of relocations. 708 if block.count & 1 != 0 { 709 self.relocs.push(U16::new(LE, 0)); 710 block.count += 1; 711 } 712 } 713 let size = self.reloc_blocks.iter().map(RelocBlock::size).sum(); 714 let range = self.reserve_section( 715 *b".reloc\0\0", 716 pe::IMAGE_SCN_CNT_INITIALIZED_DATA 717 | pe::IMAGE_SCN_MEM_READ 718 | pe::IMAGE_SCN_MEM_DISCARDABLE, 719 size, 720 size, 721 ); 722 let dir = &mut self.data_directories[pe::IMAGE_DIRECTORY_ENTRY_BASERELOC]; 723 debug_assert_eq!(dir.virtual_address, 0); 724 *dir = DataDirectory { 725 virtual_address: range.virtual_address, 726 size, 727 }; 728 self.reloc_offset = range.file_offset; 729 range 730 } 731 732 /// Write a `.reloc` section. 733 /// 734 /// This contains the base relocations that were added with `add_reloc`. write_reloc_section(&mut self)735 pub fn write_reloc_section(&mut self) { 736 if self.reloc_offset == 0 { 737 return; 738 } 739 self.pad_until(self.reloc_offset); 740 741 let mut total = 0; 742 for block in &self.reloc_blocks { 743 self.buffer.write(&pe::ImageBaseRelocation { 744 virtual_address: U32::new(LE, block.virtual_address), 745 size_of_block: U32::new(LE, block.size()), 746 }); 747 self.buffer 748 .write_slice(&self.relocs[total..][..block.count as usize]); 749 total += block.count as usize; 750 } 751 debug_assert_eq!(total, self.relocs.len()); 752 753 self.write_align(self.file_alignment); 754 } 755 756 /// Reserve the certificate table. 757 /// 758 /// This also sets the `pe::IMAGE_DIRECTORY_ENTRY_SECURITY` data directory. 759 // TODO: reserve individual certificates reserve_certificate_table(&mut self, size: u32)760 pub fn reserve_certificate_table(&mut self, size: u32) { 761 let size = util::align_u32(size, 8); 762 let offset = self.reserve(size, 8); 763 let dir = &mut self.data_directories[pe::IMAGE_DIRECTORY_ENTRY_SECURITY]; 764 debug_assert_eq!(dir.virtual_address, 0); 765 *dir = DataDirectory { 766 virtual_address: offset, 767 size, 768 }; 769 } 770 771 /// Write the certificate table. 772 // TODO: write individual certificates write_certificate_table(&mut self, data: &[u8])773 pub fn write_certificate_table(&mut self, data: &[u8]) { 774 let dir = self.data_directories[pe::IMAGE_DIRECTORY_ENTRY_SECURITY]; 775 self.pad_until(dir.virtual_address); 776 self.write(data); 777 self.pad_until(dir.virtual_address + dir.size); 778 } 779 } 780 781 /// Information required for writing [`pe::ImageNtHeaders32`] or [`pe::ImageNtHeaders64`]. 782 #[allow(missing_docs)] 783 #[derive(Debug, Clone)] 784 pub struct NtHeaders { 785 // ImageFileHeader 786 pub machine: u16, 787 pub time_date_stamp: u32, 788 pub characteristics: u16, 789 // ImageOptionalHeader 790 pub major_linker_version: u8, 791 pub minor_linker_version: u8, 792 pub address_of_entry_point: u32, 793 pub image_base: u64, 794 pub major_operating_system_version: u16, 795 pub minor_operating_system_version: u16, 796 pub major_image_version: u16, 797 pub minor_image_version: u16, 798 pub major_subsystem_version: u16, 799 pub minor_subsystem_version: u16, 800 pub subsystem: u16, 801 pub dll_characteristics: u16, 802 pub size_of_stack_reserve: u64, 803 pub size_of_stack_commit: u64, 804 pub size_of_heap_reserve: u64, 805 pub size_of_heap_commit: u64, 806 } 807 808 #[derive(Default, Clone, Copy)] 809 struct DataDirectory { 810 virtual_address: u32, 811 size: u32, 812 } 813 814 /// Information required for writing [`pe::ImageSectionHeader`]. 815 #[allow(missing_docs)] 816 #[derive(Debug, Clone)] 817 pub struct Section { 818 pub name: [u8; pe::IMAGE_SIZEOF_SHORT_NAME], 819 pub characteristics: u32, 820 pub range: SectionRange, 821 } 822 823 /// The file range and virtual address range for a section. 824 #[allow(missing_docs)] 825 #[derive(Debug, Default, Clone, Copy)] 826 pub struct SectionRange { 827 pub virtual_address: u32, 828 pub virtual_size: u32, 829 pub file_offset: u32, 830 pub file_size: u32, 831 } 832 833 struct RelocBlock { 834 virtual_address: u32, 835 count: u32, 836 } 837 838 impl RelocBlock { size(&self) -> u32839 fn size(&self) -> u32 { 840 mem::size_of::<pe::ImageBaseRelocation>() as u32 + self.count * mem::size_of::<u16>() as u32 841 } 842 } 843