pe_assembler/helpers/pe_reader/
mod.rs

1use crate::types::{
2    CoffHeader, DosHeader, ExportTable, ImportTable, NtHeader, OptionalHeader, PeHeader, PeInfo, PeProgram, PeSection,
3    SectionHeader,
4};
5use byteorder::{LittleEndian, ReadBytesExt, LE};
6use gaia_types::{helpers::Architecture, BinaryReader, GaiaError};
7use std::io::{Read, Seek, SeekFrom};
8
9/// PE 文件读取器的通用 trait
10pub trait PeReader<R: Read + Seek> {
11    /// 获取二进制读取器的可变引用
12    fn get_viewer(&mut self) -> &mut R;
13
14    /// 获取诊断信息的可变引用
15    fn add_diagnostics(&mut self, error: impl Into<GaiaError>);
16
17    fn get_position(&mut self) -> Result<u64, GaiaError>
18    where
19        R: Seek,
20    {
21        let position = self.get_viewer().stream_position()?;
22        Ok(position)
23    }
24
25    fn set_position(&mut self, offset: u64) -> Result<u64, GaiaError>
26    where
27        R: Seek,
28    {
29        let position = self.get_viewer().seek(SeekFrom::Start(offset))?;
30        Ok(position)
31    }
32
33    /// 获取缓存的节头信息
34    fn get_section_headers(&mut self) -> Result<&[SectionHeader], GaiaError>;
35
36    /// 设置缓存的节头信息
37    fn set_section_headers(&mut self, headers: Vec<SectionHeader>) -> Vec<SectionHeader>;
38
39    /// 读取 PE 头部信息（通用实现）
40    fn get_pe_header(&mut self) -> Result<&PeHeader, GaiaError>;
41
42    fn set_pe_header(&mut self, head: PeHeader) -> Option<PeHeader>;
43
44    /// 将 RVA 转换为文件偏移（通用实现）
45    fn rva_to_file_offset(&self, rva: u32, sections: &[PeSection]) -> Result<u32, GaiaError> {
46        for section in sections {
47            if rva >= section.virtual_address && rva < section.virtual_address + section.virtual_size {
48                let offset_in_section = rva - section.virtual_address;
49                return Ok(section.pointer_to_raw_data + offset_in_section);
50            }
51        }
52        Err(GaiaError::invalid_data(&format!("无法将 RVA 0x{:08X} 转换为文件偏移", rva)))
53    }
54
55    /// 解析导入表（通用实现）
56    fn parse_import_table(&mut self, header: &PeHeader, sections: &[PeSection]) -> Result<ImportTable, GaiaError> {
57        // 检查数据目录表是否包含导入表信息
58        if header.optional_header.data_directories.len() < 2 {
59            return Ok(ImportTable::new());
60        }
61
62        let import_dir = &header.optional_header.data_directories[1]; // 导入表是第2个数据目录
63        if import_dir.virtual_address == 0 || import_dir.size == 0 {
64            return Ok(ImportTable::new());
65        }
66
67        // 将 RVA 转换为文件偏移
68        let file_offset = self.rva_to_file_offset(import_dir.virtual_address, sections)?;
69
70        // 保存当前位置
71        let current_pos = self.get_position()?;
72
73        // 定位到导入表
74        self.set_position(file_offset as u64)?;
75
76        let mut import_table = ImportTable::new();
77
78        // 读取导入描述符
79        loop {
80            let import_lookup_table = self.get_viewer().read_u32::<LittleEndian>()?;
81            let time_date_stamp = self.get_viewer().read_u32::<LittleEndian>()?;
82            let forwarder_chain = self.get_viewer().read_u32::<LittleEndian>()?;
83            let name_rva = self.get_viewer().read_u32::<LittleEndian>()?;
84            let import_address_table = self.get_viewer().read_u32::<LittleEndian>()?;
85
86            // 如果所有字段都为0，表示导入表结束
87            if import_lookup_table == 0
88                && time_date_stamp == 0
89                && forwarder_chain == 0
90                && name_rva == 0
91                && import_address_table == 0
92            {
93                break;
94            }
95
96            let mut dll_name = String::new();
97            let mut functions = Vec::new();
98
99            // 读取 DLL 名称
100            if name_rva != 0 {
101                let name_offset = self.rva_to_file_offset(name_rva, sections)?;
102                let saved_pos = self.get_position()?;
103                self.set_position(name_offset as u64)?;
104
105                let mut name_bytes = Vec::new();
106                loop {
107                    let byte = self.get_viewer().read_u8()?;
108                    if byte == 0 {
109                        break;
110                    }
111                    name_bytes.push(byte);
112                }
113                dll_name = String::from_utf8_lossy(&name_bytes).to_string();
114                self.set_position(saved_pos)?;
115            }
116
117            // 读取函数名称（从导入查找表）
118            if import_lookup_table != 0 {
119                let lookup_offset = self.rva_to_file_offset(import_lookup_table, sections)?;
120                let saved_pos = self.get_position()?;
121                self.set_position(lookup_offset as u64)?;
122
123                loop {
124                    let entry = if header.optional_header.magic == 0x20b {
125                        // PE32+
126                        self.get_viewer().read_u64::<LittleEndian>()?
127                    }
128                    else {
129                        // PE32
130                        self.get_viewer().read_u32::<LittleEndian>()? as u64
131                    };
132
133                    if entry == 0 {
134                        break;
135                    }
136
137                    // 检查是否是按名称导入（最高位为0）
138                    let is_ordinal = if header.optional_header.magic == 0x20b {
139                        (entry & 0x8000000000000000) != 0
140                    }
141                    else {
142                        (entry & 0x80000000) != 0
143                    };
144
145                    if !is_ordinal {
146                        let hint_name_rva =
147                            entry & if header.optional_header.magic == 0x20b { 0x7FFFFFFFFFFFFFFF } else { 0x7FFFFFFF };
148                        let hint_name_offset = self.rva_to_file_offset(hint_name_rva as u32, sections)?;
149                        let func_pos = self.get_position()?;
150                        self.set_position(hint_name_offset as u64)?;
151
152                        // 跳过 hint（2字节）
153                        self.get_viewer().read_u16::<LittleEndian>()?;
154
155                        // 读取函数名
156                        let mut func_name_bytes = Vec::new();
157                        loop {
158                            let byte = self.get_viewer().read_u8()?;
159                            if byte == 0 {
160                                break;
161                            }
162                            func_name_bytes.push(byte);
163                        }
164                        let func_name = String::from_utf8_lossy(&func_name_bytes).to_string();
165                        functions.push(func_name);
166
167                        self.set_position(func_pos)?;
168                    }
169                    else {
170                        // 按序号导入
171                        let ordinal = entry & 0xFFFF;
172                        functions.push(format!("Ordinal_{}", ordinal));
173                    }
174                }
175
176                self.set_position(saved_pos)?;
177            }
178
179            // 添加导入条目
180            if !dll_name.is_empty() {
181                use crate::types::tables::ImportEntry;
182                let entry = ImportEntry { dll_name, functions };
183                import_table.entries.push(entry);
184            }
185        }
186
187        // 恢复位置
188        self.set_position(current_pos)?;
189
190        Ok(import_table)
191    }
192
193    /// 解析导出表（通用实现）
194    fn parse_export_table(&mut self, header: &PeHeader, sections: &[PeSection]) -> Result<ExportTable, GaiaError> {
195        // 检查数据目录表是否包含导出表信息
196        if header.optional_header.data_directories.is_empty() {
197            return Ok(ExportTable { name: String::new(), functions: Vec::new() });
198        }
199
200        let export_dir = &header.optional_header.data_directories[0]; // 导出表是第1个数据目录
201        if export_dir.virtual_address == 0 || export_dir.size == 0 {
202            return Ok(ExportTable { name: String::new(), functions: Vec::new() });
203        }
204
205        // 将 RVA 转换为文件偏移
206        let file_offset = self.rva_to_file_offset(export_dir.virtual_address, sections)?;
207
208        // 保存当前位置
209        let current_pos = self.get_position()?;
210
211        // 定位到导出表
212        self.set_position(file_offset as u64)?;
213
214        // 读取导出目录表
215        let _export_flags = self.get_viewer().read_u32::<LittleEndian>()?;
216        let _time_date_stamp = self.get_viewer().read_u32::<LittleEndian>()?;
217        let _major_version = self.get_viewer().read_u16::<LittleEndian>()?;
218        let _minor_version = self.get_viewer().read_u16::<LittleEndian>()?;
219        let name_rva = self.get_viewer().read_u32::<LittleEndian>()?;
220        let _ordinal_base = self.get_viewer().read_u32::<LittleEndian>()?;
221        let _number_of_functions = self.get_viewer().read_u32::<LittleEndian>()?;
222        let number_of_names = self.get_viewer().read_u32::<LittleEndian>()?;
223        let _address_of_functions = self.get_viewer().read_u32::<LittleEndian>()?;
224        let address_of_names = self.get_viewer().read_u32::<LittleEndian>()?;
225        let _address_of_name_ordinals = self.get_viewer().read_u32::<LittleEndian>()?;
226
227        // 读取模块名称
228        let mut name = String::new();
229        if name_rva != 0 {
230            let name_offset = self.rva_to_file_offset(name_rva, sections)?;
231            let saved_pos = self.get_position()?;
232            self.set_position(name_offset as u64)?;
233
234            let mut name_bytes = Vec::new();
235            loop {
236                let byte = self.get_viewer().read_u8()?;
237                if byte == 0 {
238                    break;
239                }
240                name_bytes.push(byte);
241            }
242            name = String::from_utf8_lossy(&name_bytes).to_string();
243            self.set_position(saved_pos)?;
244        }
245
246        // 读取函数名称
247        let mut functions = Vec::new();
248        if address_of_names != 0 && number_of_names > 0 {
249            let names_offset = self.rva_to_file_offset(address_of_names, sections)?;
250            let saved_pos = self.get_position()?;
251            self.set_position(names_offset as u64)?;
252
253            for _ in 0..number_of_names {
254                let name_rva = self.get_viewer().read_u32::<LittleEndian>()?;
255                if name_rva != 0 {
256                    let func_name_offset = self.rva_to_file_offset(name_rva, sections)?;
257                    let func_pos = self.get_position()?;
258                    self.set_position(func_name_offset as u64)?;
259
260                    let mut func_name_bytes = Vec::new();
261                    loop {
262                        let byte = self.get_viewer().read_u8()?;
263                        if byte == 0 {
264                            break;
265                        }
266                        func_name_bytes.push(byte);
267                    }
268                    let func_name = String::from_utf8_lossy(&func_name_bytes).to_string();
269                    functions.push(func_name);
270
271                    self.set_position(func_pos)?;
272                }
273            }
274
275            self.set_position(saved_pos)?;
276        }
277
278        // 恢复位置
279        self.set_position(current_pos)?;
280
281        Ok(ExportTable { name, functions })
282    }
283
284    /// 创建 PE 信息视图（通用实现）
285    fn create_pe_info(&mut self) -> Result<PeInfo, GaiaError> {
286        let header = self.get_pe_header()?.clone();
287
288        // 根据机器类型确定架构
289        let target_arch = match header.coff_header.machine {
290            0x014c => Architecture::X86,
291            0x8664 => Architecture::X86_64,
292            0x01c0 => Architecture::ARM32,
293            0xaa64 => Architecture::ARM64,
294            unknown => {
295                tracing::warn!("未知的机器类型: {:04x}", unknown);
296                Architecture::Unknown
297            }
298        };
299
300        // 获取当前文件大小
301        let current_pos = self.get_position()?;
302        self.get_viewer().seek(std::io::SeekFrom::End(0))?;
303        let file_size = self.get_position()?;
304        self.set_position(current_pos)?;
305
306        Ok(PeInfo {
307            target_arch,
308            subsystem: header.optional_header.subsystem,
309            entry_point: header.optional_header.address_of_entry_point,
310            image_base: header.optional_header.image_base,
311            section_count: header.coff_header.number_of_sections,
312            file_size,
313        })
314    }
315    /// 强制读取完整的 [PeProgram]，并缓存结果
316    fn get_program(&mut self) -> Result<&PeProgram, GaiaError>;
317
318    fn set_program(&mut self, program: PeProgram) -> Option<PeProgram>;
319}
320
321/// 解析 PE 头部（通用实现）
322pub(crate) fn read_pe_head<R: Read + Seek>(reader: &mut impl PeReader<R>) -> Result<PeHeader, GaiaError> {
323    // 保存当前位置
324    let original_pos = reader.get_position()?;
325
326    // 重置到文件开头
327    reader.set_position(0)?;
328
329    // 读取 DOS 头
330    let dos_header = DosHeader::read(reader.get_viewer())?;
331
332    // 验证 DOS 签名 (MZ)
333    if dos_header.e_magic != 0x5A4D {
334        let error = GaiaError::invalid_data("无效的 DOS 签名 (MZ)");
335        reader.add_diagnostics(error);
336    }
337
338    // 跳转到 NT 头位置
339    reader.set_position(dos_header.e_lfanew as u64)?;
340
341    // 读取 NT 头
342    let nt_header = NtHeader::read(reader.get_viewer())?;
343
344    // 验证 PE 签名 (PE\0\0)
345    if nt_header.signature != 0x00004550 {
346        let error = GaiaError::invalid_data("无效的 PE 签名 (PE)");
347        reader.add_diagnostics(error);
348    }
349
350    // 读取 COFF 头
351    let coff_header = CoffHeader::read(reader.get_viewer())?;
352
353    // 验证 COFF 头中的节数量
354    if coff_header.number_of_sections == 0 {
355        let error = GaiaError::invalid_data("PE 文件必须至少有一个节");
356        reader.add_diagnostics(error);
357    }
358
359    // 读取可选头
360    let optional_header = OptionalHeader::read(reader.get_viewer())?;
361
362    // 验证可选头的魔数
363    match optional_header.magic {
364        0x10b => {} // PE32
365        0x20b => {} // PE32+
366        _ => {
367            let error = GaiaError::invalid_data("无效的可选头魔数");
368            reader.add_diagnostics(error);
369            return Err(GaiaError::invalid_data("无效的可选头魔数"));
370        }
371    }
372
373    // 恢复原始位置
374    reader.set_position(original_pos)?;
375
376    Ok(PeHeader { dos_header, nt_header, coff_header, optional_header })
377}
378
379/// 读取节头信息（通用实现）
380pub(crate) fn read_pe_section_headers<R: Read + Seek>(reader: &mut impl PeReader<R>) -> Result<(), GaiaError> {
381    // 先读取主头部
382    let header = reader.get_pe_header()?.clone();
383    let original_pos = reader.get_position()?;
384
385    // 读取节头部
386    let mut section_headers = Vec::with_capacity(header.coff_header.number_of_sections as usize);
387
388    // 定位到节头部位置
389    let section_header_offset = header.dos_header.e_lfanew as u64
390        + 4 // PE signature
391        + std::mem::size_of::<CoffHeader>() as u64
392        + header.coff_header.size_of_optional_header as u64;
393
394    reader.set_position(section_header_offset)?;
395
396    for _ in 0..header.coff_header.number_of_sections {
397        let section_header = SectionHeader::read(reader.get_viewer())?;
398        section_headers.push(section_header);
399    }
400
401    // 恢复原始位置
402    reader.set_position(original_pos)?;
403
404    reader.set_section_headers(section_headers);
405    Ok(())
406}
407
408/// 从节头读取节数据（通用实现）
409pub(crate) fn read_section_from_header<R: Read + Seek>(
410    reader: &mut impl PeReader<R>,
411    header: &SectionHeader,
412) -> Result<PeSection, GaiaError> {
413    let name = String::from_utf8_lossy(&header.name).trim_end_matches('\0').to_string();
414
415    let mut data = Vec::new();
416    if header.size_of_raw_data > 0 && header.pointer_to_raw_data > 0 {
417        let original_pos = reader.get_position()?;
418        reader.set_position(header.pointer_to_raw_data as u64)?;
419        data.resize(header.size_of_raw_data as usize, 0);
420        reader.get_viewer().read_exact(&mut data)?;
421        reader.set_position(original_pos)?;
422    }
423
424    Ok(PeSection {
425        name,
426        virtual_size: header.virtual_size,
427        virtual_address: header.virtual_address,
428        size_of_raw_data: header.size_of_raw_data,
429        pointer_to_raw_data: header.pointer_to_raw_data,
430        pointer_to_relocations: header.pointer_to_relocations,
431        pointer_to_line_numbers: header.pointer_to_line_numbers,
432        number_of_relocations: header.number_of_relocations,
433        number_of_line_numbers: header.number_of_line_numbers,
434        characteristics: header.characteristics,
435        data,
436    })
437}
438
439pub(crate) fn read_pe_program<R: Read + Seek>(reader: &mut impl PeReader<R>) -> Result<(), GaiaError> {
440    let header = reader.get_pe_header()?.clone();
441    let section_headers = reader.get_section_headers()?.to_vec();
442
443    // 读取节数据
444    let mut sections = Vec::new();
445    for section_header in section_headers {
446        let section = read_section_from_header(reader, &section_header)?;
447        sections.push(section);
448    }
449
450    // 解析导入表
451    let imports = reader.parse_import_table(&header, &sections)?;
452
453    // 解析导出表（EXE 文件通常没有导出表）
454    let exports = reader.parse_export_table(&header, &sections)?;
455
456    reader.set_program(PeProgram { header, sections, imports, exports });
457    Ok(())
458}
pe_assembler/helpers/pe_reader/mod.rs

pe_assembler/helpers/pe_reader/
mod.rs
