feat: 增加从内存解析PE的功能

src/constants.rs

@@ -0,0 +1,15 @@
+// DOS的Magic值
+pub const DOS_MAGIC: u16 = 0x5A4D;
+// DOS头部大小
+pub const DOS_HEADER_SIZE: usize = 0x40;
+// DOS头部中PE头部的偏移
+pub const DOS_HEADER_PE_OFFSET: usize = 0x3C;
+// NT头的Magic值
+pub const NT_SIGNATURE: u32 = 0x00004550;
+// 节对齐值的偏移
+pub const SECTION_ALIGNMENT_OFFSET: usize = 0x38;
+// 镜像大小的偏移
+pub const IMAGE_SIZE_OFFSET: usize = 0x50;
+
+// 可选头Magic如果为32位的值
+pub const OPTIONAL_MAGIC_32: u16 = 0x10B;

src/error.rs

@@ -14,7 +14,7 @@ pub enum PEParseError {
     OutOfBounds,
     #[error("错误的RVA: {0}")]
     RvaToFoaError(u32),
-    #[error("RVA对应的是一个空节区数据！RVA: {0}")]
+    #[error("RVA对应的是一个空节区数据!RVA: {0}")]
     RvaToEmptySection(u32),
     #[error("错误的数据目录索引")]
     InvalidDataDirectoryIndex,

src/header.rs

@@ -1,4 +1,3 @@
-use super::types::*;
 use bitflags::bitflags;
 use serde::Serialize;
 
@@ -23,7 +22,7 @@ pub struct ImageDosHeader {
     pub e_oemid: u16,
     pub e_oeminfo: u16,
     pub e_res2: [u16; 10],
-    pub e_lfanew: Offset, // File address of new exe header nt头的偏移
+    pub e_lfanew: u32, // File address of new exe header nt头的偏移
 }
 
 #[repr(C)]
@@ -77,7 +76,7 @@ pub struct ImageFileHeader {
     pub machine: u16,
     pub number_of_sections: u16,
     pub time_date_stamp: u32,
-    pub pointer_to_symbol_table: Offset,
+    pub pointer_to_symbol_table: u32,
     pub number_of_symbols: u32,
     pub size_of_optional_header: u16,
     pub characteristics: FileCharacteristics,

src/lib.rs

@@ -1,8 +1,9 @@
 pub mod error;
 pub mod header;
 pub mod pe;
-pub mod types;
-
+pub mod memory_pe;
+mod constants;
+mod utils;
 #[cfg(feature = "memmap2_impl")]
 mod third_party {
     use memmap2::{MmapMut, Mmap};

src/memory_pe.rs

@@ -0,0 +1,128 @@
+use std::ops::{Deref, DerefMut};
+
+use crate::{align_to, error::PEParseError};
+
+use super::pe::{MutablePE, ReadOnlyPE};
+use crate::constants::*;
+pub struct MemoryPE {
+    len: usize,
+    buf: *const u8, // 缓冲区的指针 指向内存中的数据
+}
+
+pub struct MemoryPEMut {
+    len: usize,
+    buf: *mut u8, // 缓冲区的指针 指向内存中的数据
+}
+// DerefMut<Target = [u8]> + AsMut<[u8]>
+impl MutablePE for MemoryPEMut{}
+
+impl DerefMut for MemoryPEMut{
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        unsafe {
+            std::slice::from_raw_parts_mut(self.buf, self.len)
+        }
+    }
+}
+
+impl AsMut<[u8]> for MemoryPEMut{
+    fn as_mut(&mut self) -> &mut [u8] {
+        unsafe {
+            std::slice::from_raw_parts_mut(self.buf, self.len)
+        }
+    }
+}
+
+impl Deref for MemoryPEMut{
+    type Target = [u8];
+    fn deref(&self) -> &Self::Target {
+        unsafe {
+            std::slice::from_raw_parts(self.buf, self.len)
+        }
+    }
+}
+impl ReadOnlyPE for MemoryPEMut{}
+
+impl AsRef<[u8]> for MemoryPEMut{
+    fn as_ref(&self) -> &[u8] {
+        unsafe {
+            std::slice::from_raw_parts(self.buf, self.len)
+        }
+    }
+}
+
+impl ReadOnlyPE for MemoryPE{}
+
+impl Deref for MemoryPE{
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        unsafe {
+            std::slice::from_raw_parts(self.buf, self.len)
+        }
+    }
+}
+
+
+impl AsRef<[u8]> for MemoryPE{
+    fn as_ref(&self) -> &[u8] {
+        unsafe {
+            std::slice::from_raw_parts(self.buf, self.len)
+        }
+    }
+}
+
+/// 需要实现一个从内存中解析PE的方法
+pub fn parse_pe_from_memory(buf: *const u8) -> Result<MemoryPE, PEParseError> {
+    // 1. 检查是否为PE文件
+    // 取出头两个字节
+    let dos_magic = unsafe { *buf as u16 };
+    if dos_magic != 0x5A4D { // 小尾
+        return Err(PEParseError::InvalidDOSMagic);
+    }
+    // 从PE偏移处读取两个字节
+    let pe_offset = unsafe { *(buf.add(DOS_HEADER_PE_OFFSET) as *const u32) };
+    // 读取PE头部的标志
+    let nt_signature = unsafe { *(buf.add(pe_offset as usize) as *const u32) };
+    if nt_signature != NT_SIGNATURE {
+        return Err(PEParseError::InvalidNTSignature);
+    }
+    // 检查结束
+    // 从NT偏移处 读取一个NT header结构出来
+    // 需要读取节对齐和镜像大小
+    let section_alignment = unsafe { *(buf.add(pe_offset as usize + SECTION_ALIGNMENT_OFFSET) as *const u32) };
+    let image_size = unsafe { *(buf.add(pe_offset as usize + IMAGE_SIZE_OFFSET) as *const u32) };
+    // 计算对齐后的内存大小
+    let aligned_image_size = align_to!(image_size, section_alignment);
+    Ok(MemoryPE {
+        len: aligned_image_size as usize,
+        buf,
+    })
+}
+
+/// 还需要一个Mut的版本
+pub fn parse_pe_from_memory_mut(buf: *mut u8) -> Result<MemoryPEMut, PEParseError> {
+    // 1. 检查是否为PE文件
+    // 取出头两个字节
+    let dos_magic = unsafe { *buf as u16 };
+    if dos_magic != 0x5A4D { // 小尾
+        return Err(PEParseError::InvalidDOSMagic);
+    }
+    // 从PE偏移处读取两个字节
+    let pe_offset = unsafe { *(buf.add(DOS_HEADER_PE_OFFSET) as *const u32) };
+    // 读取PE头部的标志
+    let nt_signature = unsafe { *(buf.add(pe_offset as usize) as *const u32) };
+    if nt_signature != NT_SIGNATURE {
+        return Err(PEParseError::InvalidNTSignature);
+    }
+    // 检查结束
+    // 从NT偏移处 读取一个NT header结构出来
+    // 需要读取节对齐和镜像大小
+    let section_alignment = unsafe { *(buf.add(pe_offset as usize + SECTION_ALIGNMENT_OFFSET) as *const u32) };
+    let image_size = unsafe { *(buf.add(pe_offset as usize + IMAGE_SIZE_OFFSET) as *const u32) };
+    // 计算对齐后的内存大小
+    let aligned_image_size = align_to!(image_size, section_alignment);
+    Ok(MemoryPEMut {
+        len: aligned_image_size as usize,
+        buf,
+    })
+}

src/pe.rs

@@ -2,6 +2,8 @@ use super::{
     error::{MutablePEError, PEParseError},
     header::*,
 };
+use crate::align_to;
+use crate::constants::*;
 use std::{
     mem::{self},
     ops::{Deref, DerefMut},
@@ -22,14 +24,14 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
     fn is_64_bit(&self) -> Result<bool, PEParseError> {
         // 先以32位加载可选头，通过可选头的Magic字段判断
         let optional_header: &ImageOptionalHeader32 = self.get_optional_header_32()?;
-        let is_64_bit = optional_header.magic == 0x20b;
+        let is_64_bit = optional_header.magic == OPTIONAL_MAGIC_32;
         Ok(is_64_bit)
     }
 
     /// 获取数据目录
     fn get_data_directories(&self) -> Result<&[ImageDataDirectory], PEParseError> {
         // 1. 获取数据目录偏移
-    
+
         let data_directories_offset = self.get_data_directories_offset()?;
         // 2. 解析数据目录
         let data_directories = unsafe {
@@ -115,8 +117,7 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
             ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.section_alignment,
             ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.section_alignment,
         };
-        let aligned_size = (size + section_alignment - 1) & !(section_alignment - 1);
-        Ok(aligned_size)
+        Ok(align_to!(size, section_alignment))
     }
 
     /// 将size与文件对齐值进行对齐，返回对齐后的值
@@ -127,11 +128,11 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
             ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.file_alignment,
             ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.file_alignment,
         };
-        let aligned_size = (size + file_alignment - 1) & !(file_alignment - 1);
-        Ok(aligned_size)
+        Ok(align_to!(size, file_alignment))
     }
 
