new pe loader

10 files changed, 402 insertions, 198 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 6d193dc..e1a6cbb 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -20,12 +20,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
 [[package]]
-name = "byte-slice-cast"
-version = "1.2.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c3ac9f8b63eca6fd385229b3675f6cc0dc5c8a5c8a54a59d4f52ffd670d87b0c"
-
-[[package]]
 name = "cfg-if"
 version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -43,7 +37,6 @@ version = "0.1.0"
 dependencies = [
  "aarch64-intrinsics",
  "bitflags",
- "byte-slice-cast",
  "const-utf16",
  "fdt",
  "linked_list_allocator",
diff --git a/Cargo.toml b/Cargo.toml
index 575155c..ef76573 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,7 +12,6 @@ mmio = "2.1.0"
 fdt = "0.1.5"
 bitflags = "1.3"
 const-utf16 = "0.2.1"
-byte-slice-cast = { version = "1.2.2", default-features = false }
 aarch64-intrinsics = { git = "https://github.com/ardbiesheuvel/aarch64-intrinsics" }
 spinning_top = "0.2.5"
 once_cell = { version = "1.18.0", default-features = false, features = ["race", "alloc"]  }
diff --git a/src/cmo.rs b/src/cmo.rs
index 38749a7..d4154a8 100644
--- a/src/cmo.rs
+++ b/src/cmo.rs
@@ -3,13 +3,14 @@
 // Author: Ard Biesheuvel <ardb@google.com>
 
 use core::arch::asm;
+use core::ops::Range;
 
 const CTR_IDC: u64 = 1 << 28;
 
 const CTR_DMINLINE_SHIFT: u64 = 16;
 const CTR_DMINLINE_MASK: u64 = 0xf;
 
-pub fn dcache_clean_to_pou(slice: &[u8]) {
+pub fn dcache_clean_to_pou(range: &Range<usize>) {
     let ctr = unsafe {
         let mut l: u64;
         asm!("mrs {reg}, ctr_el0", // CTR: cache type register
@@ -23,18 +24,18 @@ pub fn dcache_clean_to_pou(slice: &[u8]) {
     if (ctr & CTR_IDC) == 0 {
         let line_shift = 2 + ((ctr >> CTR_DMINLINE_SHIFT) & CTR_DMINLINE_MASK);
         let line_size: usize = 1 << line_shift;
-        let num_lines = (slice.len() + line_size - 1) >> line_shift;
-        let mut offset: usize = 0;
+        let len = range.end - range.start;
+        let num_lines = (len + line_size - 1) >> line_shift;
+        let mut line: usize = range.start;
 
         for _ in 0..num_lines {
             unsafe {
-                let base = slice.as_ptr();
                 asm!("dc cvau, {reg}",
-                    reg = in(reg) base.offset(offset as isize),
+                    reg = in(reg) line,
                     options(nomem, nostack, preserves_flags),
                 );
             }
-            offset += line_size;
+            line += line_size;
         }
     }
 }
diff --git a/src/efi/loadedimage.rs b/src/efi/loadedimage.rs
index 1492a3c..c467f4c 100644
--- a/src/efi/loadedimage.rs
+++ b/src/efi/loadedimage.rs
@@ -7,6 +7,7 @@ use crate::efi::install_protocol;
 use crate::efi::new_handle;
 use crate::efi::uninstall_protocol;
 use crate::efi::{memorytype::*, status::*, systemtable::*, Guid, Handle};
+use crate::PeImage;
 
 use core::marker::PhantomData;
 use core::ptr;
@@ -60,11 +61,10 @@ pub struct LoadedImageData<'a> {
 
 impl<'a> LoadedImageData<'a> {
     pub fn new(
-        buffer: &'a [u8],
+        pe_image: &'a PeImage,
         load_options: &'a [u16],
         code_type: MemoryType,
         data_type: MemoryType,
-        entrypoint: *const u8,
         randomized: bool,
     ) -> LoadedImageData<'a> {
         let handle: Handle = new_handle();
@@ -77,8 +77,8 @@ impl<'a> LoadedImageData<'a> {
             reserved: usize::MAX,
             load_options_size: (load_options.len() * core::mem::size_of::<u16>()) as u32,
             load_options: load_options.as_ptr() as *const (),
-            image_base: buffer.as_ptr() as *const (),
-            image_size: buffer.len() as u64,
+            image_base: pe_image.image_base(),
+            image_size: pe_image.image_size(),
             image_code_type: code_type,
             image_data_type: data_type,
             unload: unload,
@@ -86,7 +86,7 @@ impl<'a> LoadedImageData<'a> {
         let li = LoadedImageData {
             loaded_image: alloc::boxed::Box::leak(p),
             image_handle: handle,
-            entrypoint: entrypoint,
+            entrypoint: pe_image.entry_point(),
             randomized: randomized,
             marker: PhantomData,
         };
diff --git a/src/efi/mod.rs b/src/efi/mod.rs
index 44e72ce..620397b 100644
--- a/src/efi/mod.rs
+++ b/src/efi/mod.rs
@@ -10,8 +10,7 @@ use crate::efi::runtimeservices::RuntimeServices;
 use crate::efi::MemoryType::*;
 use crate::efi::{loadedimage::*, memorytype::*, systemtable::*};
 use crate::RTSCODE;
-
-use crate::pecoff::Parser;
+use crate::PeImage;
 
 use core::sync::atomic::{AtomicUsize, Ordering};
 
@@ -43,6 +42,7 @@ mod simpletext;
 pub mod status;
 mod systemtable;
 mod tableheader;
+pub mod peloader;
 
 pub type Bool = u8;
 pub type Char16 = u16;
@@ -94,6 +94,13 @@ pub trait FileLoader {
     fn get_size(&self) -> usize;
 
     fn load_file<'a>(&self, loadbuffer: &'a mut [MaybeUninit<u8>]) -> Result<&'a [u8], &str>;
+
+    unsafe fn load_range<'a>(
+        &self,
+        loadbuffer: *mut (),
+        offset: usize,
+        size: usize,
+    ) -> Result<(), &str>;
 }
 
 pub const EFI_RT_PROPERTIES_TABLE_GUID: Guid = guid!(
@@ -273,16 +280,15 @@ impl EfiContext {
 
     pub fn load_image<'a>(
         &self,
-        pe_image: &'a Parser,
+        pe_image: &'a PeImage,
         cmdline: &'a Vec<u16>,
         randomized: bool,
     ) -> LoadedImageData<'a> {
         LoadedImageData::new(
-            pe_image.get_image(),
+            pe_image,
             cmdline.as_slice(),
             EfiLoaderCode,
             EfiLoaderData,
-            pe_image.get_entrypoint(),
             randomized,
         )
     }
diff --git a/src/efi/peloader.rs b/src/efi/peloader.rs
new file mode 100644
index 0000000..c3445ea
--- /dev/null
+++ b/src/efi/peloader.rs
@@ -0,0 +1,285 @@
+// SPDX-License-Identifier: GPL-2.0
+// Copyright 2023 Google LLC
+// Author: Ard Biesheuvel <ardb@google.com>
+
+use crate::efi::memmap;
+use crate::{FileLoader, MemoryType, Placement};
+
+use alloc::boxed::Box;
+use alloc::vec::Vec;
+use core::mem::{size_of, MaybeUninit};
+use core::ops::Range;
+use core::str::from_utf8;
+use log::{debug, trace};
+
+#[derive(Copy, Clone)]
+#[repr(C)]
+struct DosHeader {
+    magic: [u8; 2],
+    dontcare: [u8; 58],
+    pe_offset: u32,
+}
+
+#[derive(Copy, Clone, Debug)]
+#[repr(C)]
+struct PeHeader {
+    signature: [u8; 4],
+    machine: u16,
+    number_of_sections: u16,
+    time_date_stamp: u32,
+    pointer_to_symbol_table: u32,
+    number_of_symbols: u32,
+    size_of_optional_header: u16,
+    characteristics: u16,
+
+    magic: u16,
+    major_linker_version: u8,
+    minor_linker_version: u8,
+    size_of_code: u32,
+    size_of_initialized_data: u32,
+    size_of_uninitialized_data: u32,
+    address_of_entrypoint: u32,
+    base_of_code: u32,
+    image_base: u64,
+    section_alignment: u32,
+    file_alignment: u32,
+    major_os_version: u16,
+    minor_os_version: u16,
+    major_image_version: u16,
+    minor_image_version: u16,
+    major_subsys_version: u16,
+    minor_subsys_version: u16,
+    win32_version_value: u32,
+    size_of_image: u32,
+    size_of_headers: u32,
+    checksum: u32,
+    dll_characteristics: u16,
+    size_of_stack_reserve: u64,
+    size_of_stack_commit: u64,
+    size_of_heap_reserve: u64,
+    size_of_heap_commit: u64,
+    loader_flags: u32,
+    number_of_rva_and_sizes: u32,
+}
+
+#[derive(Copy, Clone, Debug)]
+#[repr(C)]
+struct PeSection {
+    name: [u8; 8],
+    virtual_size: u32,
+    virtual_address: u32,
+    size_of_raw_data: u32,
+    pointer_to_raw_data: u32,
+    pointer_to_relocations: u32,
+    pointer_to_line_numbers: u32,
+    number_of_relocations: u16,
+    number_of_line_numbers: u16,
+    characteristics: u32,
+}
+
+pub const EFI_IMAGE_SCN_MEM_EXECUTE: u32 = 0x20000000;
+#[allow(dead_code)]
+pub const EFI_IMAGE_SCN_MEM_READ: u32 = 0x40000000;
+#[allow(dead_code)]
+pub const EFI_IMAGE_SCN_MEM_WRITE: u32 = 0x80000000;
+
+impl PeSection {
+    fn get_name(&self) -> &str {
+        from_utf8(&self.name).unwrap()
+    }
+}
+
+pub struct PeLoader<'a> {
+    pe_header: PeHeader,
+    sections: Vec<PeSection>,
+    file_loader: &'a dyn FileLoader,
+}
+
+impl<'a> PeLoader<'a> {
+    pub fn new(loader: &'a dyn FileLoader) -> Option<PeLoader> {
+        let doshdr = {
+            let mut h = Box::pin(MaybeUninit::<DosHeader>::uninit());
+            unsafe {
+                loader
+                    .load_range(&mut *h as *mut _ as *mut (), 0, size_of::<DosHeader>())
+                    .ok()?;
+                h.assume_init()
+            }
+        };
+
+        if doshdr.magic != ['M' as u8, 'Z' as u8] {
+            debug!("Invalid DOS magic 0x{:x?}\n", doshdr.magic);
+            return None;
+        }
+
+        if (doshdr.pe_offset as usize) < size_of::<DosHeader>()
+            || (doshdr.pe_offset as usize) + size_of::<PeHeader>() > loader.get_size()
+        {
+            debug!("Invalid PE header offset 0x{:x?}\n", doshdr.pe_offset);
+            return None;
+        }
+
+        let pehdr = {
+            let mut h = Box::pin(MaybeUninit::<PeHeader>::uninit());
+            unsafe {
+                loader
+                    .load_range(
+                        &mut *h as *mut _ as *mut (),
+                        doshdr.pe_offset as usize,
+                        size_of::<PeHeader>(),
+                    )
+                    .ok()?;
+                h.assume_init()
+            }
+        };
+
+        if pehdr.signature != ['P' as u8, 'E' as u8, 0u8, 0u8] {
+            debug!("Invalid PE magic 0x{:x?}\n", pehdr.signature);
+            return None;
+        }
+
+        trace!(
+            "PE header at offset 0x{:x?}: {:x?}\n",
+            doshdr.pe_offset,
+            pehdr
+        );
+
+        let section_offset = doshdr.pe_offset + 24 + pehdr.size_of_optional_header as u32;
+        let section_count = pehdr.number_of_sections as usize;
+        let sections_size = size_of::<PeSection>() * section_count;
+        if section_offset as usize + sections_size > loader.get_size() {
+            debug!("Section array runs past the end of the image\n");
+            return None;
+        }
+        let mut sections = {
+            let mut v = Vec::<PeSection>::with_capacity(section_count + 1);
+            unsafe {
+                loader
+                    .load_range(
+                        v.as_mut_ptr() as *mut (),
+                        section_offset as usize,
+                        sections_size,
+                    )
+                    .ok()?;
+                v.set_len(section_count);
+            }
+            v
+        };
+        trace!("Section headers: {:x?}\n", sections);
+
+        // Insert a dummy section for the header region
+        sections.push(PeSection {
+            name: [0u8; 8],
+            virtual_size: pehdr.size_of_headers,
+            virtual_address: 0,
+            size_of_raw_data: pehdr.size_of_headers,
+            pointer_to_raw_data: 0,
+            pointer_to_relocations: 0,
+            pointer_to_line_numbers: 0,
+            number_of_relocations: 0,
+            number_of_line_numbers: 0,
+            characteristics: EFI_IMAGE_SCN_MEM_READ,
+        });
+
+        for s in sections.iter() {
+            if (s.pointer_to_raw_data | s.size_of_raw_data) & (pehdr.file_alignment - 1) != 0 {
+                debug!(
+                    "Section {} violates file alignment {:x}\n",
+                    s.get_name(),
+                    pehdr.file_alignment
+                );
+                return None;
+            }
+
+            if s.virtual_address & (pehdr.section_alignment - 1) != 0 {
+                debug!(
+                    "Section {} violates section alignment {:x}\n",
+                    s.get_name(),
+                    pehdr.section_alignment
+                );
+                return None;
+            }
+        }
+
+        Some(PeLoader {
+            pe_header: pehdr,
+            sections: sections,
+            file_loader: loader,
+        })
+    }
+
+    pub fn load(self, memory_type: MemoryType, placement: Placement) -> Option<PeImage<'a>> {
+        let buf = memmap::allocate_pages(
+            memmap::size_to_pages(self.pe_header.size_of_image as usize),
+            memory_type,
+            placement,
+        )?;
+
+        for s in self.sections.iter() {
+            unsafe {
+                self.file_loader
+                    .load_range(
+                        buf[s.virtual_address as usize].as_mut_ptr() as *mut (),
+                        s.pointer_to_raw_data as usize,
+                        s.size_of_raw_data as usize,
+                    )
+                    .ok()?;
+            }
+            // TODO zero out remaining space
+        }
+
+        // TODO apply relocations
+
+        Some(PeImage {
+            pe_loader: self,
+            loaded_image: buf,
+        })
+    }
+}
+
+pub struct PeImage<'a> {
+    pe_loader: PeLoader<'a>,
+    loaded_image: &'a [MaybeUninit<u8>],
+}
+
+impl PeImage<'_> {
+    pub fn image_base(&self) -> *const () {
+        self.loaded_image.as_ptr() as _
+    }
+
+    pub fn image_size(&self) -> u64 {
+        self.pe_loader.pe_header.size_of_image as _
+    }
+
+    pub fn entry_point(&self) -> *const u8 {
+        self.loaded_image[self.pe_loader.pe_header.address_of_entrypoint as usize].as_ptr() as _
+    }
+
+    pub fn sections(&self) -> PeImageSectionIterator {
+        PeImageSectionIterator {
+            index: 0,
+            pe_image: self,
+        }
+    }
+}
+
+pub struct PeImageSectionIterator<'a> {
+    index: usize,
+    pe_image: &'a PeImage<'a>,
+}
+
+impl Iterator for PeImageSectionIterator<'_> {
+    type Item = (Range<usize>, u32);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index >= self.pe_image.pe_loader.sections.len() {
+            return None;
+        }
+        let s = &self.pe_image.pe_loader.sections[self.index];
+        let start = self.pe_image.loaded_image[s.virtual_address as usize].as_ptr() as usize;
+        let end = start + s.virtual_size as usize;
+        self.index += 1;
+        Some((start..end, s.characteristics))
+    }
+}
+
diff --git a/src/fwcfg.rs b/src/fwcfg.rs
index c028caf..a65f8fd 100644
--- a/src/fwcfg.rs
+++ b/src/fwcfg.rs
@@ -117,7 +117,9 @@ impl<T: Copy> Iterator for FwCfgFileIterator<'_, T> {
         mmio.selector.write(u16::to_be(self.select));
         fence(Ordering::Release);
 
