More serialization implementation

src/abi.rs

@@ -134,7 +134,7 @@ pub enum ModuleInfoType {
 }
 
 /// Value in the module info TLV
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub enum ModuleInfoValue {
     StaticString(&'static str),
     String(String),
@@ -177,6 +177,18 @@ impl Serialize for ModuleInfoValue {
     }
 }
 
+impl core::fmt::Debug for ModuleInfoValue {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            Self::StaticString(arg0) => f.debug_tuple("StaticString").field(arg0).finish(),
+            Self::String(arg0) => f.debug_tuple("String").field(arg0).finish(),
+            Self::Pointer(arg0) => f.debug_tuple("Pointer").field(arg0).finish(),
+            Self::Size(arg0) => f.debug_tuple("Size").field(arg0).finish(),
+            Self::Buffer(_) => write!(f, "Buffer(elided)"),
+        }
+    }
+}
+
 /// Key and value in the module info TLV
 /// TODO: figure out how to make this properly typesafe
 #[derive(Debug, Clone)]

src/main.rs

@@ -5,7 +5,8 @@ mod abi;
 mod object;
 mod staging;
 
-use log::info;
+use abi::{ModuleInfoItem, Serialize};
+use log::{info};
 use object::{Kernel, Module};
 use uefi::{fs::FileSystem, prelude::*};
 
@@ -29,13 +30,35 @@ fn main(image_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
     info!("Read kernel");
 
     let kernel = Kernel::from_elf_bytes(&kernel_data);
+    let kernel_metadata = kernel.metadata();
+    info!("Kernel metadata is {:x?}", kernel_metadata);
 
-    info!(
+    let metadata_size: usize = kernel_metadata
-        "Kernel metadata is {:#?}",
+        .iter()
-        kernel.metadata(core::ptr::null::<u8>())
+        .map(|item| item.size().next_multiple_of(ModuleInfoItem::alignment()))
-    );
+        .sum();
+
+    // Allocate an extra 16KiB for items we have to add to metadata
+    let alloc_size = kernel.size() + Kernel::alignment() + metadata_size + 16 * 1024;
+
+    let staging = staging::allocate_staging(system_table.boot_services(), alloc_size);
+
+    {
+        let mut aligned = align_slice(staging, Kernel::alignment());
+        kernel.serialize(aligned);
+        for item in kernel_metadata.iter() {
+            aligned = align_slice(aligned, ModuleInfoItem::alignment());
+            item.serialize(aligned);
+        }
+    }
 
     loop {
         system_table.boot_services().stall(10_000_000);
     }
 }
+
+fn align_slice(unaligned: &mut [u8], alignment: usize) -> &mut [u8] {
+    let base_addr = &unaligned[0] as *const u8 as usize;
+    let offset = base_addr.next_multiple_of(alignment) - base_addr;
+    &mut unaligned[offset..]
+}

src/object.rs

@@ -2,15 +2,17 @@ extern crate alloc;
 
 use crate::abi::{ModuleInfoItem, ModuleInfoType, ModuleInfoValue, Serialize};
 use alloc::vec::Vec;
-use goblin::elf::Elf;
+use goblin::elf::{program_header, section_header, Elf, SectionHeader};
-use log::info;
+use log::{debug, info};
 
 #[cfg(target_arch = "x86_64")]
 const NATIVE_ELF_MACHINE: u16 = goblin::elf::header::EM_X86_64;
 
+const USIZE_SIZE: usize = core::mem::size_of::<usize>();
+
 pub trait Module {
     /// Associated metadata for this module, including name and type
-    fn metadata(&self, load_address: *const u8) -> Vec<ModuleInfoItem>;
+    fn metadata(&self) -> Vec<ModuleInfoItem>;
 }
 
 pub struct Kernel<'a> {
@@ -18,8 +20,12 @@ pub struct Kernel<'a> {
     elf: &'a [u8],
     /// Parsed ELF
     parsed: Elf<'a>,