-    /// Get the DOS header without verifying its contents.
+    /// 获取DOS头
+    /// 此方法不会验证DOS头的合法性，它仅从0偏移处获取DOS头
     fn get_dos_header(&self) -> Result<&ImageDosHeader, PEParseError> {
         let result = self.get_ref::<ImageDosHeader>(0)?;
         Ok(result)
@@ -159,14 +160,14 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         Ok(result)
     }
 
-    // 获取nt头的偏移
+    /// 获取nt头的偏移
     fn get_nt_headers_offset(&self) -> Result<usize, PEParseError> {
         let dos_header = self.get_dos_header()?;
-        let nt_offset = dos_header.e_lfanew.0 as usize;
+        let nt_offset = dos_header.e_lfanew as usize;
         Ok(nt_offset)
     }
 
-    // 获取文件头的偏移
+    /// 获取文件头的偏移
     fn get_file_header_offset(&self) -> Result<usize, PEParseError> {
         // 1. 获取nt头偏移
         let nt_offset = self.get_nt_headers_offset()?;
@@ -175,20 +176,20 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         Ok(file_header_offset)
     }
 
-    // 获取文件头数据
+    /// 获取文件头数据
     fn get_file_header(&self) -> Result<&ImageFileHeader, PEParseError> {
         let file_header_offset = self.get_file_header_offset()?;
         let result = self.get_ref::<ImageFileHeader>(file_header_offset)?;
         Ok(result)
     }
 
