memexec/peparser/
header.rs

1use super::def::*;
2use super::error::{Error, Result};
3use std::mem;
4use std::os::raw::c_void;
5
6// Zero copy
7#[derive(Debug)]
8pub enum ImageNtHeaders<'a> {
9    x86(&'a IMAGE_NT_HEADERS32),
10    x64(&'a IMAGE_NT_HEADERS64),
11}
12
13impl<'a> ImageNtHeaders<'a> {
14    pub fn get_file_header(&self) -> &IMAGE_FILE_HEADER {
15        match *self {
16            ImageNtHeaders::x86(h) => &h.FileHeader,
17            ImageNtHeaders::x64(h) => &h.FileHeader,
18        }
19    }
20
21    pub fn get_address_of_entry_point(&self) -> isize {
22        match *self {
23            ImageNtHeaders::x86(h) => h.OptionalHeader.AddressOfEntryPoint as isize,
24            ImageNtHeaders::x64(h) => h.OptionalHeader.AddressOfEntryPoint as isize,
25        }
26    }
27
28    pub fn get_image_base(&self) -> *const c_void {
29        match *self {
30            ImageNtHeaders::x86(h) => h.OptionalHeader.ImageBase as *const c_void,
31            ImageNtHeaders::x64(h) => h.OptionalHeader.ImageBase as *const c_void,
32        }
33    }
34
35    pub fn get_size_of_image(&self) -> usize {
36        match *self {
37            ImageNtHeaders::x86(h) => h.OptionalHeader.SizeOfImage as usize,
38            ImageNtHeaders::x64(h) => h.OptionalHeader.SizeOfImage as usize,
39        }
40    }
41
42    pub fn get_file_alignment(&self) -> u32 {
43        match *self {
44            ImageNtHeaders::x86(h) => h.OptionalHeader.FileAlignment,
45            ImageNtHeaders::x64(h) => h.OptionalHeader.FileAlignment,
46        }
47    }
48
49    pub fn get_section_alignment(&self) -> u32 {
50        match *self {
51            ImageNtHeaders::x86(h) => h.OptionalHeader.SectionAlignment,
52            ImageNtHeaders::x64(h) => h.OptionalHeader.SectionAlignment,
53        }
54    }
55
56    pub fn get_data_directory(&self) -> &[IMAGE_DATA_DIRECTORY; IMAGE_NUMBEROF_DIRECTORY_ENTRIES] {
57        match *self {
58            ImageNtHeaders::x86(h) => &h.OptionalHeader.DataDirectory,
59            ImageNtHeaders::x64(h) => &h.OptionalHeader.DataDirectory,
60        }
61    }
62}
63
64// Zero copy
65#[repr(C)]
66#[derive(Debug)]
67pub struct PeHeader<'a> {
68    pub dos_header: &'a IMAGE_DOS_HEADER,
69    pub dos_stub: &'a [u8],
70    pub nt_header: ImageNtHeaders<'a>,
71}
72
73impl<'a> PeHeader<'a> {
74    pub fn new(bs: &'a [u8]) -> Result<PeHeader<'a>> {
75        let dos_header =
76            unsafe { &*mem::transmute::<*const u8, *const IMAGE_DOS_HEADER>(bs.as_ptr()) };
77        dos_header.is_valid()?;
78
79        let dos_stub = &bs[mem::size_of::<IMAGE_DOS_HEADER>()..(dos_header.e_lfanew as _)];
80
81        let nt_headers = unsafe {
82            match *mem::transmute::<*const u8, *const WORD>(bs.as_ptr().offset(
83                dos_header.e_lfanew as isize
84                    + mem::size_of::<DWORD>() as isize
85                    + mem::size_of::<IMAGE_FILE_HEADER>() as isize,
86            )) {
87                IMAGE_NT_OPTIONAL_HDR64_MAGIC => {
88                    ImageNtHeaders::x64(mem::transmute::<*const u8, &IMAGE_NT_HEADERS64>(
89                        bs.as_ptr().offset(dos_header.e_lfanew as isize),
90                    ))
91                }
92                IMAGE_NT_OPTIONAL_HDR32_MAGIC => {
93                    ImageNtHeaders::x86(mem::transmute::<*const u8, &IMAGE_NT_HEADERS32>(
94                        bs.as_ptr().offset(dos_header.e_lfanew as isize),
95                    ))
96                }
97                _ => return Err(Error::InvalidOptionalHeaderMagic),
98            }
99        };
100        nt_headers.is_valid()?;
101
102        Ok(PeHeader {
103            dos_header: dos_header,
104            dos_stub: dos_stub,
105            nt_header: nt_headers,
106        })
107    }
108
109    pub fn size(&self) -> usize {
110        mem::size_of::<IMAGE_DOS_HEADER>()
111            + self.dos_stub.len()
112            + if let ImageNtHeaders::x86(_) = self.nt_header {
113                mem::size_of::<IMAGE_NT_HEADERS32>()
114            } else {
115                mem::size_of::<IMAGE_NT_HEADERS64>()
116            }
117    }
118}
memexec/peparser/header.rs

memexec/peparser/
header.rs
