1//! Helper for writing COFF files.
2use alloc::string::String;
3use alloc::vec::Vec;
4use core::mem;
5
6use crate::endian::{LittleEndian as LE, U16Bytes, U32Bytes, U16, U32};
7use crate::pe;
8use crate::write::string::{StringId, StringTable};
9use crate::write::util;
10use crate::write::{Error, Result, WritableBuffer};
11
12/// A helper for writing COFF files.
13///
14/// Writing uses a two phase approach. The first phase builds up all of the information
15/// that may need to be known ahead of time:
16/// - build string table
17/// - reserve section indices
18/// - reserve symbol indices
19/// - reserve file ranges for headers and sections
20///
21/// Some of the information has ordering requirements. For example, strings must be added
22/// to the string table before reserving the file range for the string table. There are debug
23/// asserts to check some of these requirements.
24///
25/// The second phase writes everything out in order. Thus the caller must ensure writing
26/// is in the same order that file ranges were reserved. There are debug asserts to assist
27/// with checking this.
28#[allow(missing_debug_implementations)]
29pub struct Writer<'a> {
30 buffer: &'a mut dyn WritableBuffer,
31 len: usize,
32
33 section_num: u16,
34
35 symtab_offset: u32,
36 symtab_num: u32,
37
38 strtab: StringTable<'a>,
39 strtab_len: usize,
40 strtab_offset: u32,
41 strtab_data: Vec<u8>,
42}
43
44impl<'a> Writer<'a> {
45 /// Create a new `Writer`.
46 pub fn new(buffer: &'a mut dyn WritableBuffer) -> Self {
47 Writer {
48 buffer,
49 len: 0,
50
51 section_num: 0,
52
53 symtab_offset: 0,
54 symtab_num: 0,
55
56 strtab: StringTable::default(),
57 strtab_len: 0,
58 strtab_offset: 0,
59 strtab_data: Vec::new(),
60 }
61 }
62
63 /// Return the current file length that has been reserved.
64 pub fn reserved_len(&self) -> usize {
65 self.len
66 }
67
68 /// Return the current file length that has been written.
69 #[allow(clippy::len_without_is_empty)]
70 pub fn len(&self) -> usize {
71 self.buffer.len()
72 }
73
74 /// Reserve a file range with the given size and starting alignment.
75 ///
76 /// Returns the aligned offset of the start of the range.
77 ///
78 /// `align_start` must be a power of two.
79 pub fn reserve(&mut self, len: usize, align_start: usize) -> u32 {
80 if align_start > 1 {
81 self.len = util::align(self.len, align_start);
82 }
83 let offset = self.len;
84 self.len += len;
85 offset as u32
86 }
87
88 /// Write alignment padding bytes.
89 pub fn write_align(&mut self, align_start: usize) {
90 if align_start > 1 {
91 util::write_align(self.buffer, align_start);
92 }
93 }
94
95 /// Write data.
96 pub fn write(&mut self, data: &[u8]) {
97 self.buffer.write_bytes(data);
98 }
99
100 /// Reserve the file range up to the given file offset.
101 pub fn reserve_until(&mut self, offset: usize) {
102 debug_assert!(self.len <= offset);
103 self.len = offset;
104 }
105
106 /// Write padding up to the given file offset.
107 pub fn pad_until(&mut self, offset: usize) {
108 debug_assert!(self.buffer.len() <= offset);
109 self.buffer.resize(offset);
110 }
111
112 /// Reserve the range for the file header.
113 ///
114 /// This must be at the start of the file.
115 pub fn reserve_file_header(&mut self) {
116 debug_assert_eq!(self.len, 0);
117 self.reserve(mem::size_of::<pe::ImageFileHeader>(), 1);
118 }
119
120 /// Write the file header.
121 ///
122 /// This must be at the start of the file.
123 ///
124 /// Fields that can be derived from known information are automatically set by this function.
125 pub fn write_file_header(&mut self, header: FileHeader) -> Result<()> {
126 debug_assert_eq!(self.buffer.len(), 0);
127
128 // Start writing.
129 self.buffer
130 .reserve(self.len)
131 .map_err(|_| Error(String::from("Cannot allocate buffer")))?;
132
133 // Write file header.
134 let header = pe::ImageFileHeader {
135 machine: U16::new(LE, header.machine),
136 number_of_sections: U16::new(LE, self.section_num),
137 time_date_stamp: U32::new(LE, header.time_date_stamp),
138 pointer_to_symbol_table: U32::new(LE, self.symtab_offset),
139 number_of_symbols: U32::new(LE, self.symtab_num),
140 size_of_optional_header: U16::default(),
141 characteristics: U16::new(LE, header.characteristics),
142 };
143 self.buffer.write(&header);
144
145 Ok(())
146 }
147
148 /// Reserve the range for the section headers.
149 pub fn reserve_section_headers(&mut self, section_num: u16) {
150 debug_assert_eq!(self.section_num, 0);
151 self.section_num = section_num;
152 self.reserve(
153 section_num as usize * mem::size_of::<pe::ImageSectionHeader>(),
154 1,
155 );
156 }
157
158 /// Write a section header.
159 pub fn write_section_header(&mut self, section: SectionHeader) {
160 let mut coff_section = pe::ImageSectionHeader {
161 name: [0; 8],
162 virtual_size: U32::default(),
163 virtual_address: U32::default(),
164 size_of_raw_data: U32::new(LE, section.size_of_raw_data),
165 pointer_to_raw_data: U32::new(LE, section.pointer_to_raw_data),
166 pointer_to_relocations: U32::new(LE, section.pointer_to_relocations),
167 pointer_to_linenumbers: U32::new(LE, section.pointer_to_linenumbers),
168 number_of_relocations: if section.number_of_relocations > 0xffff {
169 U16::new(LE, 0xffff)
170 } else {
171 U16::new(LE, section.number_of_relocations as u16)
172 },
173 number_of_linenumbers: U16::default(),
174 characteristics: U32::new(LE, section.characteristics),
175 };
176 match section.name {
177 Name::Short(name) => coff_section.name = name,
178 Name::Long(str_id) => {
179 let mut str_offset = self.strtab.get_offset(str_id);
180 if str_offset <= 9_999_999 {
181 let mut name = [0; 7];
182 let mut len = 0;
183 if str_offset == 0 {
184 name[6] = b'0';
185 len = 1;
186 } else {
187 while str_offset != 0 {
188 let rem = (str_offset % 10) as u8;
189 str_offset /= 10;
190 name[6 - len] = b'0' + rem;
191 len += 1;
192 }
193 }
194 coff_section.name = [0; 8];
195 coff_section.name[0] = b'/';
196 coff_section.name[1..][..len].copy_from_slice(&name[7 - len..]);
197 } else {
198 debug_assert!(str_offset as u64 <= 0xf_ffff_ffff);
199 coff_section.name[0] = b'/';
200 coff_section.name[1] = b'/';
201 for i in 0..6 {
202 let rem = (str_offset % 64) as u8;
203 str_offset /= 64;
204 let c = match rem {
205 0..=25 => b'A' + rem,
206 26..=51 => b'a' + rem - 26,
207 52..=61 => b'0' + rem - 52,
208 62 => b'+',
209 63 => b'/',
210 _ => unreachable!(),
211 };
212 coff_section.name[7 - i] = c;
213 }
214 }
215 }
216 }
217 self.buffer.write(&coff_section);
218 }
219
220 /// Reserve the range for the section data.
221 ///
222 /// Returns the aligned offset of the start of the range.
223 /// Does nothing and returns 0 if the length is zero.
224 pub fn reserve_section(&mut self, len: usize) -> u32 {
225 if len == 0 {
226 return 0;
227 }
228 // TODO: not sure what alignment is required here, but this seems to match LLVM
229 self.reserve(len, 4)
230 }
231
232 /// Write the alignment bytes prior to section data.
233 ///
234 /// This is unneeded if you are using `write_section` or `write_section_zeroes`
235 /// for the data.
236 pub fn write_section_align(&mut self) {
237 util::write_align(self.buffer, 4);
238 }
239
240 /// Write the section data.
241 ///
242 /// Writes alignment bytes prior to the data.
243 /// Does nothing if the data is empty.
244 pub fn write_section(&mut self, data: &[u8]) {
245 if data.is_empty() {
246 return;
247 }
248 self.write_section_align();
249 self.buffer.write_bytes(data);
250 }
251
252 /// Write the section data using zero bytes.
253 ///
254 /// Writes alignment bytes prior to the data.
255 /// Does nothing if the length is zero.
256 pub fn write_section_zeroes(&mut self, len: usize) {
257 if len == 0 {
258 return;
259 }
260 self.write_section_align();
261 self.buffer.resize(self.buffer.len() + len);
262 }
263
264 /// Reserve a file range for the given number of relocations.
265 ///
266 /// This will automatically reserve an extra relocation if there are more than 0xffff.
267 ///
268 /// Returns the offset of the range.
269 /// Does nothing and returns 0 if the count is zero.
270 pub fn reserve_relocations(&mut self, mut count: usize) -> u32 {
271 if count == 0 {
272 return 0;
273 }
274 if count > 0xffff {
275 count += 1;
276 }
277 self.reserve(count * mem::size_of::<pe::ImageRelocation>(), 1)
278 }
279
280 /// Write a relocation containing the count if required.
281 ///
282 /// This should be called before writing the first relocation for a section.
283 pub fn write_relocations_count(&mut self, count: usize) {
284 if count > 0xffff {
285 let coff_relocation = pe::ImageRelocation {
286 virtual_address: U32Bytes::new(LE, count as u32 + 1),
287 symbol_table_index: U32Bytes::new(LE, 0),
288 typ: U16Bytes::new(LE, 0),
289 };
290 self.buffer.write(&coff_relocation);
291 }
292 }
293
294 /// Write a relocation.
295 pub fn write_relocation(&mut self, reloc: Relocation) {
296 let coff_relocation = pe::ImageRelocation {
297 virtual_address: U32Bytes::new(LE, reloc.virtual_address),
298 symbol_table_index: U32Bytes::new(LE, reloc.symbol),
299 typ: U16Bytes::new(LE, reloc.typ),
300 };
301 self.buffer.write(&coff_relocation);
302 }
303
304 /// Reserve a symbol table entry.
305 ///
306 /// This must be called before [`Self::reserve_symtab_strtab`].
307 pub fn reserve_symbol_index(&mut self) -> u32 {
308 debug_assert_eq!(self.symtab_offset, 0);
309 let index = self.symtab_num;
310 self.symtab_num += 1;
311 index
312 }
313
314 /// Reserve a number of symbol table entries.
315 pub fn reserve_symbol_indices(&mut self, count: u32) {
316 debug_assert_eq!(self.symtab_offset, 0);
317 self.symtab_num += count;
318 }
319
320 /// Write a symbol table entry.
321 pub fn write_symbol(&mut self, symbol: Symbol) {
322 let mut coff_symbol = pe::ImageSymbol {
323 name: [0; 8],
324 value: U32Bytes::new(LE, symbol.value),
325 section_number: U16Bytes::new(LE, symbol.section_number),
326 typ: U16Bytes::new(LE, symbol.typ),
327 storage_class: symbol.storage_class,
328 number_of_aux_symbols: symbol.number_of_aux_symbols,
329 };
330 match symbol.name {
331 Name::Short(name) => coff_symbol.name = name,
332 Name::Long(str_id) => {
333 let str_offset = self.strtab.get_offset(str_id);
334 coff_symbol.name[4..8].copy_from_slice(&u32::to_le_bytes(str_offset as u32));
335 }
336 }
337 self.buffer.write(&coff_symbol);
338 }
339
340 /// Reserve auxiliary symbols for a file name.
341 ///
342 /// Returns the number of auxiliary symbols required.
343 ///
344 /// This must be called before [`Self::reserve_symtab_strtab`].
345 pub fn reserve_aux_file_name(&mut self, name: &[u8]) -> u8 {
346 debug_assert_eq!(self.symtab_offset, 0);
347 let aux_count = (name.len() + pe::IMAGE_SIZEOF_SYMBOL - 1) / pe::IMAGE_SIZEOF_SYMBOL;
348 self.symtab_num += aux_count as u32;
349 aux_count as u8
350 }
351
352 /// Write auxiliary symbols for a file name.
353 pub fn write_aux_file_name(&mut self, name: &[u8], aux_count: u8) {
354 let aux_len = aux_count as usize * pe::IMAGE_SIZEOF_SYMBOL;
355 debug_assert!(aux_len >= name.len());
356 let old_len = self.buffer.len();
357 self.buffer.write_bytes(name);
358 self.buffer.resize(old_len + aux_len);
359 }
360
361 /// Reserve an auxiliary symbol for a section.
362 ///
363 /// Returns the number of auxiliary symbols required.
364 ///
365 /// This must be called before [`Self::reserve_symtab_strtab`].
366 pub fn reserve_aux_section(&mut self) -> u8 {
367 debug_assert_eq!(self.symtab_offset, 0);
368 self.symtab_num += 1;
369 1
370 }
371
372 /// Write an auxiliary symbol for a section.
373 pub fn write_aux_section(&mut self, section: AuxSymbolSection) {
374 let aux = pe::ImageAuxSymbolSection {
375 length: U32Bytes::new(LE, section.length),
376 number_of_relocations: if section.number_of_relocations > 0xffff {
377 U16Bytes::new(LE, 0xffff)
378 } else {
379 U16Bytes::new(LE, section.number_of_relocations as u16)
380 },
381 number_of_linenumbers: U16Bytes::new(LE, section.number_of_linenumbers),
382 check_sum: U32Bytes::new(LE, section.check_sum),
383 number: U16Bytes::new(LE, section.number as u16),
384 selection: section.selection,
385 reserved: 0,
386 high_number: U16Bytes::new(LE, (section.number >> 16) as u16),
387 };
388 self.buffer.write(&aux);
389 }
390
391 /// Return the number of reserved symbol table entries.
392 pub fn symbol_count(&self) -> u32 {
393 self.symtab_num
394 }
395
396 /// Add a string to the string table.
397 ///
398 /// This must be called before [`Self::reserve_symtab_strtab`].
399 pub fn add_string(&mut self, name: &'a [u8]) -> StringId {
400 debug_assert_eq!(self.strtab_offset, 0);
401 self.strtab.add(name)
402 }
403
404 /// Add a section or symbol name to the string table if required.
405 ///
406 /// This must be called before [`Self::reserve_symtab_strtab`].
407 pub fn add_name(&mut self, name: &'a [u8]) -> Name {
408 if name.len() > 8 {
409 Name::Long(self.add_string(name))
410 } else {
411 let mut short_name = [0; 8];
412 short_name[..name.len()].copy_from_slice(name);
413 Name::Short(short_name)
414 }
415 }
416
417 /// Reserve the range for the symbol table and string table.
418 ///
419 /// This must be called after functions that reserve symbol
420 /// indices or add strings.
421 pub fn reserve_symtab_strtab(&mut self) {
422 debug_assert_eq!(self.symtab_offset, 0);
423 self.symtab_offset = self.reserve(self.symtab_num as usize * pe::IMAGE_SIZEOF_SYMBOL, 1);
424
425 debug_assert_eq!(self.strtab_offset, 0);
426 // First 4 bytes of strtab are the length.
427 self.strtab.write(4, &mut self.strtab_data);
428 self.strtab_len = self.strtab_data.len() + 4;
429 self.strtab_offset = self.reserve(self.strtab_len, 1);
430 }
431
432 /// Write the string table.
433 pub fn write_strtab(&mut self) {
434 debug_assert_eq!(self.strtab_offset, self.buffer.len() as u32);
435 self.buffer
436 .write_bytes(&u32::to_le_bytes(self.strtab_len as u32));
437 self.buffer.write_bytes(&self.strtab_data);
438 }
439}
440
441/// Shortened and native endian version of [`pe::ImageFileHeader`].
442#[allow(missing_docs)]
443#[derive(Debug, Default, Clone)]
444pub struct FileHeader {
445 pub machine: u16,
446 pub time_date_stamp: u32,
447 pub characteristics: u16,
448}
449
450/// A section or symbol name.
451#[derive(Debug, Clone, Copy)]
452pub enum Name {
453 /// An inline name.
454 Short([u8; 8]),
455 /// An id of a string table entry.
456 Long(StringId),
457}
458
459impl Default for Name {
460 fn default() -> Name {
461 Name::Short([0; 8])
462 }
463}
464
465// From isn't useful.
466#[allow(clippy::from_over_into)]
467impl<'a> Into<Name> for &'a [u8; 8] {
468 fn into(self) -> Name {
469 Name::Short(*self)
470 }
471}
472
473/// Native endian version of [`pe::ImageSectionHeader`].
474#[allow(missing_docs)]
475#[derive(Debug, Default, Clone)]
476pub struct SectionHeader {
477 pub name: Name,
478 pub size_of_raw_data: u32,
479 pub pointer_to_raw_data: u32,
480 pub pointer_to_relocations: u32,
481 pub pointer_to_linenumbers: u32,
482 /// This will automatically be clamped if there are more than 0xffff.
483 pub number_of_relocations: u32,
484 pub number_of_linenumbers: u16,
485 pub characteristics: u32,
486}
487
488/// Native endian version of [`pe::ImageSymbol`].
489#[allow(missing_docs)]
490#[derive(Debug, Default, Clone)]
491pub struct Symbol {
492 pub name: Name,
493 pub value: u32,
494 pub section_number: u16,
495 pub typ: u16,
496 pub storage_class: u8,
497 pub number_of_aux_symbols: u8,
498}
499
500/// Native endian version of [`pe::ImageAuxSymbolSection`].
501#[allow(missing_docs)]
502#[derive(Debug, Default, Clone)]
503pub struct AuxSymbolSection {
504 pub length: u32,
505 /// This will automatically be clamped if there are more than 0xffff.
506 pub number_of_relocations: u32,
507 pub number_of_linenumbers: u16,
508 pub check_sum: u32,
509 pub number: u32,
510 pub selection: u8,
511}
512
513/// Native endian version of [`pe::ImageRelocation`].
514#[allow(missing_docs)]
515#[derive(Debug, Default, Clone)]
516pub struct Relocation {
517 pub virtual_address: u32,
518 pub symbol: u32,
519 pub typ: u16,
520}
521