Skip to main content

pe_assembler/formats/exe/reader/
mod.rs

1use crate::{
2    helpers::{
3        pe_reader::{read_pe_head, read_pe_program, read_pe_section_headers},
4        PeReader,
5    },
6    types::{PeHeader, PeInfo, PeProgram, SectionHeader},
7};
8use gaia_types::{GaiaDiagnostics, GaiaError};
9use std::io::{Read, Seek, SeekFrom};
10
11/// EXE structure, lazy reader for parsing PE format executable files.
12#[derive(Debug)]
13pub struct ExeReader<R> {
14    /// The underlying reader.
15    reader: R,
16    /// The parsed PE header.
17    exe_header: Option<PeHeader>,
18    /// The parsed PE info.
19    exe_info: Option<PeInfo>,
20    /// The parsed section headers.
21    exe_section_headers: Option<Vec<SectionHeader>>,
22    /// The fully parsed PE program.
23    exe_program: Option<PeProgram>,
24    /// Collected errors during parsing.
25    errors: Vec<GaiaError>,
26}
27
28impl<R: Read> Read for ExeReader<R> {
29    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
30        self.reader.read(buffer)
31    }
32}
33
34impl<R: Seek> Seek for ExeReader<R> {
35    fn seek(&mut self, position: SeekFrom) -> std::io::Result<u64> {
36        self.reader.seek(position)
37    }
38}
39
40impl<R> ExeReader<R> {
41    /// Creates a new ExeReader with the specified reader.
42    pub fn new(reader: R) -> Self {
43        Self { reader, exe_header: None, exe_section_headers: None, exe_program: None, exe_info: None, errors: vec![] }
44    }
45
46    /// Finalizes the reading process and returns the parsed PE program.
47    pub fn finish(mut self) -> GaiaDiagnostics<PeProgram>
48    where
49        R: Read + Seek,
50    {
51        if self.exe_program.is_none() {
52            if let Err(e) = read_pe_program(&mut self) {
53                return GaiaDiagnostics { result: Err(e), diagnostics: self.errors };
54            }
55        }
56        unsafe {
57            let exe = self.exe_program.unwrap_unchecked();
58            GaiaDiagnostics { result: Ok(exe), diagnostics: self.errors }
59        }
60    }
61}
62
63impl<R: Read + Seek> PeReader<R> for ExeReader<R> {
64    fn get_viewer(&mut self) -> &mut R {
65        &mut self.reader
66    }
67
68    fn add_diagnostics(&mut self, error: impl Into<GaiaError>) {
69        self.errors.push(error.into());
70    }
71
72    fn get_section_headers(&mut self) -> Result<&[SectionHeader], GaiaError> {
73        if self.exe_section_headers.is_none() {
74            self.exe_section_headers = Some(read_pe_section_headers(self)?);
75        }
76        unsafe { Ok(self.exe_section_headers.as_ref().unwrap_unchecked()) }
77    }
78
79    fn get_pe_header(&mut self) -> Result<&PeHeader, GaiaError> {
80        if self.exe_header.is_none() {
81            self.exe_header = Some(read_pe_head(self)?)
82        }
83        unsafe { Ok(self.exe_header.as_ref().unwrap_unchecked()) }
84    }
85
86    fn get_program(&mut self) -> Result<&PeProgram, GaiaError> {
87        if self.exe_program.is_none() {
88            self.exe_program = Some(read_pe_program(self)?);
89        }
90        unsafe { Ok(self.exe_program.as_ref().unwrap_unchecked()) }
91    }
92}