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