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