-        mmio.dma_transfer(0, offset as u32).ok()?;
+        if offset > 0 {
+            mmio.dma_transfer(0, offset as u32).ok()?;
+        }
         mmio.dma_transfer(&*out as *const _ as u64, itemsz as u32)
             .ok()?;
         self.next += 1;
@@ -285,15 +287,20 @@ impl FwCfg {
         Some(Self::new(addr))
     }
 
-    fn dma_read(
+    fn dma_read<T: Copy>(
         &self,
-        loadbuffer: &mut [MaybeUninit<u8>],
+        loadbuffer: &mut [MaybeUninit<T>],
+        offset: usize,
         size: usize,
         config_item: u16,
     ) -> Result<(), ()> {
         let mut mmio = self.0.lock();
         mmio.selector.write(u16::to_be(config_item));
         fence(Ordering::Release);
+
+        if offset > 0 {
+            mmio.dma_transfer(0, offset as u32).or(Err(()))?;
+        }
         mmio.dma_transfer(loadbuffer.as_ptr() as u64, size as u32)
     }
 
@@ -304,14 +311,15 @@ impl FwCfg {
         mmio.data.read() as usize
     }
 
-    fn load_file<'a>(
+    fn load_file<'a, T: Copy>(
         &self,
-        loadbuffer: &'a mut [MaybeUninit<u8>],
+        loadbuffer: &'a mut [MaybeUninit<T>],
+        offset: usize,
         size: usize,
         data_cfg: u16,
-    ) -> Result<&'a [u8], ()> {
-        let size = size.min(loadbuffer.len());
-        self.dma_read(loadbuffer, size, data_cfg)?;
+    ) -> Result<&'a [T], ()> {
+        let size = size.min(loadbuffer.len() / size_of::<T>());
+        self.dma_read(loadbuffer, offset, size, data_cfg)?;
         if size < loadbuffer.len() {
             loadbuffer.split_at_mut(size).1.fill(MaybeUninit::zeroed());
         }
