object/write/
pe.rs

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}
object/write/pe.rs

object/write/
pe.rs
