From ac68c7e932397a8022f05f03097a6419d90383dc Mon Sep 17 00:00:00 2001
From: pjht <pterpstra@pterpstra.com>
Date: Fri, 28 Feb 2025 19:30:22 -0600
Subject: [PATCH] Read and store identify data on port creation

---
 src/hba.rs  | 11 +++++++--
 src/port.rs | 65 ++++++++++++++++++++++++++++++++++++-----------------
 2 files changed, 53 insertions(+), 23 deletions(-)

diff --git a/src/hba.rs b/src/hba.rs
index 834c82c..4db5ace 100644
--- a/src/hba.rs
+++ b/src/hba.rs
@@ -2,7 +2,7 @@ use std::os::mikros::address_space::ACTIVE_SPACE;
 
 use crate::{
     ahci_structs::{CommandHeader, FisBuf, GenHC, PortRegs, CAP, GHC},
-    port::AhciPort,
+    port::{AhciPort, CommandIssueError},
 };
 
 pub struct Hba {
@@ -17,6 +17,13 @@ unsafe impl Send for Hba {}
 #[derive(Clone, Copy, Debug)]
 pub enum HBAInitError {
     No64BitDma,
+    CommandIssueError(CommandIssueError),
+}
+
+impl From<CommandIssueError> for HBAInitError {
+    fn from(v: CommandIssueError) -> Self {
+        Self::CommandIssueError(v)
+    }
 }
 
 impl Hba {
@@ -106,7 +113,7 @@ impl Hba {
                 cmd_list,
                 cmd_list_phys,
                 max_cmd_slots,
-            ));
+            )?);
         }
         Ok(hba)
     }
diff --git a/src/port.rs b/src/port.rs
index dca1bce..d654f4d 100644
--- a/src/port.rs
+++ b/src/port.rs
@@ -11,6 +11,7 @@ use crate::{
     ata_command::{AtaCommandDataIn, IdentifyCommand, IdentifyPacketCommand},
     ata_dev::{AtaDevice, AtaNewError},
     atapi_dev::AtapiDevice,
+    identify::{IdentifyData, IdentifyPacketData},
 };
 
 #[allow(clippy::module_name_repetitions)]
@@ -22,6 +23,13 @@ pub struct AhciPort {
     #[allow(dead_code)]
     cmd_list_len: usize,
     has_device: bool,
+    identify_data: Option<AhciIdentifyData>,
+}
+
+#[derive(Clone)]
+enum AhciIdentifyData {
+    Normal(IdentifyData),
+    Packet(IdentifyPacketData),
 }
 
 unsafe impl Sync for AhciPort {}
@@ -68,7 +76,7 @@ impl AhciPort {
         cmd_list: *mut CommandHeader,
         cmd_list_phys: u64,
         cmd_list_len: usize,
-    ) -> Self {
+    ) -> Result<Self, CommandIssueError> {
         let cmd_reg = &regs.PxCMD;
         if !(cmd_reg.read(PxCMD::ST) == 0
             && cmd_reg.read(PxCMD::CR) == 0
@@ -96,14 +104,37 @@ impl AhciPort {
             regs.PxCMD.modify(PxCMD::ST::SET);
         }
 
-        Self {
+        let mut slf = Self {
             regs,
             phys_no,
             fis_buf,
             cmd_list,
             cmd_list_len,
             has_device,
-        }
+            identify_data: None,
+        };
+
+        slf.identify_data = if has_device {
+            let mut ident_buf = [0; 512];
+
+            if let Ok(identify_data) = slf.issue_data_in_command(&IdentifyCommand, &mut ident_buf) {
+                Some(AhciIdentifyData::Normal(identify_data))
+            } else {
+                let regs = slf.regs().unwrap();
+
+                if regs.lba1 != 0x14 || regs.lba2 != 0xEB {
+                    None
+                } else {
+                    let identify_data =
+                        slf.issue_data_in_command(&IdentifyPacketCommand, &mut ident_buf)?;
+                    Some(AhciIdentifyData::Packet(identify_data))
+                }
+            }
+        } else {
+            None
+        };
+
+        Ok(slf)
     }
 
     unsafe fn issue_command_prdt(
@@ -242,24 +273,16 @@ impl AhciPort {
             return Ok(None);
         }
 
-        let mut ident_buf = [0; 512];
-
-        if let Ok(identify_data) = self.issue_data_in_command(&IdentifyCommand, &mut ident_buf) {
-            Ok(Some(DeviceEnum::Ata(AtaDevice::new(self, identify_data)?)))
-        } else {
-            let regs = self.regs().unwrap();
-
-            if regs.lba1 != 0x14 || regs.lba2 != 0xEB {
-                return Ok(Some(DeviceEnum::Other));
-            }
-
-            let identify_data =
-                self.issue_data_in_command(&IdentifyPacketCommand, &mut ident_buf)?;
-
-            Ok(Some(DeviceEnum::Atapi(AtapiDevice::new(
-                self,
-                identify_data,
-            ))))
+        match self.identify_data.clone() {
+            Some(data) => match data {
+                AhciIdentifyData::Normal(data) => {
+                    Ok(Some(DeviceEnum::Ata(AtaDevice::new(self, data)?)))
+                }
+                AhciIdentifyData::Packet(data) => {
+                    Ok(Some(DeviceEnum::Atapi(AtapiDevice::new(self, data))))
+                }
+            },
+            None => Ok(None),
         }
     }