@@ -326,11 +334,12 @@ impl FwCfg {
     fn load_file_mut<'a>(
         &self,
         loadbuffer: &'a mut [MaybeUninit<u8>],
+        offset: usize,
         size: usize,
         data_cfg: u16,
     ) -> Result<&'a mut [u8], &'static str> {
         let size = size.min(loadbuffer.len());
-        self.dma_read(loadbuffer, size, data_cfg)
+        self.dma_read(loadbuffer, offset, size, data_cfg)
             .or(Err("DMA read failed"))?;
         if size < loadbuffer.len() {
             loadbuffer.split_at_mut(size).1.fill(MaybeUninit::zeroed());
@@ -393,7 +402,7 @@ impl<'a> FwCfgTableLoader<'a> {
             Placement::Anywhere,
         )
         .ok_or("Failed to allocate blob memory")?;
-        let b = self.fwcfg.load_file_mut(b, f.size(), f.select())?;
+        let b = self.fwcfg.load_file_mut(b, 0, f.size(), f.select())?;
         self.loaded_tables.insert(f.filename, b);
         Ok(())
     }
@@ -469,50 +478,28 @@ trait LeBytes<T, const N: usize> {
     fn to_le_bytes(val: T) -> [u8; N];
 }
 
-impl LeBytes<u8, 1> for u8 {
-    fn from_le_bytes(val: &[u8; 1]) -> u8 {
-        val[0]
-    }
-
-    fn to_le_bytes(val: u8) -> [u8; 1] {
-        [val]
-    }
-}
-
-impl LeBytes<u16, 2> for u16 {
-    fn from_le_bytes(val: &[u8; 2]) -> u16 {
-        u16::from_le_bytes(*val)
-    }
-
-    fn to_le_bytes(val: u16) -> [u8; 2] {
-        u16::to_le_bytes(val)
-    }
-}
-
-impl LeBytes<u32, 4> for u32 {
-    fn from_le_bytes(val: &[u8; 4]) -> u32 {
-        u32::from_le_bytes(*val)
-    }
-
-    fn to_le_bytes(val: u32) -> [u8; 4] {
-        u32::to_le_bytes(val)
-    }
+macro_rules! le_bytes {
+    ($a: ty) => {
+        impl LeBytes<$a, { size_of::<$a>() }> for $a {
+            fn from_le_bytes(val: &[u8; { size_of::<$a>() }]) -> $a {
+                <$a>::from_le_bytes(*val)
+            }
+            fn to_le_bytes(val: $a) -> [u8; { size_of::<$a>() }] {
+                <$a>::to_le_bytes(val)
+            }
+        }
+    };
 }
 
-impl LeBytes<u64, 8> for u64 {
-    fn from_le_bytes(val: &[u8; 8]) -> u64 {
-        u64::from_le_bytes(*val)
-    }
-
-    fn to_le_bytes(val: u64) -> [u8; 8] {
-        u64::to_le_bytes(val)
-    }
-}
+le_bytes!(u8);
+le_bytes!(u16);
+le_bytes!(u32);
+le_bytes!(u64);
 
 fn add_value_at_offset<T: LeBytes<T, N> + core::ops::Add<Output = T>, const N: usize>(
     addend: T,
     offset: usize,
-    buf: &mut &mut [u8],
+    buf: &mut [u8],
 ) {
     let end = offset + N;
     let val: T = T::from_le_bytes(buf[offset..end].try_into().unwrap()) + addend;
@@ -532,7 +519,22 @@ impl efi::FileLoader for FwCfgFileLoader<'_> {
 
     fn load_file<'a>(&self, loadbuffer: &'a mut [MaybeUninit<u8>]) -> Result<&'a [u8], &str> {
         self.fwcfg
-            .load_file(loadbuffer, self.size, self.data_cfg)
+            .load_file(loadbuffer, 0, self.size, self.data_cfg)
             .or(Err("Failed to load file from fwcfg"))
     }
+
+    unsafe fn load_range<'a>(
+        &self,
+        ptr: *mut (),
+        offset: usize,
+        size: usize,
+    ) -> Result<(), &str> {
+        if offset > self.size {
+            return Err("Offset out of range");
+        }
+        let loadbuffer = slice::from_raw_parts_mut(ptr as *mut MaybeUninit<u8>, size);
+        self.fwcfg
+            .load_file(loadbuffer, offset, size, self.data_cfg)
+            .or(Err("Failed to load range from fwcfg")).and(Ok(()))
+    }
 }
