pe_assembler/helpers/builder/
mod.rs

1use crate::{
2    formats::exe::writer::ExeWriter,
3    helpers::pe_writer::PeWriter,
4    types::{
5        tables::{ImportEntry, ImportTable},
6        CoffHeader, DataDirectory, DosHeader, NtHeader, OptionalHeader, PeHeader, PeProgram, PeSection, SubsystemType,
7    },
8};
9use gaia_types::{helpers::Architecture, GaiaError};
10use std::io::Cursor;
11/// PE assembler builder
12#[derive(Debug)]
13pub struct PeBuilder {
14    architecture: Option<Architecture>,
15    subsystem: Option<SubsystemType>,
16    entry_point: Option<u32>,
17    image_base: Option<u64>,
18    imports: Vec<(String, Vec<String>)>, // (dll_name, functions)
19    code: Option<Vec<u8>>,
20    data: Option<Vec<u8>>,
21    sections: Vec<PeSection>, // Add this field
22}
23
24impl PeBuilder {
25    /// Create a new PE assembler builder
26    pub fn new() -> Self {
27        Self {
28            architecture: None,
29            subsystem: None,
30            entry_point: None,
31            image_base: None,
32            imports: Vec::new(),
33            code: None,
34            data: None,
35            sections: Vec::new(), // Initialize the new field
36        }
37    }
38
39    /// Set target architecture
40    pub fn architecture(mut self, arch: Architecture) -> Self {
41        self.architecture = Some(arch);
42        self
43    }
44
45    /// Set subsystem type
46    pub fn subsystem(mut self, subsystem: SubsystemType) -> Self {
47        self.subsystem = Some(subsystem);
48        self
49    }
50
51    /// Set entry point address
52    pub fn entry_point(mut self, entry_point: u32) -> Self {
53        self.entry_point = Some(entry_point);
54        self
55    }
56
57    /// Set image base address
58    pub fn image_base(mut self, image_base: u64) -> Self {
59        self.image_base = Some(image_base);
60        self
61    }
62
63    /// Import a single function
64    pub fn import_function(mut self, dll_name: &str, function_name: &str) -> Self {
65        // Find if the DLL already exists
66        if let Some(entry) = self.imports.iter_mut().find(|(name, _)| name == dll_name) {
67            entry.1.push(function_name.to_string());
68        }
69        else {
70            self.imports.push((dll_name.to_string(), vec![function_name.to_string()]));
71        }
72        self
73    }
74
75    /// Import multiple functions
76    pub fn import_functions(mut self, dll_name: &str, function_names: &[&str]) -> Self {
77        let functions: Vec<String> = function_names.iter().map(|&s| s.to_string()).collect();
78
79        // Find if the DLL already exists
80        if let Some(entry) = self.imports.iter_mut().find(|(name, _)| name == dll_name) {
81            entry.1.extend(functions);
82        }
83        else {
84            self.imports.push((dll_name.to_string(), functions));
85        }
86        self
87    }
88
89    /// Set code data
90    pub fn code(mut self, code: Vec<u8>) -> Self {
91        self.code = Some(code);
92        self
93    }
94
95    /// Set data
96    pub fn data(mut self, data: Vec<u8>) -> Self {
97        self.data = Some(data);
98        self
99    }
100
101    /// Get import information (for creating import table)
102    pub fn get_imports(&self) -> &Vec<(String, Vec<String>)> {
103        &self.imports
104    }
105
106    /// Generate PE header info
107    pub fn build_header(&self) -> Result<PeHeader, GaiaError> {
108        let architecture = self
109            .architecture
110            .as_ref()
111            .ok_or_else(|| GaiaError::syntax_error("Architecture is required", gaia_types::SourceLocation::default()))?;
112        let pointer_size: u32 = if *architecture == Architecture::X86_64 { 8 } else { 4 };
113        let subsystem = self
114            .subsystem
115            .ok_or_else(|| GaiaError::syntax_error("Subsystem is required", gaia_types::SourceLocation::default()))?;
116        let entry_point = self.entry_point.unwrap_or(0x1000);
117        let image_base = self.image_base.unwrap_or(match architecture {
118            Architecture::X86 => 0x400000,
119            Architecture::X86_64 => 0x140000000,
120            _ => 0x400000,
121        });
122
123        // Create DOS header
124        let dos_header = DosHeader::new(0x80);
125
126        // Create NT header
127        let nt_header = NtHeader {
128            signature: 0x00004550, // "PE\0\0"
129        };
130
131        // Create COFF header
132        let machine = match architecture {
133            Architecture::X86 => 0x014C,
134            Architecture::X86_64 => 0x8664,
135            _ => 0x014C,
136        };
137
138        let mut section_count = 0;
139        if self.code.is_some() {
140            section_count += 1;
141        }
142        if self.data.is_some() {
143            section_count += 1;
144        }
145        if !self.imports.is_empty() {
146            section_count += 1;
147        }
148
149        let optional_header_size = match architecture {
150            Architecture::X86_64 => 240,
151            _ => 224,
152        };
153
154        // Set COFF characteristics based on architecture:
155        // - x86: Executable image | 32-bit machine
156        // - x64: Executable image | Large address aware (don't set 32-bit machine bit)
157        let characteristics = match architecture {
158            Architecture::X86 => 0x0102,    // IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_32BIT_MACHINE
159            Architecture::X86_64 => 0x0022, // IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_LARGE_ADDRESS_AWARE
160            _ => 0x0102,
161        };
162
163        let coff_header = CoffHeader::new(machine, section_count)
164            .with_timestamp(0)
165            .with_symbol_table(0, 0)
166            .with_optional_header_size(optional_header_size)
167            .with_characteristics(characteristics);
168
169        // Normalize size_of_code and size_of_initialized_data, using aligned size of sections
170        let size_of_code = if let Some(code) = &self.code { ((code.len() + 0x1FF) / 0x200 * 0x200) as u32 } else { 0 };
171        let mut size_of_initialized_data = 0;
172        if let Some(data) = &self.data {
173            size_of_initialized_data += ((data.len() + 0x1FF) / 0x200 * 0x200) as u32;
174            // .data
175        }
176        if !self.imports.is_empty() {
177            // We can't accurately predict .idata size here, but 0x200 should be enough for most cases
178            // Or we could build sections first to get the exact size
179            size_of_initialized_data += 0x200; // .idata
180        }
181
182        // Calculate size_of_image: starting from 0x1000, each existing section adds 0x1000
183        let mut size_of_image = 0x1000; // DOS/Headers take one alignment page
184        if self.code.is_some() {
185            size_of_image += 0x1000;
186        }
187        if self.data.is_some() {
188            size_of_image += 0x1000;
189        }
190        if !self.imports.is_empty() {
191            size_of_image += 0x1000;
192        }
193
194        let mut optional_header = OptionalHeader::new_for_architecture(
195            architecture,
196            entry_point,
197            image_base,
198            size_of_code,
199            0x200,         // size_of_headers
200            size_of_image, // Dynamically calculated image size
201            subsystem,
202        );
203        optional_header.size_of_initialized_data = size_of_initialized_data;
204
205        // Disable ASLR (DYNAMIC_BASE), otherwise the absolute addresses we patch will be invalid due to base randomization
206        // DYNAMIC_BASE bit value is 0x0040
207        optional_header.dll_characteristics &= !0x0040;
208
209        // Set import table data directory (dynamically calculate RVA)
210        // Compatibility mode (x64): both IAT and INT initially point to the RVA of IMAGE_IMPORT_BY_NAME (Hint+Name).
211        // - x64: INT = array of name pointers, IAT = name RVA (overwritten with real address after loader resolution)
212        // - x86: OFT/INT = 0, IAT = name RVA (common compatible layout)
213        if !self.imports.is_empty() {
214            // Find the actual RVA of the .idata section
215            let idata_section = self
216                .sections
217                .iter()
218                .find(|s| s.name == ".idata")
219                .ok_or_else(|| GaiaError::syntax_error("Missing .idata section", gaia_types::SourceLocation::default()))?;
220            let import_rva_base = idata_section.virtual_address;
221
222            // Subsequent calculations remain unchanged, but use import_rva_base instead of hardcoded values
223            let mut current_rva = import_rva_base + ((self.imports.len() + 1) as u32) * 20;
224            for (dll_name, _) in &self.imports {
225                current_rva += (dll_name.len() as u32) + 1;
226            }
227            if current_rva % 2 != 0 {
228                current_rva += 1;
229            }
230            // Function Hint/Name
231            for (_, functions) in &self.imports {
232                for func in functions {
233                    // Name aligned to 2 bytes
234                    if current_rva % 2 != 0 {
235                        current_rva += 1;
236                    }
237                    current_rva += 2 + (func.len() as u32) + 1;
238                }
239            }
240            // Remove redundant alignment here, keep consistent with write_import_table
241            // INT
242            if current_rva % pointer_size != 0 {
243                current_rva = (current_rva + pointer_size - 1) & !(pointer_size - 1);
244            }
245            for (_, functions) in &self.imports {
246                current_rva += ((functions.len() as u32) + 1) * pointer_size;
247            }
248            // IAT (start of IAT at this point)
249            if current_rva % pointer_size != 0 {
250                current_rva = (current_rva + pointer_size - 1) & !(pointer_size - 1);
251            }
252            let iat_rva_start = current_rva; // Record IAT start
253                                             // IAT length
254            let mut end_rva = current_rva;
255            for (_, functions) in &self.imports {
256                end_rva += ((functions.len() as u32) + 1) * pointer_size;
257            }
258            optional_header.data_directories[1] =
259                DataDirectory { virtual_address: import_rva_base, size: end_rva - import_rva_base };
260            // Also fill IAT Directory (index 12) so loader knows IAT range
261            let mut iat_rva_end = iat_rva_start;
262            for (_, functions) in &self.imports {
263                iat_rva_end += ((functions.len() as u32) + 1) * pointer_size;
264            }
265            optional_header.data_directories[12] =
266                DataDirectory { virtual_address: iat_rva_start, size: iat_rva_end - iat_rva_start };
267        }
268
269        Ok(PeHeader { dos_header, nt_header, coff_header, optional_header })
270    }
271
272    /// Generate section list
273    pub fn build_sections(&mut self) -> Vec<PeSection> {
274        let mut sections = Vec::new();
275        let mut next_virtual_address = 0x1000;
276        let mut next_raw_data_offset = 0x200;
277
278        // Pre-calculate section RVAs for fix_code_relocations
279        // Note: Logic here must be exactly consistent with section addition order below
280        let mut code_rva = None;
281        let mut data_rva = None;
282
283        if self.code.is_some() {
284            code_rva = Some(next_virtual_address);
285            next_virtual_address += 0x1000;
286        }
287        if self.data.is_some() {
288            data_rva = Some(next_virtual_address);
289            next_virtual_address += 0x1000;
290        }
291
292        // Reset for actual building
293        next_virtual_address = 0x1000;
294
295        // Add code section
296        if let Some(code) = &self.code {
297            let mut code_data = code.clone();
298
299            // Fix relocations
300            self.fix_code_relocations_with_rvas(&mut code_data, code_rva.unwrap_or(0x1000), data_rva);
301
302            // Align to 512 bytes
303            let raw_size = ((code_data.len() + 0x1FF) / 0x200 * 0x200) as u32;
304            while code_data.len() < raw_size as usize {
305                code_data.push(0);
306            }
307
308            let text_section = PeSection {
309                name: ".text".to_string(),
310                virtual_size: 0x1000,
311                virtual_address: next_virtual_address,
312                size_of_raw_data: raw_size,
313                pointer_to_raw_data: next_raw_data_offset,
314                pointer_to_relocations: 0,
315                pointer_to_line_numbers: 0,
316                number_of_relocations: 0,
317                number_of_line_numbers: 0,
318                characteristics: 0x60000020,
319                data: code_data,
320            };
321            sections.push(text_section);
322            next_virtual_address += 0x1000;
323            next_raw_data_offset += raw_size;
324        }
325
326        // Add data section
327        if let Some(data) = &self.data {
328            let mut data_bytes = data.clone();
329            // Align to 512 bytes
330            let raw_size = ((data_bytes.len() + 0x1FF) / 0x200 * 0x200) as u32;
331            while data_bytes.len() < raw_size as usize {
332                data_bytes.push(0);
333            }
334
335            let data_section = PeSection {
336                name: ".data".to_string(),
337                virtual_size: 0x1000,
338                virtual_address: next_virtual_address,
339                size_of_raw_data: raw_size,
340                pointer_to_raw_data: next_raw_data_offset,
341                pointer_to_relocations: 0,
342                pointer_to_line_numbers: 0,
343                number_of_relocations: 0,
344                number_of_line_numbers: 0,
345                characteristics: 0xC0000040,
346                data: data_bytes,
347            };
348            sections.push(data_section);
349            next_virtual_address += 0x1000;
350            next_raw_data_offset += raw_size;
351        }
352
353        // Add import table section (if there are imports)
354        if !self.imports.is_empty() {
355            let mut idata_section = self.build_import_section();
356            idata_section.virtual_address = next_virtual_address;
357            idata_section.pointer_to_raw_data = next_raw_data_offset;
358            sections.push(idata_section);
359        }
360
361        sections
362    }
363
364    /// Build import table section
365    fn build_import_section(&self) -> PeSection {
366        // Data is not filled here, let write_import_table handle it.
367        // Note: write_import_table uses "compatible mode" (see above), where IAT initially fills in Hint/Name RVA on x64.
368        PeSection {
369            name: ".idata".to_string(),
370            virtual_size: 0x1000,
371            virtual_address: 0x3000, // This value will be overwritten in build_sections
372            size_of_raw_data: 0x200,
373            pointer_to_raw_data: 0x600, // This value will be overwritten in build_sections
374            pointer_to_relocations: 0,
375            pointer_to_line_numbers: 0,
376            number_of_relocations: 0,
377            number_of_line_numbers: 0,
378            characteristics: 0xC0000040, // IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE
379            data: Vec::new(),            // Empty data, filled by write_import_table method
380        }
381    }
382
383    /// Fix relocations in code
384    fn fix_code_relocations_with_rvas(&self, code: &mut Vec<u8>, code_section_rva: u32, data_section_rva: Option<u32>) {
385        // Find placeholders for CALL instructions and replace them with correct addresses
386        let mut i = 0;
387
388        let arch = self.architecture.as_ref().unwrap_or(&Architecture::X86).clone();
389        let pointer_size: usize = if arch == Architecture::X86_64 { 8 } else { 4 };
390        let image_base: u64 = self.image_base.as_ref().copied().unwrap_or(match arch {
391            Architecture::X86 => 0x400000,
392            Architecture::X86_64 => 0x140000000,
393            _ => 0x400000,
394        });
395
396        // Calculate start RVA of .idata section (consistent with build_sections)
397        let mut import_rva_base: u64 = 0x1000; // Starting RVA
398        if self.code.is_some() {
399            import_rva_base += 0x1000;
400        }
401        if self.data.is_some() {
402            import_rva_base += 0x1000;
403        }
404
405        // Calculate actual position of IAT (must be exactly consistent with calculation in write_import_table)
406        let mut current_rva: u64 = import_rva_base;
407
408        // Calculate IAT offset (keep consistent with write_import_table)
409        if !self.imports.is_empty() {
410            // Import descriptor table size: (number of DLLs + 1) * 20 bytes
411            current_rva += ((self.imports.len() + 1) * 20) as u64;
412
413            // DLL name size
414            for (dll_name, _) in &self.imports {
415                current_rva += (dll_name.len() + 1) as u64; // include null terminator
416            }
417
418            // Align to 2 bytes, consistent with write_import_table
419            if current_rva % 2 != 0 {
420                current_rva += 1;
421            }
422
423            // Function name (Hint(2) + Name + '\0'), accumulate one by one
424            for (_, functions) in &self.imports {
425                for function in functions {
426                    // Name aligned to 2 bytes
427                    if current_rva % 2 != 0 {
428                        current_rva += 1;
429                    }
430                    current_rva += 2 + (function.len() + 1) as u64;
431                }
432            }
433
434            // INT (OriginalFirstThunk) aligned to pointer_size bytes and allocated
435            if current_rva % pointer_size as u64 != 0 {
436                current_rva = (current_rva + pointer_size as u64 - 1) & !(pointer_size as u64 - 1);
437            }
438            for (_, functions) in &self.imports {
439                current_rva += ((functions.len() as u64) + 1) * pointer_size as u64;
440                // include terminator
441            }
442
443            // IAT (FirstThunk) aligned to pointer_size bytes and allocated -> this is the IAT start RVA
444            if current_rva % pointer_size as u64 != 0 {
445                current_rva = (current_rva + pointer_size as u64 - 1) & !(pointer_size as u64 - 1);
446            }
447        }
448
449        let iat_start_rva = current_rva;
450
451        while i < code.len() {
452            // x64: patch multiple instructions using [rip+disp32] addressing
453            if arch == Architecture::X86_64 && i + 2 < code.len() {
454                let mut pos = i;
455                let mut _rex_prefix = None;
456                // Check if it's a REX prefix (0x40 - 0x4F)
457                if code[pos] >= 0x40 && code[pos] <= 0x4F {
458                    _rex_prefix = Some(code[pos]);
459                    pos += 1;
460                }
461
462                if pos + 1 < code.len() {
463                    let mut opcode = code[pos] as u32;
464                    let mut modrm_pos = pos + 1;
465
466                    // Handle 2-byte instructions (starting with 0x0F)
467                    if opcode == 0x0F && pos + 2 < code.len() {
468                        opcode = (opcode << 8) | (code[pos + 1] as u32);
469                        modrm_pos = pos + 2;
470                    }
471
472                    if modrm_pos < code.len() {
473                        let modrm = code[modrm_pos];
474
475                        // Check if ModR/M is [RIP + disp32] (mod=00, rm=101)
476                        if (modrm & 0xC7) == 0x05 && modrm_pos + 4 <= code.len() {
477                            let disp_offset = modrm_pos + 1;
478                            let next_instr_offset = disp_offset + 4;
479
480                            // Branch 1: Indirect call/jmp (CALL/JMP [RIP+disp32]) -> point to Import Address Table (IAT)
481                            // CALL: FF /2, JMP: FF /4
482                            if opcode == 0xFF && ((modrm >> 3) & 7 == 2 || (modrm >> 3) & 7 == 4) {
483                                // Read current displacement as import index
484                                let import_index = u32::from_le_bytes([
485                                    code[disp_offset],
486                                    code[disp_offset + 1],
487                                    code[disp_offset + 2],
488                                    code[disp_offset + 3],
489                                ]) as usize;
490
491                                // Calculate actual IAT RVA
492                                let mut current_iat_offset = 0;
493                                let mut found_rva = None;
494                                let mut flat_idx = 0;
495
496                                'import_search: for (_, functions) in &self.imports {
497                                    for _ in functions {
498                                        if flat_idx == import_index {
499                                            found_rva = Some(iat_start_rva + current_iat_offset);
500                                            break 'import_search;
501                                        }
502                                        flat_idx += 1;
503                                        current_iat_offset += pointer_size as u64;
504                                    }
505                                    current_iat_offset += pointer_size as u64; // skip DLL terminator
506                                }
507
508                                let target_rva =
509                                    found_rva.unwrap_or(iat_start_rva + (import_index as u64 * pointer_size as u64));
510
511                                let rip_rva: u64 = (code_section_rva as u64) + (next_instr_offset as u64);
512                                let disp_i32 = (target_rva as i64 - rip_rva as i64) as i32;
513
514                                code[disp_offset..disp_offset + 4].copy_from_slice(&disp_i32.to_le_bytes());
515
516                                // Debug output
517                                tracing::trace!(
518                                    "IMPORT patch (RIP-relative): i={}, opcode={:04X}, rip_rva={:08X}, target_rva={:08X}, disp={:08X}, idx={}",
519                                    i,
520                                    opcode,
521                                    rip_rva as u32,
522                                    target_rva as u32,
523                                    disp_i32 as u32,
524                                    import_index
525                                );
526
527                                i = next_instr_offset;
528                                continue;
529                            }
530                            // Branch 2: Data access instructions -> point to data section (.data)
531                            // Any [RIP + disp32] that is not an import call is treated as data access
532                            else {
533                                // Special handling for LEA instruction relocation logic
534                                // Check if it's LEA instruction (0x8D)
535                                if opcode == 0x8D {
536                                    // Check if REX prefix contains W bit (0x08)
537                                    // For 64-bit LEA, REX should be 0x48 (W=1) or 0x4C (W=1, R=1) etc.
538                                    let is_64_lea = if let Some(rex) = _rex_prefix { (rex & 0x08) != 0 } else { false };
539
540                                    if is_64_lea {
541                                        // Read current displacement as data index
542                                        let data_index = u32::from_le_bytes([
543                                            code[disp_offset],
544                                            code[disp_offset + 1],
545                                            code[disp_offset + 2],
546                                            code[disp_offset + 3],
547                                        ]) as usize;
548
549                                        // Calculate data RVA
550                                        let data_rva = data_section_rva.unwrap_or(code_section_rva + 0x1000) as u64;
551                                        let target_rva = data_rva + (data_index as u64);
552                                        let rip_rva: u64 = (code_section_rva as u64) + (next_instr_offset as u64);
553                                        let disp_i32 = (target_rva as i64 - rip_rva as i64) as i32;
554
555                                        code[disp_offset..disp_offset + 4].copy_from_slice(&disp_i32.to_le_bytes());
556
557                                        // Debug output
558                                        tracing::trace!(
559                                            "LEA DATA patch (RIP-relative): i={}, opcode={:04X}, rip_rva={:08X}, target_rva={:08X}, disp={:08X}",
560                                            i,
561                                            opcode,
562                                            rip_rva as u32,
563                                            target_rva as u32,
564                                            disp_i32 as u32
565                                        );
566                                    }
567                                    else {
568                                        // Non-64-bit LEA (e.g., using 32-bit register lea eax, [rip+disp32])
569                                        // Read current displacement as offset within data section
570                                        let current_disp = u32::from_le_bytes([
571                                            code[disp_offset],
572                                            code[disp_offset + 1],
573                                            code[disp_offset + 2],
574                                            code[disp_offset + 3],
575                                        ]);
576
577                                        let target_rva: u64 = (data_section_rva.unwrap_or(0) as u64) + (current_disp as u64);
578                                        let rip_rva: u64 = (code_section_rva as u64) + (next_instr_offset as u64);
579                                        let disp_i32 = (target_rva as i64 - rip_rva as i64) as i32;
580
581                                        code[disp_offset..disp_offset + 4].copy_from_slice(&disp_i32.to_le_bytes());
582
583                                        // Debug output
584                                        tracing::trace!(
585                                            "LEA DATA patch (RVA): i={}, opcode={:04X}, rip_rva={:08X}, target_rva={:08X}, disp={:08X}",
586                                            i,
587                                            opcode,
588                                            rip_rva as u32,
589                                            target_rva as u32,
590                                            disp_i32 as u32
591                                        );
592                                    }
593                                }
594                                else {
595                                    // Other [RIP + disp32] instructions (e.g., mov, cmp, etc.)
596                                    // Read current displacement as offset within data section
597                                    let current_disp = u32::from_le_bytes([
598                                        code[disp_offset],
599                                        code[disp_offset + 1],
600                                        code[disp_offset + 2],
601                                        code[disp_offset + 3],
602                                    ]);
603
604                                    let target_rva: u64 = (data_section_rva.unwrap_or(0) as u64) + (current_disp as u64);
605                                    let rip_rva: u64 = (code_section_rva as u64) + (next_instr_offset as u64);
606                                    let disp_i32 = (target_rva as i64 - rip_rva as i64) as i32;
607
608                                    code[disp_offset..disp_offset + 4].copy_from_slice(&disp_i32.to_le_bytes());
609
610                                    // Debug output
611                                    tracing::trace!(
612                                        "DATA patch (RVA): i={}, opcode={:04X}, rip_rva={:08X}, target_rva={:08X}, disp={:08X}",
613                                        i,
614                                        opcode,
615                                        rip_rva as u32,
616                                        target_rva as u32,
617                                        disp_i32 as u32
618                                    );
619                                }
620                            }
621
622                            i = next_instr_offset;
623                            continue;
624                        }
625                    }
626                }
627            }
628
629            // x86 (32-bit) logic remains unchanged
630            if arch == Architecture::X86 && i + 1 < code.len() && code[i] == 0xFF && code[i + 1] == 0x15 && i + 5 < code.len() {
631                // Read current displacement as import index
632                let import_index = u32::from_le_bytes([code[i + 2], code[i + 3], code[i + 4], code[i + 5]]) as usize;
633
634                // Calculate actual IAT RVA
635                let mut current_iat_offset = 0;
636                let mut found_rva = None;
637                let mut flat_idx = 0;
638
639                'import_search_x86: for (_, functions) in &self.imports {
640                    for _ in functions {
641                        if flat_idx == import_index {
642                            found_rva = Some(iat_start_rva + current_iat_offset);
643                            break 'import_search_x86;
644                        }
645                        flat_idx += 1;
646                        current_iat_offset += pointer_size as u64;
647                    }
648                    current_iat_offset += pointer_size as u64;
649                }
650
651                let target_rva = found_rva.unwrap_or(iat_start_rva + (import_index as u64 * pointer_size as u64));
652                let target_va = image_base + target_rva;
653
654                let disp: u32 = target_va as u32;
655                let address_bytes = disp.to_le_bytes();
656                code[i + 2..i + 6].copy_from_slice(&address_bytes);
657                i += 6;
658            }
659            // Find direct CALL instruction (0xE8) - reserved for internal function calls
660            else if code[i] == 0xE8 && i + 4 < code.len() {
661                // Read relative offset
662                let rel_offset = i32::from_le_bytes([code[i + 1], code[i + 2], code[i + 3], code[i + 4]]);
663
664                // Debug output: Since AOT currently doesn't use internal labels, normal direct CALL shouldn't appear here unless it's a reserved placeholder.
665                // If offset is 0, it might be an internal function label placeholder that needs patching.
666                // We keep the rel_offset value for future expansion.
667                tracing::trace!("Direct CALL (E8): i={}, rel_offset={:08X}", i, rel_offset);
668
669                i += 5; // Skip the entire direct CALL instruction
670            }
671            // Find PUSH imm32 instruction (0x68) - used for push_label or data address placeholder
672            else if code[i] == 0x68 && i + 4 < code.len() {
673                // Read immediate value
674                let imm = u32::from_le_bytes([code[i + 1], code[i + 2], code[i + 3], code[i + 4]]);
675
676                if arch == Architecture::X86 {
677                    // If this is a 0 placeholder, check if it should be patched to .data start address
678                    if imm == 0 {
679                        // Check if there's a push_imm8 + push_label pattern before
680                        // push_imm(msg_len) compiles to 6a <len> (2 bytes)
681                        // push_label("msg") compiles to 68 00 00 00 00 (5 bytes)
682                        let mut should_patch = false;
683
684                        // Look back to see if there's a push imm8 <msg_len> immediately followed by current push 0
685                        if i >= 2 {
686                            // Ensure enough space to look back
687                            // Find previous push imm8 instruction (6a)
688                            let prev_push_pos = i - 2;
689                            if prev_push_pos < code.len() && code[prev_push_pos] == 0x6a {
690                                let prev_imm8 = code[prev_push_pos + 1] as u32;
691
692                                // Check if this immediate value equals message length
693                                let msg_len = if let Some(data) = &self.data {
694                                    data.iter().position(|&b| b == 0).map(|p| p as u32).unwrap_or(0)
695                                }
696                                else {
697                                    0
698                                };
699
700                                if prev_imm8 == msg_len {
701                                    should_patch = true;
702                                }
703                            }
704                        }
705
706                        if should_patch {
707                            // Data section start VA
708                            let data_section_va: u64 = image_base + (data_section_rva.unwrap_or(0) as u64);
709                            let addr_u32 = data_section_va as u32;
710                            code[i + 1..i + 5].copy_from_slice(&addr_u32.to_le_bytes());
711                        }
712                    }
713                }
714
715                i += 5; // Skip the entire PUSH instruction
716            }
717            else {
718                i += 1;
719            }
720        }
721    }
722
723    /// Generate PE file byte array
724    pub fn generate(&mut self) -> Result<Vec<u8>, GaiaError> {
725        // Build sections
726        self.sections = self.build_sections(); // Populate sections first
727
728        // Build header
729        let header = self.build_header()?;
730
731        // Build import table
732        let mut import_table = ImportTable::new();
733        for (dll_name, functions) in &self.imports {
734            let entry = ImportEntry { dll_name: dll_name.clone(), functions: functions.clone() };
735            import_table.entries.push(entry);
736        }
737
738        // Create PE program
739        let program = PeProgram {
740            header,
741            sections: self.sections.clone(),
742            imports: import_table,
743            exports: crate::types::tables::ExportTable::new(),
744        };
745
746        // Write to byte array
747        let mut buffer = Vec::new();
748        let cursor = Cursor::new(&mut buffer);
749        let mut writer = ExeWriter::new(cursor);
750        writer.write_program(&program)?;
751
752        Ok(buffer)
753    }
754}
755
756impl Default for PeBuilder {
757    fn default() -> Self {
758        Self::new()
759    }
760}
pe_assembler/helpers/builder/mod.rs

pe_assembler/helpers/builder/
mod.rs
