kernel/ATA: Calculate block sizes for ATA devices as well

kernel/src/arch/x86_64/disk/ATA.cpp

@@ -4,6 +4,7 @@
 #include "arch/Timer.h"
 #include "arch/x86_64/IO.h"
 #include "memory/MemoryManager.h"
+#include <luna/SafeArithmetic.h>
 #include <luna/Vector.h>
 
 SharedPtr<ATA::Controller> g_controller;
@@ -419,13 +420,38 @@ namespace ATA
                 return false;
             }
 
+            m_is_lba48 = true;
+
             // FIXME: This assumes the host machine is little-endian.
             u32 last_lba = __builtin_bswap32(reply.last_lba);
             u32 sector_size = __builtin_bswap32(reply.sector_size);
 
-            kinfoln("ata: ATAPI drive %d capacity information: Last LBA=%u, Sector Size=%u, Total Capacity=%u",
+            m_block_count = last_lba + 1;
-                    m_drive_index, last_lba, sector_size, (last_lba + 1) * sector_size);
+            m_block_size = sector_size;
         }
+        else
+        {
+            if (m_identify_data.big_lba) m_is_lba48 = true;
+
+            if (m_is_lba48) m_block_count = m_identify_data.sectors_48;
+            else
+                m_block_count = m_identify_data.sectors_28;
+
+            // FIXME: Should we check for CHS?
+
+            // FIXME: Maybe a different block size is in use? Detect that.
+            m_block_size = 512;
+        }
+
+        u64 total_capacity;
+        if (!safe_mul(m_block_count, m_block_size).try_set_value(total_capacity))
+        {
+            kwarnln("ata: Drive %d's total capacity is too large", m_drive_index);
+            return false;
+        }
+
+        kinfoln("ata: Drive %d capacity information: Block Count=%lu, Block Size=%lu, Total Capacity=%lu",
+                m_drive_index, m_block_count, m_block_size, total_capacity);
 
         return true;
     }

kernel/src/arch/x86_64/disk/ATA.h

@@ -68,6 +68,31 @@ namespace ATA
         BMS_DMAMode = 0x1
     };
 
+    struct ATAIdentify
+    {
+        u16 flags;
+        u16 unused1[9];
+        char serial[20];
+        u16 unused2[3];
+        char firmware[8];
+        char model[40];
+        u16 sectors_per_int;
+        u16 unused3;
+        u16 capabilities[2];
+        u16 unused4[2];
+        u16 valid_ext_data;
+        u16 unused5[5];
+        u16 size_of_rw_mult;
+        u32 sectors_28;
+        u16 unused6[21];
+        u16 unused7 : 10;
+        u16 big_lba : 1;
+        u16 unused8 : 5;
+        u16 unused9[17];
+        u64 sectors_48;
+        u16 unused10[152];
+    };
+
     enum ATAPICommand : u8
     {
         ATAPI_ReadCapacity = 0x25,
@@ -118,12 +143,16 @@ namespace ATA
         u8 m_drive_index;
         union {
             u16 m_identify_words[256];
-            u8 m_identify_data[512];
+            ATAIdentify m_identify_data;
         };
 
         bool m_is_atapi { false };
         bool m_uses_dma { true };
 
+        bool m_is_lba48;
+        u64 m_block_count;
+        u64 m_block_size;
+
         volatile prdt_entry* m_dma_prdt;
         u64 m_dma_prdt_phys;
         volatile void* m_dma_mem;