coffeeldr/
coff.rs

1//! COFF parsing utilities for the CoffeeLdr loader.
2
3use core::ffi::{CStr, c_void};
4use alloc::{
5    string::{String, ToString},
6    vec::Vec,
7};
8
9use log::{debug, warn};
10use binrw::{BinRead, binread};
11use binrw::io::Cursor;
12
13use crate::error::{CoffError, CoffeeLdrError};
14
15// Architecture definitions for x64
16const COFF_MACHINE_X64: u16 = 0x8664;
17
18// Architecture definitions for x32
19const COFF_MACHINE_X32: u16 = 0x14c;
20
21/// Limit of sections supported by the Windows loader
22const MAX_SECTIONS: u16 = 96;
23
24/// Represents a COFF (Common Object File Format) file.
25pub struct Coff<'a> {
26    /// The COFF file header (`IMAGE_FILE_HEADER`).
27    pub file_header: IMAGE_FILE_HEADER,
28
29    /// A vector of COFF symbols (`IMAGE_SYMBOL`).
30    pub symbols: Vec<IMAGE_SYMBOL>,
31
32    /// A vector of section headers (`IMAGE_SECTION_HEADER`).
33    pub sections: Vec<IMAGE_SECTION_HEADER>,
34
35    /// The raw contents of the file read into memory
36    pub buffer: &'a [u8],
37
38    /// Architecture of the COFF File (x64 or x32)
39    pub arch: CoffMachine,
40}
41
42impl<'a> Default for Coff<'a> {
43    fn default() -> Self {
44        Self {
45            file_header: IMAGE_FILE_HEADER::default(),
46            symbols: Vec::new(),
47            sections: Vec::new(),
48            buffer: &[],
49            arch: CoffMachine::X64,
50        }
51    }
52}
53
54impl<'a> Coff<'a> {
55    /// Parses a COFF object from a byte slice.
56    ///
57    /// The slice must contain a valid COFF file header, section table and
58    /// symbol table. Minimal validation is performed to ensure the layout is
59    /// consistent.
60    ///
61    /// # Errors
62    ///
63    /// Fails when the file is too small, the header is invalid, the
64    /// architecture is unsupported, or any section or symbol fails to decode.
65    pub fn parse(buffer: &'a [u8]) -> Result<Self, CoffError> {
66        debug!("Parsing COFF file header, buffer size: {}", buffer.len());
67
68        // Validates that the file has the minimum size to contain a COFF header
69        if buffer.len() < size_of::<IMAGE_FILE_HEADER>() {
70            return Err(CoffError::InvalidCoffFile);
71        }
72
73        // Creating a cursor
74        let mut cursor = Cursor::new(buffer);
75
76        // The COFF file header
77        let file_header = IMAGE_FILE_HEADER::read(&mut cursor)
78            .map_err(|_| CoffError::InvalidCoffFile)?;
79
80        // Detects the architecture of the COFF file and returns an enum `CoffMachine`
81        let arch = Self::validate_architecture(file_header)?;
82
83        // Checks that the number of sections and symbols is valid
84        let num_sections = file_header.NumberOfSections;
85        let num_symbols = file_header.NumberOfSymbols;
86        if num_sections == 0 || num_symbols == 0 {
87            return Err(CoffError::InvalidSectionsOrSymbols);
88        }
89
90        // Validation of the maximum number of sections (Windows limit)
91        if num_sections > MAX_SECTIONS {
92            warn!("Exceeded maximum number of sections: {} > {}", num_sections, MAX_SECTIONS);
93            return Err(CoffError::SectionLimitExceeded);
94        }
95
96        // A vector of COFF symbols
97        let symbol_offset = file_header.PointerToSymbolTable as usize;
98        let mut cursor = Cursor::new(&buffer[symbol_offset..]);
99        let symbols = (0..num_symbols)
100            .map(|_| {
101                IMAGE_SYMBOL::read(&mut cursor)
102                    .map_err(|_| CoffError::InvalidCoffSymbolsFile)
103            })
104            .collect::<Result<Vec<IMAGE_SYMBOL>, _>>()?;
105
106        // A vector of COFF sections
107        let section_offset = size_of::<IMAGE_FILE_HEADER>() + file_header.SizeOfOptionalHeader as usize;
108        let mut section_cursor = Cursor::new(&buffer[section_offset..]);
109        let sections = (0..num_sections)
110            .map(|_| {
111                IMAGE_SECTION_HEADER::read(&mut section_cursor)
112                    .map_err(|_| CoffError::InvalidCoffSectionFile)
113            })
114            .collect::<Result<Vec<IMAGE_SECTION_HEADER>, _>>()?;
115
116        Ok(Self {
117            file_header,
118            symbols,
119            sections,
120            buffer,
121            arch,
122        })
123    }
124
125    /// Determines the machine architecture from the COFF header.
126    ///
127    /// # Errors
128    ///
129    /// Fails if the machine field does not match any supported architecture.
130    #[inline]
131    fn validate_architecture(file_header: IMAGE_FILE_HEADER) -> Result<CoffMachine, CoffError> {
132        match file_header.Machine {
133            COFF_MACHINE_X64 => Ok(CoffMachine::X64),
134            COFF_MACHINE_X32 => Ok(CoffMachine::X32),
135            _ => {
136                warn!("Unsupported COFF architecture: {:?}", file_header.Machine);
137                Err(CoffError::UnsupportedArchitecture)
138            }
139        }
140    }
141
142    /// Computes the total allocation size needed to load the COFF image.
143    pub fn size(&self) -> usize {
144        let length = self
145            .sections
146            .iter()
147            .filter(|section| section.SizeOfRawData > 0)
148            .map(|section| Self::page_align(section.SizeOfRawData as usize))
149            .sum();
150
151        let total_length = self
152            .sections
153            .iter()
154            .fold(length, |mut total_length, section| {
155                let relocations = self.get_relocations(section);
156                relocations.iter().for_each(|relocation| {
157                    let sym = &self.symbols[relocation.SymbolTableIndex as usize];
158                    let name = self.get_symbol_name(sym);
159                    if name.starts_with("__imp_") {
160                        total_length += size_of::<*const c_void>();
161                    }
162                });
163
164                total_length
165            });
166
167        debug!("Total image size after alignment: {} bytes", total_length);
168        Self::page_align(total_length)
169    }
170
171    /// Returns relocation entries for the specified section.
172    ///
173    /// Invalid relocations are logged and skipped, allowing parsing to proceed.
174    pub fn get_relocations(&self, section: &IMAGE_SECTION_HEADER) -> Vec<IMAGE_RELOCATION> {
175        let reloc_offset = section.PointerToRelocations as usize;
176        let num_relocs = section.NumberOfRelocations as usize;
177        let mut relocations = Vec::with_capacity(num_relocs);
178        let mut cursor = Cursor::new(&self.buffer[reloc_offset..]);
179
180        for _ in 0..num_relocs {
181            match IMAGE_RELOCATION::read(&mut cursor) {
182                Ok(reloc) => relocations.push(reloc),
183                Err(_e) => {
184                    debug!("Failed to read relocation: {_e:?}");
185                    continue;
186                }
187            }
188        }
189
190        relocations
191    }
192
193    /// Reads the symbol name, handling short names and long names from the string table.
194    pub fn get_symbol_name(&self, symtbl: &IMAGE_SYMBOL) -> String {
195        unsafe {
196            let name = if symtbl.N.ShortName[0] != 0 {
197                String::from_utf8_lossy(&symtbl.N.ShortName).into_owned()
198            } else {
199                let long_name_offset = symtbl.N.Name.Long as usize;
200                let string_table_offset = self.file_header.PointerToSymbolTable as usize
201                    + self.file_header.NumberOfSymbols as usize * size_of::<IMAGE_SYMBOL>();
202
203                // Retrieve the name from the string table
204                let offset = string_table_offset + long_name_offset;
205                let name_ptr = &self.buffer[offset] as *const u8;
206                CStr::from_ptr(name_ptr.cast())
207                    .to_string_lossy()
208                    .into_owned()
209            };
210
211            name.trim_end_matches('\0').to_string()
212        }
213    }
214
215    /// Rounds a value up to the next page boundary.
216    #[inline]
217    pub fn page_align(page: usize) -> usize {
218        const SIZE_OF_PAGE: usize = 0x1000;
219        (page + SIZE_OF_PAGE - 1) & !(SIZE_OF_PAGE - 1)
220    }
221
222    /// Extracts the section name, trimming trailing NUL bytes.
223    #[inline]
224    pub fn get_section_name(section: &IMAGE_SECTION_HEADER) -> String {
225        let s = String::from_utf8_lossy(&section.Name);
226        s.trim_end_matches('\0').to_string()
227    }
228
229    /// Determines whether a symbol type describes a function.
230    #[inline]
231    pub fn is_fcn(ty: u16) -> bool {
232        (ty & 0x30) == (2 << 4)
233    }
234}
235
236/// Represents the architecture of the COFF file.
237#[derive(Debug, PartialEq, Hash, Clone, Copy, Eq, PartialOrd, Ord)]
238pub enum CoffMachine {
239    /// 64-bit architecture.
240    X64,
241
242    /// 32-bit architecture.
243    X32,
244}
245
246impl CoffMachine {
247    /// Validates that the COFF architecture matches the host process.
248    ///
249    /// # Errors
250    ///
251    /// Fails if the COFF architecture does not match the current pointer width.
252    #[inline]
253    pub fn check_architecture(&self) -> Result<(), CoffeeLdrError> {
254        match self {
255            CoffMachine::X32 => {
256                if cfg!(target_pointer_width = "64") {
257                    return Err(CoffeeLdrError::ArchitectureMismatch {
258                        expected: "x32",
259                        actual: "x64",
260                    });
261                }
262            }
263            CoffMachine::X64 => {
264                if cfg!(target_pointer_width = "32") {
265                    return Err(CoffeeLdrError::ArchitectureMismatch {
266                        expected: "x64",
267                        actual: "x32",
268                    });
269                }
270            }
271        }
272
273        Ok(())
274    }
275}
276
277/// Represents the COFF data source.
278pub enum CoffSource<'a> {
279    /// COFF file indicated by a string representing the file path.
280    File(&'a str),
281
282    /// Memory buffer containing COFF data.
283    Buffer(&'a [u8]),
284}
285
286impl<'a> From<&'a str> for CoffSource<'a> {
287    fn from(file: &'a str) -> Self {
288        CoffSource::File(file)
289    }
290}
291
292impl<'a, const N: usize> From<&'a [u8; N]> for CoffSource<'a> {
293    fn from(buffer: &'a [u8; N]) -> Self {
294        CoffSource::Buffer(buffer)
295    }
296}
297
298impl<'a> From<&'a [u8]> for CoffSource<'a> {
299    fn from(buffer: &'a [u8]) -> Self {
300        CoffSource::Buffer(buffer)
301    }
302}
303
304/// Represents the file header of a COFF (Common Object File Format) file.
305#[binread]
306#[derive(Default, Debug, Clone, Copy)]
307#[br(little)]
308#[repr(C)]
309pub struct IMAGE_FILE_HEADER {
310    /// The target machine architecture (e.g., x64, x32).
311    pub Machine: u16,
312
313    /// The number of sections in the COFF file.
314    pub NumberOfSections: u16,
315
316    /// The timestamp when the file was created.
317    pub TimeDateStamp: u32,
318
319    /// The pointer to the symbol table.
320    pub PointerToSymbolTable: u32,
321
322    /// The number of symbols in the COFF file.
323    pub NumberOfSymbols: u32,
324
325    /// The size of the optional header.
326    pub SizeOfOptionalHeader: u16,
327
328    /// The characteristics of the file.
329    pub Characteristics: u16,
330}
331
332/// Represents a symbol in the COFF symbol table.
333#[binread]
334#[derive(Clone, Copy)]
335#[br(little)]
336#[repr(C, packed(2))]
337pub struct IMAGE_SYMBOL {
338    #[br(temp)]
339    name_raw: [u8; 8],
340
341    /// The value associated with the symbol.
342    pub Value: u32,
343
344    /// The section number that contains the symbol.
345    pub SectionNumber: i16,
346
347    /// The type of the symbol.
348    pub Type: u16,
349
350    /// The storage class of the symbol (e.g., external, static).
351    pub StorageClass: u8,
352
353    /// The number of auxiliary symbol records.
354    pub NumberOfAuxSymbols: u8,
355
356    #[br(calc = unsafe {
357        core::ptr::read_unaligned(name_raw.as_ptr() as *const IMAGE_SYMBOL_0)
358    })]
359    pub N: IMAGE_SYMBOL_0,
360}
361
362/// A union representing different ways a symbol name can be stored.
363#[repr(C, packed(2))]
364#[derive(Clone, Copy)]
365pub union IMAGE_SYMBOL_0 {
366    /// A short symbol name (8 bytes).
367    pub ShortName: [u8; 8],
368
369    /// A long symbol name stored in a different structure.
370    pub Name: IMAGE_SYMBOL_0_0,
371
372    /// Long symbol name stored as a pair of u32 values.
373    pub LongName: [u32; 2],
374}
375
376/// Represents the long name of a symbol as a pair of values.
377#[repr(C, packed(2))]
378#[derive(Clone, Copy)]
379pub struct IMAGE_SYMBOL_0_0 {
380    /// The offset to the symbol name.
381    pub Short: u32,
382
383    /// The length of the symbol name.
384    pub Long: u32,
385}
386
387/// Represents a section header in a COFF file.
388#[binread]
389#[repr(C)]
390#[br(little)]
391#[derive(Clone, Copy)]
392pub struct IMAGE_SECTION_HEADER {
393    /// The name of the section (8 bytes).
394    pub Name: [u8; 8],
395
396    #[br(temp)]
397    misc_raw: u32,
398
399    /// The virtual address of the section in memory.
400    pub VirtualAddress: u32,
401
402    /// The size of the section's raw data.
403    pub SizeOfRawData: u32,
404
405    /// The pointer to the raw data in the file.
406    pub PointerToRawData: u32,
407
408    /// The pointer to relocation entries.
409    pub PointerToRelocations: u32,
410
411    /// The pointer to line numbers (if any).
412    pub PointerToLinenumbers: u32,
413
414    /// The number of relocations in the section.
415    pub NumberOfRelocations: u16,
416
417    /// The number of line numbers in the section.
418    pub NumberOfLinenumbers: u16,
419
420    /// Characteristics that describe the section (e.g., executable, writable).
421    pub Characteristics: u32,
422
423    #[br(calc = IMAGE_SECTION_HEADER_0 {
424        PhysicalAddress: misc_raw
425    })]
426    pub Misc: IMAGE_SECTION_HEADER_0,
427}
428
429/// A union representing either the physical or virtual size of the section.
430#[repr(C)]
431#[derive(Clone, Copy)]
432pub union IMAGE_SECTION_HEADER_0 {
433    /// The physical address of the section.
434    pub PhysicalAddress: u32,
435
436    /// The virtual size of the section.
437    pub VirtualSize: u32,
438}
439
440/// Represents a relocation entry in a COFF file.
441#[binread]
442#[br(little)]
443#[repr(C, packed(2))]
444pub struct IMAGE_RELOCATION {
445    #[br(temp)]
446    va_raw: u32,
447
448    /// The index of the symbol in the symbol table.
449    pub SymbolTableIndex: u32,
450
451    /// The type of relocation.
452    pub Type: u16,
453
454    #[br(calc = IMAGE_RELOCATION_0 {
455        VirtualAddress: va_raw
456    })]
457    pub Anonymous: IMAGE_RELOCATION_0,
458}
459
460/// A union representing either the virtual address or relocation count.
461#[repr(C, packed(2))]
462pub union IMAGE_RELOCATION_0 {
463    /// The virtual address of the relocation.
464    pub VirtualAddress: u32,
465
466    /// The relocation count.
467    pub RelocCount: u32,
468}