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

pe_assembler/formats/dll/reader/
mod.rs
