pe_assembler/helpers/coff_reader/
mod.rs

1use crate::types::{
2    coff::{CoffFileType, CoffHeader, CoffInfo, CoffObject, CoffRelocation, CoffSymbol, SectionHeader},
3    CoffSection,
4};
5use byteorder::{LittleEndian, ReadBytesExt};
6use gaia_types::{helpers::Architecture, GaiaError};
7use std::io::{Read, Seek, SeekFrom};
8
9/// COFF 文件读取器的通用 trait
10pub trait CoffReader<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        Ok(self.get_viewer().stream_position()?)
22    }
23
24    fn set_position(&mut self, offset: u64) -> Result<u64, GaiaError>
25    where
26        R: Seek,
27    {
28        Ok(self.get_viewer().seek(SeekFrom::Start(offset))?)
29    }
30
31    /// 读取 COFF 头部信息（通用实现）
32    fn get_coff_header(&mut self) -> Result<&CoffHeader, GaiaError>;
33
34    /// 读取 COFF 头部信息（通用实现）
35    fn set_coff_header(&mut self, head: CoffHeader) -> Option<CoffHeader>;
36
37    /// 获取缓存的节头信息
38    fn get_section_headers(&mut self) -> Result<&[SectionHeader], GaiaError>;
39
40    /// 设置缓存的节头信息
41    fn set_section_headers(&mut self, headers: Vec<SectionHeader>) -> Vec<SectionHeader>;
42
43    /// 强制读取完整的 COFF 对象，并缓存结果
44    fn get_coff_object(&mut self) -> Result<&CoffObject, GaiaError>;
45
46    /// 强制读取完整的 COFF 对象，并缓存结果
47    fn set_coff_object(&mut self, object: CoffObject) -> Option<CoffObject>;
48
49    /// 获取缓存的 COFF 信息
50    fn get_coff_info(&mut self) -> Result<&CoffInfo, GaiaError>;
51
52    /// 设置缓存的 COFF 信息
53    fn set_coff_info(&mut self, info: CoffInfo) -> Option<CoffInfo>;
54
55    /// 创建 COFF 信息视图（通用实现）
56    fn create_coff_info(&mut self) -> Result<CoffInfo, GaiaError> {
57        let header = self.get_coff_header()?.clone();
58
59        // 根据机器类型确定架构
60        let target_arch = match header.machine {
61            0x014c => Architecture::X86,
62            0x8664 => Architecture::X86_64,
63            0x01c0 => Architecture::ARM32,
64            0xaa64 => Architecture::ARM64,
65            _ => Architecture::Unknown,
66        };
67
68        // 获取当前文件大小
69        let current_pos = self.get_position()?;
70        self.get_viewer().seek(SeekFrom::End(0))?;
71        let file_size = self.get_position()?;
72        self.set_position(current_pos)?;
73
74        Ok(CoffInfo {
75            file_type: CoffFileType::Object,
76            target_arch,
77            section_count: header.number_of_sections,
78            symbol_count: header.number_of_symbols,
79            file_size,
80            timestamp: header.time_date_stamp,
81        })
82    }
83}
84
85impl CoffRelocation {
86    /// 从 CoffViewer 读取重定位项
87    pub fn read<R: Read>(mut reader: R) -> Result<Self, GaiaError> {
88        Ok(CoffRelocation {
89            virtual_address: reader.read_u32::<LittleEndian>()?,
90            symbol_table_index: reader.read_u32::<LittleEndian>()?,
91            relocation_type: reader.read_u16::<LittleEndian>()?,
92        })
93    }
94}
95
96impl CoffSymbol {
97    /// 从 CoffViewer 读取符号
98    pub fn read<R: Read>(mut reader: R) -> Result<Self, GaiaError> {
99        let mut name_bytes = [0u8; 8];
100        reader.read_exact(&mut name_bytes)?;
101
102        let name = if name_bytes[0..4] == [0, 0, 0, 0] {
103            // 长名称，存储在字符串表中
104            format!("symbol_{}", u32::from_le_bytes([name_bytes[4], name_bytes[5], name_bytes[6], name_bytes[7]]))
105        }
106        else {
107            String::from_utf8_lossy(&name_bytes).trim_end_matches('\0').to_string()
108        };
109
110        Ok(CoffSymbol {
111            name,
112            value: reader.read_u32::<LittleEndian>()?,
113            section_number: reader.read_i16::<LittleEndian>()?,
114            symbol_type: reader.read_u16::<LittleEndian>()?,
115            storage_class: reader.read_u8()?,
116            number_of_aux_symbols: reader.read_u8()?,
117        })
118    }
119}
120
121impl SectionHeader {
122    /// 从 ExeReader 读取节头
123    pub fn read<R: Read>(mut reader: R) -> Result<Self, GaiaError> {
124        let mut name = [0u8; 8];
125        reader.read_exact(&mut name)?;
126
127        Ok(SectionHeader {
128            name,
129            virtual_size: reader.read_u32::<LittleEndian>()?,
130            virtual_address: reader.read_u32::<LittleEndian>()?,
131            size_of_raw_data: reader.read_u32::<LittleEndian>()?,
132            pointer_to_raw_data: reader.read_u32::<LittleEndian>()?,
133            pointer_to_relocations: reader.read_u32::<LittleEndian>()?,
134            pointer_to_line_numbers: reader.read_u32::<LittleEndian>()?,
135            number_of_relocations: reader.read_u16::<LittleEndian>()?,
136            number_of_line_numbers: reader.read_u16::<LittleEndian>()?,
137            characteristics: reader.read_u32::<LittleEndian>()?,
138        })
139    }
140}
141
142/// 强制读取 COFF 头部（通用实现）
143pub(crate) fn read_coff_header<R: Read + Seek>(reader: &mut impl CoffReader<R>) -> Result<CoffHeader, GaiaError> {
144    // 保存当前位置
145    let original_pos = reader.get_position()?;
146
147    // 重置到文件开头
148    reader.set_position(0)?;
149
150    // 读取 COFF 头
151    let coff_header = CoffHeader::read(reader.get_viewer())?;
152
153    // 验证机器类型
154    match coff_header.machine {
155        0x014c | 0x8664 | 0x01c0 | 0xaa64 => {} // 支持的架构
156        unknown => {
157            let error = GaiaError::invalid_data(&format!("不支持的机器类型: 0x{:04x}", unknown));
158            reader.add_diagnostics(error);
159        }
160    }
161
162    // 验证节数量
163    if coff_header.number_of_sections == 0 {
164        let error = GaiaError::invalid_data("COFF 文件必须至少有一个节");
165        reader.add_diagnostics(error);
166    }
167
168    // 恢复原始位置
169    reader.set_position(original_pos)?;
170
171    Ok(coff_header)
172}
173
174/// 读取节头信息（通用实现）
175pub(crate) fn read_section_headers<R: Read + Seek>(reader: &mut impl CoffReader<R>) -> Result<Vec<SectionHeader>, GaiaError> {
176    // 先读取主头部
177    let header = reader.get_coff_header()?.clone();
178    let original_pos = reader.get_position()?;
179
180    // 读取节头部
181    let mut section_headers = Vec::with_capacity(header.number_of_sections as usize);
182
183    // 定位到节头部位置（COFF 头后面）
184    let section_header_offset = std::mem::size_of::<CoffHeader>() as u64 + header.size_of_optional_header as u64;
185
186    reader.set_position(section_header_offset)?;
187
188    for _ in 0..header.number_of_sections {
189        let section_header = SectionHeader::read(reader.get_viewer())?;
190        section_headers.push(section_header);
191    }
192
193    // 恢复原始位置
194    reader.set_position(original_pos)?;
195
196    Ok(section_headers)
197}
198
199/// 从节头读取节数据（通用实现）
200pub(crate) fn read_section_from_header<R: Read + Seek>(
201    reader: &mut impl CoffReader<R>,
202    header: &SectionHeader,
203) -> Result<CoffSection, GaiaError> {
204    let mut data = Vec::new();
205    let mut relocations = Vec::new();
206
207    // 读取节数据
208    if header.size_of_raw_data > 0 && header.pointer_to_raw_data > 0 {
209        let original_pos = reader.get_position()?;
210        reader.set_position(header.pointer_to_raw_data as u64)?;
211        data.resize(header.size_of_raw_data as usize, 0);
212        reader.get_viewer().read_exact(&mut data)?;
213        reader.set_position(original_pos)?;
214    }
215
216    // 读取重定位信息
217    if header.number_of_relocations > 0 && header.pointer_to_relocations > 0 {
218        let original_pos = reader.get_position()?;
219        reader.set_position(header.pointer_to_relocations as u64)?;
220
221        for _ in 0..header.number_of_relocations {
222            let relocation = CoffRelocation::read(reader.get_viewer())?;
223            relocations.push(relocation);
224        }
225
226        reader.set_position(original_pos)?;
227    }
228
229    Ok(CoffSection { header: *header, data, relocations })
230}
231
232/// 读取符号表（通用实现）
233pub(crate) fn read_symbols<R: Read + Seek>(reader: &mut impl CoffReader<R>) -> Result<Vec<CoffSymbol>, GaiaError> {
234    let header = reader.get_coff_header()?.clone();
235    let original_pos = reader.get_position()?;
236
237    // 计算符号表偏移
238    let symbol_table_offset = header.pointer_to_symbol_table as u64;
239    reader.set_position(symbol_table_offset)?;
240
241    let mut symbols = Vec::new();
242    for _ in 0..header.number_of_symbols {
243        let symbol = CoffSymbol::read(reader.get_viewer())?;
244        symbols.push(symbol);
245    }
246
247    reader.set_position(original_pos)?;
248    Ok(symbols)
249}
250
251/// 强制读取完整的 COFF 对象
252pub(crate) fn read_coff_object<R: Read + Seek>(reader: &mut impl CoffReader<R>) -> Result<CoffObject, GaiaError> {
253    let header = reader.get_coff_header()?.clone();
254    let section_headers = read_section_headers(reader)?;
255
256    // 读取节数据
257    let mut sections = Vec::new();
258    for section_header in &section_headers {
259        let section = read_section_from_header(reader, section_header)?;
260        sections.push(section);
261    }
262
263    // 读取符号表
264    let symbols = read_symbols(reader)?;
265
266    // 读取字符串表（简化实现）
267    let string_table = Vec::new();
268
269    Ok(CoffObject { header, sections, symbols, string_table })
270}
pe_assembler/helpers/coff_reader/mod.rs

pe_assembler/helpers/coff_reader/
mod.rs