-    // 获取节区数量
+    /// 获取节区数量
     fn get_number_of_sections(&self) -> Result<usize, PEParseError> {
         let file_header = self.get_file_header()?;
         Ok(file_header.number_of_sections as usize)
     }
 
-    // 获取可选头大小
+    /// 获取可选头大小
     fn get_size_of_optional_header(&self) -> Result<usize, PEParseError> {
         let file_header = self.get_file_header()?;
         Ok(file_header.size_of_optional_header as usize)
@@ -201,7 +202,7 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         Ok(result)
     }
 
-    // 作为64位的可选头解析
+    /// 作为64位的可选头解析
     fn get_optional_header_64(&self) -> Result<&ImageOptionalHeader64, PEParseError> {
         let file_header_offset = self.get_file_header_offset()?;
         let optional_header_offset = file_header_offset + std::mem::size_of::<ImageFileHeader>();
@@ -225,16 +226,16 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
     }
 
     /// 获取节区头的偏移
+    /// 算法：节区头偏移 = 可选头偏移 + 可选头大小
     fn get_section_headers_offset(&self) -> Result<usize, PEParseError> {
         // 节区头偏移在可选头之后，可选头大小是可变的，所以需要计算
         let optional_header_size = self.get_size_of_optional_header()?;
-        let file_header_offset = self.get_file_header_offset()?;
-        let section_header_offset =
-            file_header_offset + std::mem::size_of::<ImageFileHeader>() + optional_header_size;
+        let optional_header_offset = self.get_optional_header_offset()?;
+        let section_header_offset = optional_header_offset + optional_header_size;
         Ok(section_header_offset)
     }
 
-    // 获取所有节区头数据
+    /// 获取所有节区头数据
     fn get_section_headers(&self) -> Result<&[ImageSectionHeader], PEParseError> {
         // 1. 获取节区数量
         let number_of_sections = self.get_number_of_sections()?;
@@ -299,9 +300,8 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         Ok(export_table)
     }
 
-
     /// 获取导出函数的RVA地址，类似于GetProcAddress
