pe_assembler/formats/exe/reader/
mod.rs

1use crate::{
2    helpers::PeReader,
3    types::{PeHeader, PeInfo, PeProgram, SectionHeader},
4};
5use byteorder::LittleEndian;
6use gaia_types::{BinaryReader, GaiaDiagnostics, GaiaError};
7use std::io::{Read, Seek};
8
9/// EXE 结构,惰性读取器
10#[derive(Debug)]
11pub struct ExeReader<R> {
12    /// 二进制读取器(已定位到 DOS 头起始位置)
13    viewer: BinaryReader<R, LittleEndian>,
14    headers_read: bool,
15    lazy_header: Option<PeHeader>,
16    lazy_section_headers: Option<Vec<SectionHeader>>,
17    lazy_program: Option<PeProgram>,
18    lazy_info: Option<PeInfo>,
19    diagnostics: Vec<GaiaError>,
20}
21
22impl<R> ExeReader<R> {
23    pub fn new(reader: R) -> Self {
24        Self {
25            viewer: BinaryReader::new(reader),
26            headers_read: false,
27            lazy_header: None,
28            lazy_section_headers: None,
29            lazy_program: None,
30            lazy_info: None,
31            diagnostics: vec![],
32        }
33    }
34}
35
36impl<W: Read + Seek> PeReader<W> for ExeReader<W> {
37    fn get_viewer(&mut self) -> &mut BinaryReader<W, LittleEndian> {
38        &mut self.viewer
39    }
40
41    fn add_diagnostics(&mut self, _error: impl Into<GaiaError>) {
42        // 暂时不实现,避免编译错误
43    }
44
45    fn get_cached_section_headers(&self) -> Option<&Vec<SectionHeader>> {
46        self.lazy_section_headers.as_ref()
47    }
48
49    fn set_cached_section_headers(&mut self, headers: Vec<SectionHeader>) {
50        self.lazy_section_headers = Some(headers);
51    }
52
53    fn read_header_once(&mut self) -> Result<&PeHeader, GaiaError> {
54        if self.lazy_header.is_none() {
55            self.lazy_header = Some(self.read_header_force()?);
56        }
57        match self.lazy_header.as_ref() {
58            Some(s) => Ok(s),
59            None => unreachable!(),
60        }
61    }
62
63    fn read_program_once(&mut self) -> Result<&PeProgram, GaiaError> {
64        if self.lazy_program.is_none() {
65            self.lazy_program = Some(self.read_program_force()?);
66        }
67        match self.lazy_program.as_ref() {
68            Some(s) => Ok(s),
69            None => unreachable!(),
70        }
71    }
72}
73
74impl<W: Read + Seek> ExeReader<W> {
75    /// 读取完整的 PE 程序(惰性读取,会缓存结果)
76    pub fn read_program(mut self) -> GaiaDiagnostics<PeProgram> {
77        match self.read_program_once() {
78            Ok(_) => match self.lazy_program {
79                Some(program) => GaiaDiagnostics { result: Ok(program), diagnostics: self.diagnostics },
80                None => unreachable!(),
81            },
82            Err(error) => GaiaDiagnostics { result: Err(error), diagnostics: self.diagnostics },
83        }
84    }
85
86    /// 查看 EXE 文件信息
87    pub fn view(&mut self) -> Result<PeInfo, GaiaError> {
88        if let Some(ref info) = self.lazy_info {
89            return Ok(info.clone());
90        }
91
92        let info = self.create_pe_info()?;
93        self.lazy_info = Some(info.clone());
94        Ok(info)
95    }
96}