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