-    fn get_export_function_rva(&self, proc_name: ProcName) -> Result<Option<u32>, PEParseError>{
+    fn get_export_function_rva(&self, proc_name: ProcName) -> Result<Option<u32>, PEParseError> {
         let export_table_dir = self.get_data_directory(ImageDirectoryEntry::Export)?;
         let export_table_rva = export_table_dir.virtual_address;
         let export_table = self.parse_export_table(export_table_rva)?;
@@ -311,7 +311,7 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         // 名称指针表的FOA
         let names_table_foa = self.rva_to_foa(export_table.address_of_names)?;
 
-        match proc_name{
+        match proc_name {
             ProcName::Named(proc_name_str) => {
                 // 遍历名称表
                 for index in 0..number_of_names {
@@ -328,12 +328,16 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
                     // 比较
                     if name == proc_name_str {
                         // 找到了这个函数 那么需要找到这个函数的序号
-                        let ordinals_table_foa = self.rva_to_foa(export_table.address_of_name_ordinals)?;
-                        let ordinal_item_foa = ordinals_table_foa + index * std::mem::size_of::<u16>() as u32;
+                        let ordinals_table_foa =
+                            self.rva_to_foa(export_table.address_of_name_ordinals)?;
+                        let ordinal_item_foa =
+                            ordinals_table_foa + index * std::mem::size_of::<u16>() as u32;
                         let ordinal = self.get_ref::<u16>(ordinal_item_foa as usize)?.clone();
                         // 找到了序号 那么就可以找到函数地址
-                        let functions_table_foa = self.rva_to_foa(export_table.address_of_functions)?;
-                        let function_item_foa = functions_table_foa + ordinal as u32 * std::mem::size_of::<u32>() as u32;
+                        let functions_table_foa =
+                            self.rva_to_foa(export_table.address_of_functions)?;
+                        let function_item_foa = functions_table_foa
+                            + ordinal as u32 * std::mem::size_of::<u32>() as u32;
                         let rva = self.get_ref::<u32>(function_item_foa as usize)?.clone();
                         return Ok(Some(rva));
                     }
@@ -347,13 +351,35 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
                     return Err(PEParseError::ExportOrdinalNotFound(order));
                 }
                 let functions_table_foa = self.rva_to_foa(export_table.address_of_functions)?;
-                let function_item_foa = functions_table_foa + find_ordinal * std::mem::size_of::<u32>() as u32;
+                let function_item_foa =
+                    functions_table_foa + find_ordinal * std::mem::size_of::<u32>() as u32;
                 let rva = self.get_ref::<u32>(function_item_foa as usize)?.clone();
                 return Ok(Some(rva));
             }
         }
         Ok(None)
     }
+
+    /// 以传入的RVA为开始偏移，从中解析导入表
+    fn parse_import_tables(&self, rva: u32) -> Result<&[ImageImportDescriptor], PEParseError> {
+        let foa = self.rva_to_foa(rva)?;
+        let mut current_foa = foa;
+        let mut numbers = 0;
+        loop {
+            let import_table = self.get_ref::<ImageImportDescriptor>(current_foa as usize)?;
+            if import_table.original_first_thunk == 0 {
+                break;
+            }
+            numbers += 1;
+            current_foa += std::mem::size_of::<ImageImportDescriptor>() as u32;
+        }
+        let result_tables = unsafe {
+            let ptr = self.as_ptr().wrapping_offset(foa as isize) as *const ImageImportDescriptor;
+            let result = std::slice::from_raw_parts(ptr, numbers);
+            result
+        };
+        Ok(result_tables)
+    }
 }
 /// 可修改的PE trait
 pub trait MutablePE: ReadOnlyPE + DerefMut<Target = [u8]> + AsMut<[u8]> {

src/types.rs

@@ -1,18 +0,0 @@
-use std::fmt::{Formatter, UpperHex};
-
-use serde::Serialize;
-
-
-#[repr(C)]
-#[derive(Debug, Clone, Copy, Serialize)]
-pub struct Offset(pub u32);
-
-impl UpperHex for Offset {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        self.0.fmt(f)
-    }
-}
-
-#[repr(C)]
-#[derive(Debug, Clone, Copy, Serialize)]
-pub struct RVA(pub u32);

src/utils/helper.rs

@@ -0,0 +1,7 @@
+/// 定义一个宏，让某值对齐到指定的字节大小
+#[macro_export]
+macro_rules! align_to {
+    ($value:expr, $align:expr) => {
+        ($value + $align - 1) & !($align - 1)
+    };
+}

src/utils/mod.rs

@@ -0,0 +1 @@
+pub mod helper;