6 files changed, 531 insertions, 0 deletions
diff --git a/tools/rust/src/chainload.rs b/tools/rust/src/chainload.rs
new file mode 100644
index 0000000..093158d
--- /dev/null
+++ b/tools/rust/src/chainload.rs
@@ -0,0 +1,101 @@
+// SPDX-License-Identifier: MIT
+#![deny(unsafe_op_in_unsafe_fn)]
+
+use crate::gpt;
+use crate::nvme;
+use crate::println;
+use alloc::vec::Vec;
+use core::ffi::c_void;
+use cstr_core::CStr;
+use cty::*;
+use fatfs::{FileSystem, FsOptions, Read, Seek, SeekFrom};
+use uuid::Uuid;
+
+#[derive(Debug)]
+pub enum Error {
+    FATError(fatfs::Error<nvme::Error>),
+    GPTError(gpt::Error<nvme::Error>),
+    BadArgs,
+    PartitionNotFound,
+    Unknown,
+}
+
+impl From<fatfs::Error<nvme::Error>> for Error {
+    fn from(err: fatfs::Error<nvme::Error>) -> Error {
+        Error::FATError(err)
+    }
+}
+
+impl From<gpt::Error<nvme::Error>> for Error {
+    fn from(err: gpt::Error<nvme::Error>) -> Error {
+        Error::GPTError(err)
+    }
+}
+
+fn load_image(spec: &str) -> Result<Vec<u8>, Error> {
+    println!("Chainloading {}", spec);
+
+    let mut args = spec.split(';');
+
+    let uuid = Uuid::parse_str(args.next().ok_or(Error::BadArgs)?).or(Err(Error::BadArgs))?;
+    let path = args.next().ok_or(Error::BadArgs)?;
+
+    let part = {
+        let storage = nvme::NVMEStorage::new(1, 0);
+        let mut pt = gpt::GPT::new(storage)?;
+
+        //println!("Partitions:");
+        //pt.dump();
+
+        println!("Searching for partition UUID: {}", uuid);
+        pt.find_by_partuuid(uuid)?.ok_or(Error::PartitionNotFound)?
+    };
+
+    let offset = part.get_starting_lba();
+
+    println!("Partition offset: {}", offset);
+
+    let storage = nvme::NVMEStorage::new(1, offset);
+    let opts = FsOptions::new().update_accessed_date(false);
+
+    let fs = FileSystem::new(storage, opts)?;
+    let mut file = fs.root_dir().open_file(path)?;
+
+    let size = file.seek(SeekFrom::End(0))? as usize;
+    file.seek(SeekFrom::Start(0))?;
+
+    println!("File size: {}", size);
+
+    let mut buf: Vec<u8> = vec![0; size];
+    let mut slice = &mut buf[..];
+    while !slice.is_empty() {
+        let read = file.read(slice)?;
+        slice = &mut slice[read..];
+    }
+    println!("File read successfully");
+
+    Ok(buf)
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn rust_load_image(
+    raw_spec: *const c_char,
+    image: *mut *mut c_void,
+    size: *mut size_t,
+) -> c_int {
+    let spec = unsafe { CStr::from_ptr(raw_spec).to_str().unwrap() };
+
+    match load_image(spec) {
+        Ok(buf) => {
+            unsafe {
+                *size = buf.len();
+                *image = buf.leak().as_mut_ptr() as *mut c_void;
+            }
+            0
+        }
+        Err(err) => {
+            println!("Chainload failed: {:?}", err);
+            -1
+        }
+    }
+}
diff --git a/tools/rust/src/dlmalloc.rs b/tools/rust/src/dlmalloc.rs
new file mode 100644
index 0000000..e60b2d1
--- /dev/null
+++ b/tools/rust/src/dlmalloc.rs
@@ -0,0 +1,79 @@
+// SPDX-License-Identifier: MIT
+
+use core::alloc::{GlobalAlloc, Layout};
+use core::ffi::c_void;
+use core::ptr;
+use cty::*;
+
+extern "C" {
+    pub fn malloc(size: size_t) -> *mut c_void;
+    pub fn realloc_in_place(p: *mut c_void, size: size_t) -> *mut c_void;
+    pub fn free(p: *mut c_void);
+    pub fn posix_memalign(p: *mut *mut c_void, alignment: size_t, size: size_t) -> c_int;
+}
+
+pub struct DLMalloc;
+
+unsafe impl GlobalAlloc for DLMalloc {
+    #[inline]
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let mut ptr = ptr::null_mut();
+        let ret = unsafe {
+            posix_memalign(
+                &mut ptr,
+                layout.align().max(core::mem::size_of::<usize>()),
+                layout.size(),
+            )
+        };
+        if ret == 0 {
+            ptr as *mut u8
+        } else {
+            ptr::null_mut()
+        }
+    }
+
+    #[inline]
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        // Unfortunately, calloc doesn't make any alignment guarantees, so the memory
+        // has to be manually zeroed-out.
+        let ptr = unsafe { self.alloc(layout) };
+        if !ptr.is_null() {
+            unsafe { ptr::write_bytes(ptr, 0, layout.size()) };
+        }
+        ptr
+    }
+
+    #[inline]
+    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
+        unsafe {
+            free(ptr as *mut c_void);
+        }
+    }
+
+    #[inline]
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        // Unfortunately, realloc doesn't make any alignment guarantees, so the memory
+        // has to be manually allocated as aligned memory if it cannot be resized
+        // in-place.
+        let mut new_ptr = unsafe { realloc_in_place(ptr as *mut c_void, new_size) as *mut u8 };
+
+        // return early if in-place resize succeeded
+        if !new_ptr.is_null() {
+            return new_ptr;
+        }
+
+        // allocate new aligned storage with correct layout
+        new_ptr =
+            unsafe { self.alloc(Layout::from_size_align_unchecked(new_size, layout.align())) };
+
+        // return early if allocation failed
+        if new_ptr.is_null() {
+            return ptr::null_mut();
+        }
+
+        // copy over the data and deallocate the old storage
+        unsafe { ptr::copy(ptr, new_ptr, layout.size().min(new_size)) };
+        unsafe { self.dealloc(ptr, layout) };
+        new_ptr
+    }
+}
diff --git a/tools/rust/src/gpt.rs b/tools/rust/src/gpt.rs
new file mode 100644
index 0000000..0c65d0b
--- /dev/null
+++ b/tools/rust/src/gpt.rs
@@ -0,0 +1,162 @@
+// SPDX-License-Identifier: MIT
+
+use crate::println;
+use core::convert::TryInto;
+use core::result::Result;
+use fatfs::{Read, Seek};
+use uuid::Uuid;
+
+const EFI_SIGNATURE: u64 = 0x5452415020494645;
+
+const SECTOR_SIZE: usize = 4096;
+
+#[derive(Debug)]
+pub enum Error<T> {
+    Io(T),
+    InvalidGPTHeader,
+}
+
+impl<T> From<T> for Error<T> {
+    fn from(err: T) -> Error<T> {
+        Error::Io(err)
+    }
+}
+
+struct TableHeader {
+    bytes: [u8; Self::SIZE],
+    my_lba: u64,
+}
+
+impl TableHeader {
+    const SIZE: usize = 0x5C;
+
+    fn read<R: Read + Seek>(rdr: &mut R, lba: u64) -> Result<Self, Error<R::Error>> {
+        let mut hdr = Self {
+            bytes: [0; Self::SIZE],
+            my_lba: lba,
+        };
+        let off = SECTOR_SIZE * (lba as usize);
+        rdr.seek(fatfs::SeekFrom::Start(off as u64))?;
+        rdr.read_exact(&mut hdr.bytes)?;
+        match hdr.is_valid() {
+            true => Ok(hdr),
+            false => Err(Error::InvalidGPTHeader),
+        }
+    }
+
+    fn get_signature(&self) -> u64 {
+        u64::from_le_bytes(self.bytes[0..8].try_into().unwrap())
+    }
+    fn get_my_lba(&self) -> u64 {
+        u64::from_le_bytes(self.bytes[24..32].try_into().unwrap())
+    }
+    fn get_partition_entry_lba(&self) -> u64 {
+        u64::from_le_bytes(self.bytes[72..80].try_into().unwrap())
+    }
+    fn get_partition_entry_count(&self) -> usize {
+        u32::from_le_bytes(self.bytes[80..84].try_into().unwrap()) as usize
+    }
+    fn get_partition_entry_size(&self) -> usize {
+        u32::from_le_bytes(self.bytes[84..88].try_into().unwrap()) as usize
+    }
+    fn is_valid(&self) -> bool {
+        self.get_signature() == EFI_SIGNATURE && self.get_my_lba() == self.my_lba
+    }
+}
+
+pub struct PartitionEntry {
+    bytes: [u8; Self::SIZE],
+}
+
+impl PartitionEntry {
+    const SIZE: usize = 0x80;
+
+    fn read<R: Read + Seek>(rdr: &mut R, off: usize) -> Result<Self, Error<R::Error>> {
+        let mut part = Self {
+            bytes: [0; Self::SIZE],
+        };
+        rdr.seek(fatfs::SeekFrom::Start(off as u64))?;
+        rdr.read_exact(&mut part.bytes)?;
+        Ok(part)
+    }
+
+    #[allow(dead_code)]
+    pub fn get_type_guid(&self) -> Uuid {
+        Uuid::from_bytes_le(self.bytes[0..16].try_into().unwrap())
+    }
+    pub fn get_partition_guid(&self) -> Uuid {
+        Uuid::from_bytes_le(self.bytes[16..32].try_into().unwrap())
+    }
+    pub fn get_starting_lba(&self) -> u64 {
+        u64::from_le_bytes(self.bytes[32..40].try_into().unwrap())
+    }
+    pub fn get_ending_lba(&self) -> u64 {
+        u64::from_le_bytes(self.bytes[40..48].try_into().unwrap())
+    }
+    pub fn get_attributes(&self) -> u64 {
+        u64::from_le_bytes(self.bytes[48..56].try_into().unwrap())
+    }
+    pub fn get_name(&self) -> &[u8] {
+        &self.bytes[56..72]
+    }
+}
+
+pub struct GPT<T: fatfs::ReadWriteSeek> {
+    disk: T,
+    hdr: TableHeader,
+}
+
+impl<IO: fatfs::ReadWriteSeek> GPT<IO> {
+    pub fn new<T: fatfs::IntoStorage<IO>>(storage: T) -> Result<Self, Error<IO::Error>> {
+        let mut disk = storage.into_storage();
+
+        let hdr = TableHeader::read(&mut disk, 1)?;
+
+        let gpt = Self { disk, hdr };
+        Ok(gpt)
+    }
+
+    pub fn count(&self) -> usize {
+        self.hdr.get_partition_entry_count()
+    }
+
+    pub fn index(&mut self, index: usize) -> Result<PartitionEntry, Error<IO::Error>> {
+        let off = (self.hdr.get_partition_entry_lba() as usize * SECTOR_SIZE)
+            + index * self.hdr.get_partition_entry_size();
+        PartitionEntry::read(&mut self.disk, off)
+    }
+
+    pub fn find_by_partuuid(
+        &mut self,
+        uuid: Uuid,
+    ) -> Result<Option<PartitionEntry>, Error<IO::Error>> {
+        for i in 0..self.count() {
+            let part = self.index(i)?;
+            if part.get_type_guid().is_nil() {
+                continue;
+            }
+            if part.get_partition_guid() == uuid {
+                return Ok(Some(part));
+            }
+        }
+        Ok(None)
+    }
+
+    pub fn dump(&mut self) {
+        for i in 0..self.count() {
+            let part = self.index(i).unwrap();
+            let guid = part.get_type_guid();
+            if guid.is_nil() {
+                continue;
+            }
+            println!(
+                "{}: {}..{} {:x} {:x}",
+                i,
+                part.get_starting_lba(),
+                part.get_ending_lba(),
+                guid,
+                part.get_partition_guid()
+            );
+        }
+    }
+}
diff --git a/tools/rust/src/lib.rs b/tools/rust/src/lib.rs
new file mode 100644
index 0000000..a578b91
--- /dev/null
+++ b/tools/rust/src/lib.rs
@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: MIT
+#![no_std]
+#![deny(unsafe_op_in_unsafe_fn)]
+#![feature(alloc_error_handler)]
+#![feature(mixed_integer_ops)]
+#![feature(new_uninit)]
+
+#[macro_use]
+extern crate alloc;
+
+pub mod chainload;
+pub mod dlmalloc;
+pub mod gpt;
+pub mod nvme;
+pub mod print;
+
+use crate::dlmalloc::DLMalloc;
+
+#[global_allocator]
+static GLOBAL: DLMalloc = dlmalloc::DLMalloc;
+
+extern "C" {
+    fn flush_and_reboot();
+}
+
+#[panic_handler]
+fn panic(info: &::core::panic::PanicInfo) -> ! {
+    println!("{}", info);
+    unsafe { flush_and_reboot() };
+    loop {}
+}
+
+#[alloc_error_handler]
+fn alloc_error(layout: core::alloc::Layout) -> ! {
+    panic!("memory allocation of {} bytes failed", layout.size())
+}
diff --git a/tools/rust/src/nvme.rs b/tools/rust/src/nvme.rs
new file mode 100644
index 0000000..fdcec0e
--- /dev/null
+++ b/tools/rust/src/nvme.rs
@@ -0,0 +1,93 @@
+// SPDX-License-Identifier: MIT
+use crate::println;
+use alloc::boxed::Box;
+use core::cmp::min;
+use core::ffi::c_void;
+use fatfs::SeekFrom;
+
+extern "C" {
+    fn nvme_read(nsid: u32, lba: u64, buffer: *mut c_void) -> bool;
+}
+
+const SECTOR_SIZE: usize = 4096;
+
+pub type Error = ();
+
+#[repr(C, align(4096))]
+struct SectorBuffer([u8; SECTOR_SIZE]);
+
+fn alloc_sector_buf() -> Box<SectorBuffer> {
+    let p: Box<SectorBuffer> = unsafe { Box::new_zeroed().assume_init() };
+    debug_assert_eq!(0, p.0.as_ptr().align_offset(4096));
+    p
+}
+
+pub struct NVMEStorage {
+    nsid: u32,
+    offset: u64,
+    lba: Option<u64>,
+    buf: Box<SectorBuffer>,
+    pos: u64,
+}
+
+impl NVMEStorage {
+    pub fn new(nsid: u32, offset: u64) -> NVMEStorage {
+        NVMEStorage {
+            nsid: nsid,
+            offset: offset,
+            lba: None,
+            buf: alloc_sector_buf(),
+            pos: 0,
+        }
+    }
+}
+
+impl fatfs::IoBase for NVMEStorage {
+    type Error = Error;
+}
+
+impl fatfs::Read for NVMEStorage {
+    fn read(&mut self, mut buf: &mut [u8]) -> Result<usize, Self::Error> {
+        let mut read = 0;
+
+        while !buf.is_empty() {
+            let lba = self.pos / SECTOR_SIZE as u64;
+            let off = self.pos as usize % SECTOR_SIZE;
+
+            if Some(lba) != self.lba {
+                self.lba = Some(lba);
+                let lba = lba + self.offset;
+                if !unsafe { nvme_read(self.nsid, lba, self.buf.0.as_mut_ptr() as *mut c_void) } {
+                    println!("nvme_read({}, {}) failed", self.nsid, lba);
+                    return Err(());
+                }
+            }
+            let copy_len = min(SECTOR_SIZE - off, buf.len());
+            buf[..copy_len].copy_from_slice(&self.buf.0[off..off + copy_len]);
+            buf = &mut buf[copy_len..];
+            read += copy_len;
+            self.pos += copy_len as u64;
+        }
+        Ok(read)
+    }
+}
+
+impl fatfs::Write for NVMEStorage {
+    fn write(&mut self, _buf: &[u8]) -> Result<usize, Self::Error> {
+        Err(())
+    }
+    fn flush(&mut self) -> Result<(), Self::Error> {
+        Err(())
+    }
+}
+
+impl fatfs::Seek for NVMEStorage {
+    fn seek(&mut self, from: SeekFrom) -> Result<u64, Self::Error> {
+        self.pos = match from {
+            SeekFrom::Start(n) => n,
+            SeekFrom::End(_n) => panic!("SeekFrom::End not supported"),
+            SeekFrom::Current(n) => self.pos.checked_add_signed(n).ok_or(())?,
+        };
+        Ok(self.pos)
+    }
+}
diff --git a/tools/rust/src/print.rs b/tools/rust/src/print.rs
new file mode 100644
index 0000000..b9f0ead
--- /dev/null
+++ b/tools/rust/src/print.rs
@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: MIT
+use core::ffi::c_void;
+
+extern "C" {
+    fn iodev_console_write(buf: *const c_void, len: u64);
+}
+
+pub struct IODevConsoleWriter;
+
+impl core::fmt::Write for IODevConsoleWriter {
+    #[inline]
+    fn write_str(&mut self, msg: &str) -> core::fmt::Result {
+        write(msg)
+    }
+}
+
+impl IODevConsoleWriter {
+    #[inline]
+    pub fn write_fmt(args: core::fmt::Arguments) -> core::fmt::Result {
+        core::fmt::Write::write_fmt(&mut Self, args)
+    }
+
+    #[inline]
+    pub fn write_str(msg: &str) -> core::fmt::Result {
+        write(msg)
+    }
+
+    #[inline]
+    pub fn write_nl() -> core::fmt::Result {
+        write("\n")
+    }
+}
+
+#[inline]
+fn write(msg: &str) -> core::fmt::Result {
+    unsafe { iodev_console_write(msg.as_ptr() as _, msg.len() as u64) };
+    Ok(())
+}
+
+#[macro_export]
+macro_rules! println {
+    () => { $crate::println!("") };
+    ($($arg:tt)*) => {
+        #[allow(unused_must_use)]
+        {
+            $crate::print::IODevConsoleWriter::write_fmt(format_args!($($arg)*));
+            $crate::print::IODevConsoleWriter::write_nl();
+        }
+    };
+}
+
+#[macro_export]
+macro_rules! print {
+    ($($arg:tt)*) => {
+        #[allow(unused_must_use)]
+        {
+            $crate::print::IODevConsoleWriter::write_fmt(format_args!($($arg)*));
+        }
+    };
+}