From b880c11ac7fa95ae7e0244be7fc87e83ebc715dd Mon Sep 17 00:00:00 2001
From: Drew Galbraith <dgalbraith33@gmail.com>
Date: Sun, 18 Aug 2024 11:36:25 -0700
Subject: [PATCH] [Yellowstone] add ability to read GPT.

---
 rust/sys/yellowstone/src/gpt.rs    | 138 +++++++++++++++++++++++++++++
 rust/sys/yellowstone/src/main.rs   |   1 +
 rust/sys/yellowstone/src/server.rs |  15 +++-
 3 files changed, 153 insertions(+), 1 deletion(-)
 create mode 100644 rust/sys/yellowstone/src/gpt.rs

diff --git a/rust/sys/yellowstone/src/gpt.rs b/rust/sys/yellowstone/src/gpt.rs
new file mode 100644
index 0000000..950baf9
--- /dev/null
+++ b/rust/sys/yellowstone/src/gpt.rs
@@ -0,0 +1,138 @@
+use denali::DenaliClient;
+use mammoth::{cap::Capability, zion::ZError};
+
+const MBR_SIG: u16 = 0xAA55;
+
+const GPT_OS_TYPE: u8 = 0xEE;
+
+#[repr(C, packed)]
+struct MbrPartition {
+    boot_indicator: u8,
+    starting_chs: [u8; 3],
+    os_type: u8,
+    ending_chs: [u8; 3],
+    starting_lba: u32,
+    ending_lba: u32,
+}
+
+#[repr(C, packed)]
+struct PartitionHeader {
+    signature: u64,
+    revision: u32,
+    header_size: u32,
+    crc_32: u32,
+    reserved: u32,
+    lba_self: u64,
+    lba_mirror: u64,
+    lba_min: u64,
+    lba_max: u64,
+    guid_low: u64,
+    guid_high: u64,
+    lba_partition_entries: u64,
+    num_partitions: u32,
+    partition_entry_size: u32,
+    partition_entry_crc32: u32,
+}
+
+#[repr(C, packed)]
+struct PartitionEntry {
+    type_guid_low: u64,
+    type_guid_high: u64,
+    part_guid_low: u64,
+    part_guid_high: u64,
+    lba_start: u64,
+    lba_end: u64,
+    attrs: u64,
+    partition_name: [u8; 72],
+}
+
+pub fn read_gpt(mut denali: DenaliClient) -> Result<u64, ZError> {
+    let resp = denali.read(&denali::ReadRequest {
+        device_id: 0,
+        block: denali::DiskBlock { lba: 0, size: 2 },
+    })?;
+
+    let first_lbas = mammoth::mem::MemoryRegion::from_cap(Capability::take(resp.memory))?;
+
+    let maybe_mbr_sig = u16::from_le_bytes(first_lbas.slice()[0x1FE..0x200].try_into().unwrap());
+
+    if maybe_mbr_sig != MBR_SIG {
+        mammoth::debug!("MBR Sig {:#x} does not match {:#x}", maybe_mbr_sig, MBR_SIG);
+        return Err(ZError::FAILED_PRECONDITION);
+    }
+
+    let mbr = unsafe {
+        first_lbas.slice::<u8>()[0x1BE..]
+            .as_ptr()
+            .cast::<MbrPartition>()
+            .as_ref()
+            .unwrap()
+    };
+
+    if mbr.os_type != GPT_OS_TYPE {
+        let os_type = mbr.os_type;
+        mammoth::debug!(
+            "MBR OS Type: {:#x} does not match GPT ({:#x})",
+            os_type,
+            GPT_OS_TYPE
+        );
+    }
+
+    let gpt_header = unsafe {
+        first_lbas.slice::<u8>()[512..]
+            .as_ptr()
+            .cast::<PartitionHeader>()
+            .as_ref()
+            .unwrap()
+    };
+
+    let num_partitions = gpt_header.num_partitions;
+    let partition_entry_size = gpt_header.partition_entry_size;
+    let lba_partition_entries = gpt_header.lba_partition_entries;
+    let num_blocks = (num_partitions * partition_entry_size).div_ceil(512) as u64;
+
+    mammoth::debug!(
+        "Reading partition table from LBA {}, {} entries of size {} equals {} blocks.",
+        lba_partition_entries,
+        num_partitions,
+        partition_entry_size,
+        num_blocks
+    );
+
+    let resp = denali.read(&denali::ReadRequest {
+        device_id: 0,
+        block: denali::DiskBlock {
+            lba: lba_partition_entries,
+            size: num_blocks,
+        },
+    })?;
+
+    let partition_table = mammoth::mem::MemoryRegion::from_cap(Capability::take(resp.memory))?;
+
+    for ind in 0..num_partitions {
+        let offset = (ind * partition_entry_size) as usize;
+        let limit = offset + (partition_entry_size as usize);
+
+        let partition_entry = unsafe {
+            partition_table.slice::<u8>()[offset..limit]
+                .as_ptr()
+                .cast::<PartitionEntry>()
+                .as_ref()
+                .unwrap()
+        };
+
+        if partition_entry.type_guid_low == 0 && partition_entry.type_guid_high == 0 {
+            continue;
+        }
+
+        const LFS_DATA_LOW: u64 = 0x477284830fc63daf;
+        const LFS_DATA_HIGH: u64 = 0xe47d47d8693d798e;
+        if partition_entry.type_guid_low == LFS_DATA_LOW
+            && partition_entry.type_guid_high == LFS_DATA_HIGH
+        {
+            return Ok(partition_entry.lba_start);
+        }
+    }
+
+    Err(ZError::NOT_FOUND)
+}
diff --git a/rust/sys/yellowstone/src/main.rs b/rust/sys/yellowstone/src/main.rs
index 2a62a1a..d3c84d1 100644
--- a/rust/sys/yellowstone/src/main.rs
+++ b/rust/sys/yellowstone/src/main.rs
@@ -13,6 +13,7 @@ use mammoth::{
 use yellowstone_yunq::YellowstoneServer;
 use yunq::server::YunqServer;
 
+mod gpt;
 mod pci;
 mod server;
 
diff --git a/rust/sys/yellowstone/src/server.rs b/rust/sys/yellowstone/src/server.rs
index 8af4673..63a9d23 100644
--- a/rust/sys/yellowstone/src/server.rs
+++ b/rust/sys/yellowstone/src/server.rs
@@ -88,6 +88,19 @@ impl YellowstoneServerHandler for YellowstoneServerImpl {
     }
 
     fn get_denali(&mut self) -> Result<DenaliInfo, ZError> {
-        todo!()
+        match self.service_map.get("denali") {
+            Some(ep_cap) => crate::gpt::read_gpt(denali::DenaliClient::new(
+                ep_cap.duplicate(Capability::PERMS_ALL).unwrap(),
+            ))
+            .map(|lba| DenaliInfo {
+                denali_endpoint: ep_cap.duplicate(Capability::PERMS_ALL).unwrap().release(),
+                device_id: 0,
+                lba_offset: lba,
+            }),
+            None => {
+                mammoth::debug!("Denali not yet registered");
+                Err(ZError::FAILED_PRECONDITION)
+            }
+        }
     }
 }