]>
Commit | Line | Data |
---|---|---|
9c376795 FG |
1 | //! Helper for writing PE files. |
2 | use alloc::string::String; | |
3 | use alloc::vec::Vec; | |
4 | use core::mem; | |
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`. | |
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(§ion_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)] | |
789 | pub 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)] | |
814 | struct 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)] | |
822 | pub 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)] | |
831 | pub struct SectionRange { | |
832 | pub virtual_address: u32, | |
833 | pub virtual_size: u32, | |
834 | pub file_offset: u32, | |
835 | pub file_size: u32, | |
836 | } | |
837 | ||
838 | struct RelocBlock { | |
839 | virtual_address: u32, | |
840 | count: u32, | |
841 | } | |
842 | ||
843 | impl RelocBlock { | |
844 | fn size(&self) -> u32 { | |
845 | mem::size_of::<pe::ImageBaseRelocation>() as u32 + self.count * mem::size_of::<u16>() as u32 | |
846 | } | |
847 | } |