diff --git a/src/idmap.rs b/src/idmap.rs
index e3f6d7a..cb145fa 100644
--- a/src/idmap.rs
+++ b/src/idmap.rs
@@ -5,6 +5,7 @@
 use crate::paging::*;
 
 use core::arch::asm;
+use core::ops::Range;
 
 // Use a different ASID for the full ID map, as well as non-global attributes
 // for all its DRAM mappings. This way, we can ignore break-before-make rules
@@ -85,6 +86,18 @@ impl IdMap {
         self.root.map_range(&c, c.0.start as u64, flags);
     }
 
+    pub fn map_range_(&mut self, range: &Range<usize>, flags: Attributes) {
+        let c = Chunk(range.start..range.end);
+
+        log::info!(
+            "Mapping memory at [0x{:X} - 0x{:X}] {:?}\n",
+            c.0.start,
+            c.0.end - 1,
+            flags
+        );
+        self.root.map_range(&c, c.0.start as u64, flags);
+    }
+
     pub fn walk_range<F>(&self, slice: &[u8], f: &mut F)
     where
         F: FnMut(&Descriptor, usize, usize),
diff --git a/src/main.rs b/src/main.rs
index 58b5a07..add4f3f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -25,13 +25,12 @@ mod idmap;
 mod initrd;
 mod pagealloc;
 mod paging;
