From 78ab38626853047babc50e23e25f3989deab2c50 Mon Sep 17 00:00:00 2001
From: Ian Douglas Scott <ian@iandouglasscott.com>
Date: Thu, 4 Jan 2018 12:25:47 -0800
Subject: [PATCH] Move capacity reading to seperate method

---
 ahcid/src/ahci/disk_atapi.rs | 25 ++++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/ahcid/src/ahci/disk_atapi.rs b/ahcid/src/ahci/disk_atapi.rs
index 3ddc531690..8159c39bf3 100644
--- a/ahcid/src/ahci/disk_atapi.rs
+++ b/ahcid/src/ahci/disk_atapi.rs
@@ -50,6 +50,19 @@ impl DiskATAPI {
             buf: buf
         })
     }
+
+    fn read_capacity(&mut self) -> Result<(u32, u32)> {
+        // TODO: only query when needed (disk changed)
+
+        let mut cmd = [0; 16];
+        cmd[0] = SCSI_READ_CAPACITY;
+        self.port.packet(&cmd, 8, &mut self.clb, &mut self.ctbas, &mut self.buf)?;
+
+        let blk_count = BigEndian::read_u32(&self.buf[0..4]);
+        let blk_size = BigEndian::read_u32(&self.buf[4..8]);
+
+        Ok((blk_count, blk_size))
+    }
 }
 
 impl Disk for DiskATAPI {
@@ -58,16 +71,10 @@ impl Disk for DiskATAPI {
     }
 
     fn size(&mut self) -> u64 {
-        let mut cmd = [0; 16];
-        cmd[0] = SCSI_READ_CAPACITY;
-        if let Err(_) = self.port.packet(&cmd, 8, &mut self.clb, &mut self.ctbas, &mut self.buf) {
-            return 0; // XXX
+        match self.read_capacity() {
+            Ok((blk_count, blk_size)) => (blk_count as u64) * (blk_size as u64),
+            Err(_) => 0 // XXX
         }
-
-        let blk_count = BigEndian::read_u32(&self.buf[0..4]);
-        let blk_size = BigEndian::read_u32(&self.buf[4..8]);
-
-        (blk_count as u64) * (blk_size as u64)
     }
 
     fn read(&mut self, _block: u64, _buffer: &mut [u8]) -> Result<usize> {