-    /// Raw ELF header bytes
+    /// Highest virtual address from program headers
-    elf_header: &'a [u8],
+    max_loadaddr: usize,
+    /// Lowest virtual address from program headers
+    min_loadaddr: usize,
+    /// .symtab and .symstr section headers
+    symbol_headers: Option<(SectionHeader, SectionHeader)>,
 }
 
 impl<'a> Kernel<'a> {
@@ -29,15 +35,105 @@ impl<'a> Kernel<'a> {
             panic!("Kernel has wrong ELF machine");
         }
 
-        let header_size = if parsed.is_64 { 0x40 } else { 0x34 };
+        let max_loadaddr = parsed
-        let elf_header = &elf[..header_size];
+            .program_headers
+            .iter()
+            .filter(|phdr| phdr.p_type == program_header::PT_LOAD)
+            .map(|phdr| phdr.p_vaddr + phdr.p_memsz)
+            .max()
+            .expect("Kernel has no program headers") as usize;
+
+        let min_loadaddr = parsed
+            .program_headers
+            .iter()
+            .filter(|phdr| phdr.p_type == program_header::PT_LOAD)
+            .map(|phdr| phdr.p_vaddr)
+            .min()
+            .unwrap() as usize;
+
+        let symbol_headers = Self::find_symbol_headers(&parsed);
 
         Self {
             elf,
             parsed,
-            elf_header,
+            max_loadaddr,
+            min_loadaddr,
+            symbol_headers,
         }
     }