-mod pecoff;
 mod psci;
 mod rng;
 
 use core::{arch::global_asm, panic::PanicInfo, slice};
 use linked_list_allocator::LockedHeap;
-use log::{error, info, debug};
+use log::{error, warn, info, debug};
 
 extern crate alloc;
 use alloc::vec::Vec;
@@ -40,9 +39,10 @@ extern crate aarch64_intrinsics;
 
 use crate::efi::{Guid, guid};
 use crate::efi::FileLoader;
-use crate::efi::memmap::Placement;
+use crate::efi::memmap::*;
 use crate::efi::memorytype::*;
 use crate::paging::Attributes;
+use crate::efi::peloader::*;
 use crate::MemoryType::EfiBootServicesData;
 
 #[macro_use]
@@ -90,6 +90,8 @@ extern "C" fn efilite_main(base: *mut u8, mapped: usize, used: isize, avail: usi
         log::set_logger(&console::OUT).ok()?;
         #[cfg(debug_assertions)]
         log::set_max_level(log::LevelFilter::Debug);
+        #[cfg(not(debug_assertions))]
+        log::set_max_level(log::LevelFilter::Warn);
         info!("Using {:?} for console output\n", n.name);
         Some(())
     });
@@ -165,50 +167,31 @@ extern "C" fn efilite_main(base: *mut u8, mapped: usize, used: isize, avail: usi
         (Placement::Aligned(0x20000), false)
     };
 
