pe_assembler/formats/dll/reader/
mod.rs

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