From fd3eae75edc261f5eb6dcc8c0e33e55bd695b428 Mon Sep 17 00:00:00 2001
From: pjht <pterpstra@pterpstra.com>
Date: Mon, 12 Aug 2024 20:48:01 -0500
Subject: [PATCH] data_in now accepts buffer and size over 4MiB, and frees
 allocated prd/command buffers

---
 Cargo.toml  |  5 ++++
 src/main.rs | 72 ++++++++++++++++++++++++++++++++++++++---------------
 2 files changed, 57 insertions(+), 20 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 9d6c57a..6fd0f77 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,3 +21,8 @@ x86_64 = "0.15.1"
 
 [patch.crates-io]
 tock-registers = {path = "./tock-registers"}
+
+[profile.release]
+strip = true
+lto = true
+opt-level = "s"
diff --git a/src/main.rs b/src/main.rs
index bbd975d..891d646 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -65,9 +65,11 @@ impl AhciPort {
     fn issue_command_data_in(
         &self,
         fis: &RegH2DFis,
-        len: usize,
-    ) -> Result<Vec<u8>, CommandIssueError> {
-        if len > 0x40_0000 {
+        buf: &mut [u8]
+    ) -> Result<(), CommandIssueError> {
+        let len = buf.len();
+
+        if len > 0x40_0000 * 0xFFFF {
             return Err(CommandIssueError::DataTooBig);
         }
 
@@ -77,7 +79,23 @@ impl AhciPort {
             .map_free_cont_phys(len.div_ceil(4096))
             .unwrap();
 
-        let prdt = [Prd::new(data_buf_phys, len as u32, false).unwrap()];
+        let num_prds = len.div_ceil(0x40_0000);
+
+        let prdt = (0..num_prds)
+            .map(|i| {
+                let size = if i == num_prds - 1 {
+                    len - (i * 0x40_0000)
+                } else {
+                    0x40_0000
+                };
+                println!(
+                    "Prd::new({:#x}, {:#x}, false)",
+                    data_buf_phys + (i * 0x40_0000) as u64,
+                    size as u32
+                );
+                Prd::new(data_buf_phys + (i * 0x40_0000) as u64, size as u32, false).unwrap()
+            })
+            .collect_vec();
 
         unsafe {
             self.issue_command_prdt(fis, &prdt)?;
@@ -89,12 +107,16 @@ impl AhciPort {
                 len,
             ))
         };
-        let mut data = vec![0; len];
 
-        data_vol.copy_into_slice(data.as_mut_slice());
+        data_vol.copy_into_slice(buf);
 
-        // TODO: Free data buffer
-        Ok(data)
+        std::os::mikros::address_space::ACTIVE_SPACE
+            .lock()
+            .unwrap()
+            .unmap(data_buf, len.div_ceil(4096))
+            .unwrap();
+
+        Ok(())
     }
 
     unsafe fn issue_command_prdt(
@@ -164,7 +186,12 @@ impl AhciPort {
             return Err(CommandIssueError::CommandFailed);
         }
 
-        // TODO: Free command table
+        std::os::mikros::address_space::ACTIVE_SPACE
+            .lock()
+            .unwrap()
+            .unmap(buf, tbl_size.div_ceil(4096))
+            .unwrap();
+
         Ok(())
     }
 
@@ -371,10 +398,10 @@ fn main() {
             .build()
             .unwrap();
 
-        if let Ok(data) = port.issue_command_data_in(&cmd_fis, 512) {
-            let identify_data = IdentifyData::read(&mut Cursor::new(data.as_slice())).unwrap();
+        let mut ident_buf = [0; 512];
 
-            dbg!(&identify_data);
+        if port.issue_command_data_in(&cmd_fis, &mut ident_buf).is_ok() {
+            let identify_data = IdentifyData::read(&mut Cursor::new(ident_buf.as_slice())).unwrap();
 
             let sect_sz_info_valid = identify_data.phys_log_sect_sz & 0xC000 == 0x4000;
             let sect_size = if sect_sz_info_valid {
@@ -427,16 +454,14 @@ fn main() {
                 .build()
                 .unwrap();
 
-            let Ok(data) = port.issue_command_data_in(&cmd_fis, 512) else {
+            if port.issue_command_data_in(&cmd_fis, &mut ident_buf).is_err() {
                 syslog_client
                     .send_text_message("ahci", format!("Port {} failed to identify", port.phys_no))
                     .unwrap();
                 continue;
             };
 
-            let identify_data = IdentifyDataATAPI::read(&mut Cursor::new(data.as_slice())).unwrap();
-
-            dbg!(&identify_data);
+            let identify_data = IdentifyDataATAPI::read(&mut Cursor::new(ident_buf.as_slice())).unwrap();
 
             syslog_client
                 .send_text_message(
@@ -466,7 +491,10 @@ fn main() {
         .build()
         .unwrap();
 
-    let mbr = avail_ports[0].issue_command_data_in(&cmd_fis, 512).unwrap();
+    let mut mbr = [0; 512];
+
+    avail_ports[0].issue_command_data_in(&cmd_fis, &mut mbr).unwrap();
+
 
     let mbr_entries = (0..4)
         .map(|i| {
@@ -517,7 +545,9 @@ fn main() {
             .build()
             .unwrap();
 
-        let gpt_header = avail_ports[0].issue_command_data_in(&cmd_fis, 512).unwrap();
+        let mut gpt_header = [0; 512]; 
+
+        avail_ports[0].issue_command_data_in(&cmd_fis, &mut gpt_header).unwrap();
 
         if &gpt_header[0..8] != b"EFI PART" {
             println!("Invalid GPT signature!");
@@ -570,8 +600,10 @@ fn main() {
             .build()
             .unwrap();
 
-        let gpt_part_table = avail_ports[0]
-            .issue_command_data_in(&cmd_fis, part_table_num_lbas * 512)
+        let mut gpt_part_table = vec![0; part_table_num_lbas * 512];
+
+        avail_ports[0]
+            .issue_command_data_in(&cmd_fis, gpt_part_table.as_mut_slice())
             .unwrap();
 
         for i in 0..num_parts {