-    let pe_image = {
-        let ldr = fwcfg.get_kernel_loader().unwrap();
-        let buf = efi::memmap::allocate_pages(
-            efi::memmap::size_to_pages(ldr.get_size()),
-            MemoryType::EfiLoaderData,
-            Placement::Anywhere,
-        )
-        .expect("Failed to allocate memory for kernel image");
-
-        let buf = ldr
-            .load_file(buf)
-            .expect("Failed to load image header");
-
-        let size = pecoff::Parser::get_image_size(&buf).expect("Failed to parse PE/COFF header");
-
-        let buf = efi::memmap::allocate_pages(
-            efi::memmap::size_to_pages(size),
-            MemoryType::EfiLoaderCode,
-            placement,
-        )
-        .expect("Failed to allocate memory for EFI program");
-
-        let loadbuffer = ldr
-            .load_file(buf)
-            .expect("Failed to load kernel image");
-
-        pecoff::Parser::from_slice(loadbuffer).expect("Failed to parse PE/COFF image")
-    };
-
-
-    // Clean the code region of the loaded image to the PoU so we
-    // can safely fetch instructions from it once the PXN/UXN
-    // attributes are cleared
-    let code = pe_image.get_code_region();
-    cmo::dcache_clean_to_pou(code);
-
-    // Switch back to the initial ID map so we can remap
-    // the loaded kernel image with different permissions
-    idmap.deactivate();
-
-    // Remap the text/rodata part of the image read-only so we will
-    // be able to execute it with WXN protections enabled
-    idmap.map_range(code, ro_flags.non_global());
-    idmap.activate();
+    let ldr = fwcfg.get_kernel_loader().expect("No kernel image provided");
+    let pe_ldr = PeLoader::new(&ldr).unwrap();
+    let pe_image = pe_ldr.load(MemoryType::EfiLoaderCode, placement).unwrap();
+
+    for s in pe_image.sections() {
+        if (s.1 & EFI_IMAGE_SCN_MEM_WRITE) == 0 {
+            let f = if (s.1 & EFI_IMAGE_SCN_MEM_EXECUTE) != 0 {
+                // Clean the code regions of the loaded image to the PoU so we
+                // can safely fetch instructions from them once the PXN/UXN
+                // attributes are cleared
+                cmo::dcache_clean_to_pou(&s.0);
+                ro_flags.non_global()
+            } else {
+                ro_flags.non_global().execute_disable()
+            };
+
+            if (s.0.start | s.0.end) & EFI_PAGE_MASK == 0 {
+                idmap.deactivate();
+                idmap.map_range_(&s.0, f);
+                idmap.activate();
+            } else {
+                warn!("Cannot remap range {:?}\n", s.0);
+            }
+        }
+    }
 
     let initrdloader = fwcfg.get_initrd_loader().unwrap();
     let _initrd = efi::initrdloadfile2::new(&initrdloader);
