v0.0.0b

sub/pe/src/imagebase.rs

@@ -0,0 +1,81 @@
+use core::ops::{Add, Range, Sub};
+use sub_core::cast_traits::Upcast;
+use crate::structs::*;
+
+#[repr(C)]
+pub struct ImageBase(());
+
+extern "C" {
+    #[link_name = "__ImageBase"]
+    static ImageBasePtr: [u8;0];
+}
+
+impl ImageBase {
+
+    #[inline(always)]
+    pub fn get() -> &'static ImageBase {
+        unsafe { &*(ImageBasePtr.as_ptr() as *const ImageBase) }
+    }
+
+
+    #[inline(always)]
+    pub fn as_ptr(&self) -> usize {
+        self as *const _ as usize
+    }
+
+    #[inline(always)]
+    pub(crate) fn offset<T: Upcast<usize>>(&self, offset: T) -> *const u8 {
+        unsafe { (self as *const _ as *const u8).add(offset.upcast()) }
+    }
+
+    pub unsafe fn dos(&self) -> &ImageDOSHeader {
+        &*(self as *const _ as *const ImageDOSHeader)
+    }
+
+    pub unsafe fn dos_mut(&mut self) -> &mut ImageDOSHeader {
+        &mut *(self as *mut _ as *mut ImageDOSHeader)
+    }
+
+    pub unsafe fn nt_header(&self) -> &'static ImageNTHeaders64 {
+        &*(self.offset(self.dos().e_lfanew) as *const ImageNTHeaders64)
+    }
+
+    pub unsafe fn nt_header_mut(&mut self) -> &'static mut ImageNTHeaders64 {
+        &mut *(self.offset(self.dos().e_lfanew) as *mut ImageNTHeaders64)
+    }
+
+    pub unsafe fn sections(&self) -> &[ImageSectionHeader] {
+        self.nt_header().sections()
+    }
+
+    pub unsafe fn as_range(&self) -> Range<usize> {
+        self.as_ptr()..self.as_ptr() + self.nt_header().optional_header.size_of_image as usize
+    }
+
+    pub unsafe fn as_slice(&self) -> &[u8] {
+        let ptr = self.as_ptr() as *const u8;
+        let size = self.nt_header().optional_header.size_of_image as usize;
+        core::slice::from_raw_parts(ptr, size)
+    }
+
+    pub unsafe fn as_slice_mut(&self) -> &mut [u8] {
+        let ptr = self.as_ptr() as *mut u8;
+        let size = self.nt_header().optional_header.size_of_image as usize;
+        core::slice::from_raw_parts_mut(ptr, size)
+    }
+
+}
+
+impl Add<usize> for &ImageBase {
+    type Output = usize;
+    fn add(self, rhs: usize) -> Self::Output {
+        (self as *const _ as usize) + rhs
+    }
+}
+
+impl Sub<usize> for &ImageBase {
+    type Output = usize;
+    fn sub(self, rhs: usize) -> Self::Output {
+        (self as *const _ as usize) + rhs
+    }
+}