+
+    pub fn entry_addr(&self, load_addr: usize) -> usize {
+        self.parsed.header.e_entry as usize - self.min_loadaddr + load_addr
+    }
+
+    fn find_symbol_headers(parsed: &Elf) -> Option<(SectionHeader, SectionHeader)> {
+        let symtab = parsed
+            .section_headers
+            .iter()
+            .find(|shdr| shdr.sh_type == section_header::SHT_SYMTAB)?;
+
+        let symstr = parsed.section_headers.get(symtab.sh_link as usize)?;
+
+        Some((symtab.clone(), symstr.clone()))
+    }
+
+    fn serialize_load_segments(&self, out: &mut [u8]) {
+        for phdr in self
+            .parsed
+            .program_headers
+            .iter()
+            .filter(|phdr| phdr.p_type == program_header::PT_LOAD)
+        {
+            let out_offset = phdr.p_vaddr as usize - self.min_loadaddr;
+            let in_offset = phdr.p_offset as usize;
+            let copy_size = phdr.p_filesz as usize;
+            out[out_offset..out_offset + copy_size]
+                .copy_from_slice(&self.elf[in_offset..in_offset + copy_size]);
+
+            debug!(
+                "kernel: {:x}@{:x} -> {:x}",
+                copy_size,
+                in_offset,
+                (&out[out_offset] as *const u8) as usize
+            );
+
+            // No need to zero bss, we zeroed during allocation
+        }
+    }
+
+    fn size_load_segments(&self) -> usize {
+        (self.max_loadaddr - self.min_loadaddr).next_multiple_of(USIZE_SIZE)
+    }
+
+    fn size_symbols(&self) -> usize {
+        let Some((ref symtab, ref symstr)) = self.symbol_headers else {
+            return 0;
+        };
+
+        self.size_shdr(symtab) + self.size_shdr(symstr)
+    }
+
+    fn size_shdr(&self, shdr: &SectionHeader) -> usize {
+        USIZE_SIZE + (shdr.sh_size as usize).next_multiple_of(USIZE_SIZE)
+    }
+
+    fn serialize_shdr(&self, shdr: &SectionHeader, out: &mut [u8]) {
+        let elf_offset = shdr.sh_offset as usize;
+        let copy_size = shdr.sh_size as usize;
+        out[..USIZE_SIZE].copy_from_slice(&usize::to_ne_bytes(shdr.sh_size as usize));
+        out[USIZE_SIZE..USIZE_SIZE + copy_size]
+            .copy_from_slice(&self.elf[elf_offset..elf_offset + copy_size])
+    }
+
+    fn serialize_symbols(&self, out: &mut [u8]) {
+        let Some((ref symtab, ref symstr)) = self.symbol_headers else {
+            info!("Kernel has no symbols, skipping");
+            return;
+        };
+
+        self.serialize_shdr(symtab, out);
+        self.serialize_shdr(symstr, &mut out[self.size_shdr(symtab)..])
+    }
 }
 
 impl<'a> Serialize for Kernel<'a> {
@@ -47,17 +143,24 @@ impl<'a> Serialize for Kernel<'a> {
     }
 
     fn size(&self) -> usize {
-        todo!()
+        self.size_load_segments() + self.size_symbols()
     }
 
     fn serialize(&self, out: &mut [u8]) {
-        todo!()
+        // Each section is rounded to sizeof(long)
+        let end_load = self.size_load_segments();
+        self.serialize_load_segments(&mut out[..end_load]);
+
+        let end_symbols = end_load + self.size_symbols();
+        self.serialize_symbols(&mut out[end_load..end_symbols]);
     }
 }
 
 impl<'a> Module for Kernel<'a> {
-    fn metadata(&self, load_address: *const u8) -> Vec<ModuleInfoItem> {
+    fn metadata(&self) -> Vec<ModuleInfoItem> {
-        alloc::vec![
+        // FreeBSD would also put constructors here but I haven't managed
+        // to find any kernels or modules with constructors
+        let mut items = alloc::vec![
             ModuleInfoItem {
                 tag: ModuleInfoType::Name,
                 value: ModuleInfoValue::StaticString("/boot/kernel/kernel")
@@ -67,9 +170,47 @@ impl<'a> Module for Kernel<'a> {
                 value: ModuleInfoValue::StaticString("elf kernel")
             },
             ModuleInfoItem {
-                tag: ModuleInfoType::ElfHeader,
+                tag: ModuleInfoType::LoadAddress,
-                value: ModuleInfoValue::Buffer(self.elf_header.to_vec())
+                value: ModuleInfoValue::Size(self.min_loadaddr)
             },
-        ]
+            ModuleInfoItem {
+                tag: ModuleInfoType::Size,
+                value: ModuleInfoValue::Size(self.size())
+            },
+            ModuleInfoItem {
+                tag: ModuleInfoType::ElfHeader,
+                value: ModuleInfoValue::Buffer({
+                    let header_size = if self.parsed.is_64 { 0x40 } else { 0x34 };
+                    self.elf[..header_size].to_vec()
+                })
+            },
+            ModuleInfoItem {
+                tag: ModuleInfoType::SectionHeader,
+                value: ModuleInfoValue::Buffer({
+                    let offset = self.parsed.header.e_shoff as usize;
+                    let header_size = self.parsed.header.e_shentsize as usize;
+                    // FreeBSD loader doesn't handle e_snum == 0x00ff,
+                    // don't bother doing it here for compatibility
+                    let num_headers = self.parsed.header.e_shnum as usize;
+                    self.elf[offset..offset + header_size * num_headers].to_vec()
+                })
+            },
+        ];
+
+        if self.symbol_headers.is_some() {
+            let start_symbols = self.min_loadaddr + self.size_load_segments();
+            let end_symbols = start_symbols + self.size_symbols();
+
+            items.push(ModuleInfoItem {
+                tag: ModuleInfoType::StartSymbols,
+                value: ModuleInfoValue::Size(start_symbols),
+            });
+            items.push(ModuleInfoItem {
+                tag: ModuleInfoType::EndSymbols,
+                value: ModuleInfoValue::Size(end_symbols),
+            });
+        }
+
+        items
     }
 }

src/staging.rs

@@ -1,13 +1,7 @@
 use log::info;
-use uefi::table::boot::{AllocateType, BootServices, MemoryType};
+use uefi::table::boot::{AllocateType, BootServices, MemoryType, PAGE_SIZE};
 
-/// Size of the staging buffer for preloaded files
+pub fn allocate_staging(boot_services: &BootServices, size: usize) -> &'static mut [u8] {
-const STAGING_BUF_SIZE: usize = 64 * 1024 * 1024;
-const STAGING_BUF_NUM_PAGES: usize = STAGING_BUF_SIZE / uefi::table::boot::PAGE_SIZE;
-
-const ADDR_ROUND: u64 = 2 * 1024 * 1024;
-
-pub fn allocate_staging(boot_services: &BootServices) -> &'static mut [u8] {
     let allocation_type = if cfg!(target_arch = "x86_64") {
         // Must be under 4GiB because we need to pass a 32-bit pointer
         AllocateType::MaxAddress(4 * 1024 * 1024 * 1024)
@@ -15,22 +9,22 @@ pub fn allocate_staging(boot_services: &BootServices) -> &'static mut [u8] {
         AllocateType::AnyPages
     };
 
+    let num_pages = size.div_ceil(PAGE_SIZE);
+    let total_size = num_pages * PAGE_SIZE;
+
     let base = boot_services
-        .allocate_pages(
+        .allocate_pages(allocation_type, MemoryType::LOADER_CODE, num_pages)
-            allocation_type,
-            MemoryType::LOADER_CODE,
-            STAGING_BUF_NUM_PAGES,
-        )
         .expect("Unable to allocate staging area");
 
-    // Some architectures want a kernel on a 2MiB boundary, do it everywhere
-    let rounded = ((base + ADDR_ROUND - 1) / ADDR_ROUND) * ADDR_ROUND;
-
-    let available_size = STAGING_BUF_SIZE - ((rounded - base) as usize);
     info!(
         "Allocated {:#018x} byte staging buffer at {:#018x}",
-        available_size, rounded
+        num_pages * PAGE_SIZE,
+        base
     );
 
-    unsafe { core::slice::from_raw_parts_mut(rounded as *mut u8, available_size) }
+    unsafe {
+        let ptr = base as *mut u8;
+        ptr.write_bytes(0, total_size);
+        core::slice::from_raw_parts_mut(ptr, total_size)
+    }
 }