diff --git a/src/pecoff.rs b/src/pecoff.rs
deleted file mode 100644
index a1dfec9..0000000
--- a/src/pecoff.rs
+++ /dev/null
@@ -1,78 +0,0 @@
-// SPDX-License-Identifier: GPL-2.0
-// Copyright 2022-2023 Google LLC
-// Author: Ard Biesheuvel <ardb@google.com>
-
-use byte_slice_cast::*;
-
-pub struct Parser<'a> {
-    image: &'a [u8],
-    code: &'a [u8],
-    entrypoint: &'a [u8],
-}
-
-impl<'a> Parser<'a> {
-    fn get_pe_header(slice: &'a [u8]) -> Option<&'a [u8]> {
-        // Check that the image starts with the magic bytes 'MZ'
-        if slice[0] != b'M' || slice[1] != b'Z' || slice.len() < 64 {
-            return None;
-        }
-        // Get the PE header offset from the MSDOS header
-        let offset: usize = slice.get(60..=63).unwrap().as_slice_of::<u32>().unwrap()[0] as _;
-        let pehdr = &slice.get(offset..).unwrap();
-        if pehdr[0] != b'P' || pehdr[1] != b'E' || pehdr.len() < 84 {
-            return None;
-        }
-        Some(pehdr)
-    }
-
-    pub fn get_image_size(slice: &'a [u8]) -> Option<usize> {
-        if let Some(pehdr) = Self::get_pe_header(slice) {
-            let size_of_image: usize =
-                pehdr.get(80..=83).unwrap().as_slice_of::<u32>().unwrap()[0] as _;
-            Some(size_of_image)
-        } else {
-            None
-        }
-    }
-
-    pub fn from_slice(slice: &'a [u8]) -> Option<Parser<'a>> {
-        let pehdr = Self::get_pe_header(slice).unwrap();
-
-        let base_of_code: usize = pehdr.get(44..=47).unwrap().as_slice_of::<u32>().unwrap()[0] as _;
-        let size_of_code: usize = pehdr.get(28..=31).unwrap().as_slice_of::<u32>().unwrap()[0] as _;
-        let entrypoint: usize = pehdr.get(40..=43).unwrap().as_slice_of::<u32>().unwrap()[0] as _;
-        let size_of_image: usize =
-            pehdr.get(80..=83).unwrap().as_slice_of::<u32>().unwrap()[0] as _;
-
-        // Check that the various bounds are within the slice
-        if base_of_code >= slice.len()
-            || (base_of_code + size_of_code) > slice.len()
-            || entrypoint >= slice.len()
-            || size_of_code == 0
-            || size_of_image == 0
-            || size_of_image > slice.len()
-        {
-            return None;
-        }
-
-        Some(Parser::<'a> {
-            image: &slice.get(0..=(size_of_image - 1)).unwrap(),
-            entrypoint: &slice.get(entrypoint..entrypoint).unwrap(),
-            code: &slice
-                .get(base_of_code..=(base_of_code + size_of_code - 1))
-                .unwrap(),
-        })
-    }
-
-    pub fn get_image(&self) -> &'a [u8] {
-        self.image
-    }
-
-    pub fn get_code_region(&self) -> &'a [u8] {
-        self.code
-    }
-
-    pub fn get_entrypoint(&self) -> *const u8 {
-        self.entrypoint as *const _ as *const u8
-    }
-}