pe_assembler/formats/exe/reader/
mod.rs

1use crate::{
2    helpers::{read_pe_head, read_pe_program, read_pe_section_headers, PeReader},
3    types::{PeHeader, PeInfo, PeProgram, SectionHeader},
4};
5use gaia_types::{GaiaDiagnostics, GaiaError};
6use std::io::{Read, Seek};
7
8/// EXE 结构,惰性读取器
9#[derive(Debug)]
10pub struct ExeReader<R> {
11    reader: R,
12    exe_header: Option<PeHeader>,
13    exe_info: Option<PeInfo>,
14    exe_section_headers: Option<Vec<SectionHeader>>,
15    exe_program: Option<PeProgram>,
16    errors: Vec<GaiaError>,
17}
18
19impl<R> ExeReader<R> {
20    pub fn new(reader: R) -> Self {
21        Self { reader, exe_header: None, exe_section_headers: None, exe_program: None, exe_info: None, errors: vec![] }
22    }
23
24    pub fn finish(mut self) -> GaiaDiagnostics<PeProgram>
25    where
26        R: Read + Seek,
27    {
28        if self.exe_program.is_none() {
29            if let Err(e) = read_pe_program(&mut self) {
30                return GaiaDiagnostics { result: Err(e), diagnostics: self.errors };
31            }
32        }
33        match self.exe_program {
34            Some(s) => GaiaDiagnostics { result: Ok(s), diagnostics: self.errors },
35            None => unreachable!(),
36        }
37    }
38}
39
40impl<R: Read + Seek> PeReader<R> for ExeReader<R> {
41    fn get_viewer(&mut self) -> &mut R {
42        &mut self.reader
43    }
44
45    fn add_diagnostics(&mut self, error: impl Into<GaiaError>) {
46        self.errors.push(error.into());
47    }
48
49    fn get_section_headers(&mut self) -> Result<&[SectionHeader], GaiaError> {
50        if self.exe_section_headers.is_none() {
51            read_pe_section_headers(self)?;
52        }
53        match self.exe_section_headers.as_ref() {
54            Some(s) => Ok(s),
55            None => unreachable!(),
56        }
57    }
58
59    fn set_section_headers(&mut self, headers: Vec<SectionHeader>) -> Vec<SectionHeader> {
60        self.exe_section_headers.replace(headers).unwrap_or_default()
61    }
62
63    fn get_pe_header(&mut self) -> Result<&PeHeader, GaiaError> {
64        if self.exe_header.is_none() {
65            match read_pe_head(self) {
66                Ok(header) => {
67                    self.exe_header = Some(header);
68                }
69                Err(e) => return Err(e),
70            }
71        }
72        match self.exe_header.as_ref() {
73            Some(s) => Ok(s),
74            None => Err(GaiaError::invalid_data("无法读取 PE 头部")),
75        }
76    }
77
78    fn set_pe_header(&mut self, head: PeHeader) -> Option<PeHeader> {
79        self.exe_header.replace(head)
80    }
81
82    fn get_program(&mut self) -> Result<&PeProgram, GaiaError> {
83        if self.exe_program.is_none() {
84            read_pe_program(self)?;
85        }
86        match self.exe_program.as_ref() {
87            Some(s) => Ok(s),
88            None => unreachable!(),
89        }
90    }
91
92    fn set_program(&mut self, program: PeProgram) -> Option<PeProgram> {
93        self.exe_program.replace(program)
94    }
95}