pe_assembler/formats/dll/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/// DLL 结构，惰性读取器
12#[derive(Debug)]
13pub struct DllReader<R> {
14    reader: R,
15    dll_header: Option<PeHeader>,
16    dll_info: Option<PeInfo>,
17    dll_section_headers: Option<Vec<SectionHeader>>,
18    dll_program: Option<PeProgram>,
19    errors: Vec<GaiaError>,
20}
21
22impl<R: Read> Read for DllReader<R> {
23    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
24        self.reader.read(buffer)
25    }
26}
27
28impl<R: Seek> Seek for DllReader<R> {
29    fn seek(&mut self, position: SeekFrom) -> std::io::Result<u64> {
30        self.reader.seek(position)
31    }
32}
33
34impl<R> DllReader<R> {
35    pub fn new(reader: R) -> Self {
36        Self { reader, dll_header: None, dll_section_headers: None, dll_program: None, dll_info: None, errors: vec![] }
37    }
38    pub fn finish(mut self) -> GaiaDiagnostics<PeProgram>
39    where
40        R: Read + Seek,
41    {
42        if self.dll_program.is_none() {
43            if let Err(e) = read_pe_program(&mut self) {
44                return GaiaDiagnostics { result: Err(e), diagnostics: self.errors };
45            }
46        }
47        unsafe {
48            let exe = self.dll_program.unwrap_unchecked();
49            GaiaDiagnostics { result: Ok(exe), diagnostics: self.errors }
50        }
51    }
52}
53
54impl<R: Read + Seek> PeReader<R> for DllReader<R> {
55    fn get_viewer(&mut self) -> &mut R {
56        &mut self.reader
57    }
58
59    fn add_diagnostics(&mut self, error: impl Into<GaiaError>) {
60        self.errors.push(error.into());
61    }
62
63    fn get_section_headers(&mut self) -> Result<&[SectionHeader], GaiaError> {
64        if self.dll_section_headers.is_none() {
65            self.dll_section_headers = Some(read_pe_section_headers(self)?);
66        }
67        unsafe { Ok(self.dll_section_headers.as_ref().unwrap_unchecked()) }
68    }
69
70    fn get_pe_header(&mut self) -> Result<&PeHeader, GaiaError> {
71        if self.dll_header.is_none() {
72            self.dll_header = Some(read_pe_head(self)?)
73        }
74        unsafe { Ok(self.dll_header.as_ref().unwrap_unchecked()) }
75    }
76
77    fn get_program(&mut self) -> Result<&PeProgram, GaiaError> {
78        if self.dll_program.is_none() {
79            self.dll_program = Some(read_pe_program(self)?);
80        }
81        unsafe { Ok(self.dll_program.as_ref().unwrap_unchecked()) }
82    }
83}
pe_assembler/formats/dll/reader/mod.rs

pe_assembler/formats/dll/reader/
mod.rs
