pe_assembler/helpers/coff_reader/
mod.rs

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

pe_assembler/helpers/coff_reader/
mod.rs
