pe-assembler 0.1.1

PE/COFF assembler for Windows instruction sets - strongly typed, object-oriented, zero-dependency core
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

use crate::{
    helpers::{
        pe_reader::{read_pe_head, read_pe_program, read_pe_section_headers},
        PeReader,
    },
    types::{PeHeader, PeInfo, PeProgram, SectionHeader},
};
use gaia_types::{GaiaDiagnostics, GaiaError};
use std::io::{Read, Seek, SeekFrom};

/// DLL reader for parsing PE format DLL files.
#[derive(Debug)]
pub struct DllReader<R> {
    /// The underlying reader.
    reader: R,
    /// The parsed PE header.
    dll_header: Option<PeHeader>,
    /// The parsed PE info.
    dll_info: Option<PeInfo>,
    /// The parsed section headers.
    dll_section_headers: Option<Vec<SectionHeader>>,
    /// The fully parsed PE program.
    dll_program: Option<PeProgram>,
    /// Collected errors during parsing.
    errors: Vec<GaiaError>,
}

impl<R: Read> Read for DllReader<R> {
    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
        self.reader.read(buffer)
    }
}

impl<R: Seek> Seek for DllReader<R> {
    fn seek(&mut self, position: SeekFrom) -> std::io::Result<u64> {
        self.reader.seek(position)
    }
}

impl<R> DllReader<R> {
    /// Creates a new DllReader with the specified reader.
    pub fn new(reader: R) -> Self {
        Self { reader, dll_header: None, dll_section_headers: None, dll_program: None, dll_info: None, errors: vec![] }
    }
    /// Finalizes the reading process and returns the parsed PE program.
    pub fn finish(mut self) -> GaiaDiagnostics<PeProgram>
    where
        R: Read + Seek,
    {
        if self.dll_program.is_none() {
            if let Err(e) = read_pe_program(&mut self) {
                return GaiaDiagnostics { result: Err(e), diagnostics: self.errors };
            }
        }
        unsafe {
            let exe = self.dll_program.unwrap_unchecked();
            GaiaDiagnostics { result: Ok(exe), diagnostics: self.errors }
        }
    }
}

impl<R: Read + Seek> PeReader<R> for DllReader<R> {
    fn get_viewer(&mut self) -> &mut R {
        &mut self.reader
    }

    fn add_diagnostics(&mut self, error: impl Into<GaiaError>) {
        self.errors.push(error.into());
    }

    fn get_section_headers(&mut self) -> Result<&[SectionHeader], GaiaError> {
        if self.dll_section_headers.is_none() {
            self.dll_section_headers = Some(read_pe_section_headers(self)?);
        }
        unsafe { Ok(self.dll_section_headers.as_ref().unwrap_unchecked()) }
    }

    fn get_pe_header(&mut self) -> Result<&PeHeader, GaiaError> {
        if self.dll_header.is_none() {
            self.dll_header = Some(read_pe_head(self)?)
        }
        unsafe { Ok(self.dll_header.as_ref().unwrap_unchecked()) }
    }

    fn get_program(&mut self) -> Result<&PeProgram, GaiaError> {
        if self.dll_program.is_none() {
            self.dll_program = Some(read_pe_program(self)?);
        }
        unsafe { Ok(self.dll_program.as_ref().unwrap_unchecked()) }
    }
}