1use 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
15const COFF_MACHINE_X64: u16 = 0x8664;
17
18const COFF_MACHINE_X32: u16 = 0x14c;
20
21const MAX_SECTIONS: u16 = 96;
23
24pub struct Coff<'a> {
26 pub file_header: IMAGE_FILE_HEADER,
28
29 pub symbols: Vec<IMAGE_SYMBOL>,
31
32 pub sections: Vec<IMAGE_SECTION_HEADER>,
34
35 pub buffer: &'a [u8],
37
38 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 pub fn parse(buffer: &'a [u8]) -> Result<Self, CoffError> {
66 debug!("Parsing COFF file header, buffer size: {}", buffer.len());
67
68 if buffer.len() < size_of::<IMAGE_FILE_HEADER>() {
70 return Err(CoffError::InvalidCoffFile);
71 }
72
73 let mut cursor = Cursor::new(buffer);
75
76 let file_header = IMAGE_FILE_HEADER::read(&mut cursor)
78 .map_err(|_| CoffError::InvalidCoffFile)?;
79
80 let arch = Self::validate_architecture(file_header)?;
82
83 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 if num_sections > MAX_SECTIONS {
92 warn!("Exceeded maximum number of sections: {} > {}", num_sections, MAX_SECTIONS);
93 return Err(CoffError::SectionLimitExceeded);
94 }
95
96 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 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 #[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 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 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 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 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 #[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 #[inline]
224 pub fn get_section_name(section: &IMAGE_SECTION_HEADER) -> String {
225 let s = String::from_utf8_lossy(§ion.Name);
226 s.trim_end_matches('\0').to_string()
227 }
228
229 #[inline]
231 pub fn is_fcn(ty: u16) -> bool {
232 (ty & 0x30) == (2 << 4)
233 }
234}
235
236#[derive(Debug, PartialEq, Hash, Clone, Copy, Eq, PartialOrd, Ord)]
238pub enum CoffMachine {
239 X64,
241
242 X32,
244}
245
246impl CoffMachine {
247 #[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
277pub enum CoffSource<'a> {
279 File(&'a str),
281
282 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#[binread]
306#[derive(Default, Debug, Clone, Copy)]
307#[br(little)]
308#[repr(C)]
309pub struct IMAGE_FILE_HEADER {
310 pub Machine: u16,
312
313 pub NumberOfSections: u16,
315
316 pub TimeDateStamp: u32,
318
319 pub PointerToSymbolTable: u32,
321
322 pub NumberOfSymbols: u32,
324
325 pub SizeOfOptionalHeader: u16,
327
328 pub Characteristics: u16,
330}
331
332#[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 pub Value: u32,
343
344 pub SectionNumber: i16,
346
347 pub Type: u16,
349
350 pub StorageClass: u8,
352
353 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#[repr(C, packed(2))]
364#[derive(Clone, Copy)]
365pub union IMAGE_SYMBOL_0 {
366 pub ShortName: [u8; 8],
368
369 pub Name: IMAGE_SYMBOL_0_0,
371
372 pub LongName: [u32; 2],
374}
375
376#[repr(C, packed(2))]
378#[derive(Clone, Copy)]
379pub struct IMAGE_SYMBOL_0_0 {
380 pub Short: u32,
382
383 pub Long: u32,
385}
386
387#[binread]
389#[repr(C)]
390#[br(little)]
391#[derive(Clone, Copy)]
392pub struct IMAGE_SECTION_HEADER {
393 pub Name: [u8; 8],
395
396 #[br(temp)]
397 misc_raw: u32,
398
399 pub VirtualAddress: u32,
401
402 pub SizeOfRawData: u32,
404
405 pub PointerToRawData: u32,
407
408 pub PointerToRelocations: u32,
410
411 pub PointerToLinenumbers: u32,
413
414 pub NumberOfRelocations: u16,
416
417 pub NumberOfLinenumbers: u16,
419
420 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#[repr(C)]
431#[derive(Clone, Copy)]
432pub union IMAGE_SECTION_HEADER_0 {
433 pub PhysicalAddress: u32,
435
436 pub VirtualSize: u32,
438}
439
440#[binread]
442#[br(little)]
443#[repr(C, packed(2))]
444pub struct IMAGE_RELOCATION {
445 #[br(temp)]
446 va_raw: u32,
447
448 pub SymbolTableIndex: u32,
450
451 pub Type: u16,
453
454 #[br(calc = IMAGE_RELOCATION_0 {
455 VirtualAddress: va_raw
456 })]
457 pub Anonymous: IMAGE_RELOCATION_0,
458}
459
460#[repr(C, packed(2))]
462pub union IMAGE_RELOCATION_0 {
463 pub VirtualAddress: u32,
465
466 pub RelocCount: u32,
468}