From dea66ae9004e34271aad8af8bb1ff3b7459b6773 Mon Sep 17 00:00:00 2001
From: Serhiy Mytrovtsiy <mitrovtsiy@ukr.net>
Date: Sat, 19 Nov 2022 16:40:49 +0100
Subject: [PATCH] fix: changed the way to get a disk capacity (#1180)

---
 Modules/Disk/readers.swift | 36 ++++++++++++++++++++++++++++--------
 1 file changed, 28 insertions(+), 8 deletions(-)

diff --git a/Modules/Disk/readers.swift b/Modules/Disk/readers.swift
index d04f12a8..5a79ff73 100644
--- a/Modules/Disk/readers.swift
+++ b/Modules/Disk/readers.swift
@@ -51,6 +51,7 @@ internal class CapacityReader: Reader<Disks> {
                         if var d = driveDetails(disk, removableState: removableState) {
                             if let path = d.path {
                                 d.free = self.freeDiskSpaceInBytes(path)
+                                d.size = self.totalDiskSpaceInBytes(path)
                             }
                             self.list.append(d)
                             self.list.sort()
@@ -92,6 +93,30 @@ internal class CapacityReader: Reader<Disks> {
         
         return 0
     }
+    
+    private func totalDiskSpaceInBytes(_ path: URL) -> Int64 {
+        do {
+            let systemAttributes = try FileManager.default.attributesOfFileSystem(forPath: path.path)
+            if let totalSpace = (systemAttributes[FileAttributeKey.systemSize] as? NSNumber)?.int64Value {
+                return totalSpace
+            }
+        } catch let err {
+            error("error retrieving total space #2: \(err.localizedDescription)", log: self.log)
+        }
+        
+        do {
+            if let url = URL(string: path.absoluteString) {
+                let values = try url.resourceValues(forKeys: [.volumeTotalCapacityKey])
+                if let space = values.volumeAvailableCapacityForImportantUsage, space != 0 {
+                    return space
+                }
+            }
+        } catch let err {
+            error("error retrieving total space #1: \(err.localizedDescription)", log: self.log)
+        }
+        
+        return 0
+    }
 }
 
 internal class ActivityReader: Reader<Disks> {
@@ -206,9 +231,6 @@ private func driveDetails(_ disk: DADisk, removableState: Bool) -> drive? {
                     }
                 }
             }
-            if let mediaSize = dict[kDADiskDescriptionMediaSizeKey as String] {
-                d.size = Int64(truncating: mediaSize as! NSNumber)
-            }
             if let deviceModel = dict[kDADiskDescriptionDeviceModelKey as String] {
                 d.model = (deviceModel as! String).trimmingCharacters(in: .whitespacesAndNewlines)
             }
@@ -256,11 +278,9 @@ public func getDeviceIOParent(_ obj: io_registry_entry_t, level: Int) -> io_regi
         return nil
     }
     
-    for _ in 1...level {
-        if IORegistryEntryGetParentEntry(parent, kIOServicePlane, &parent) != KERN_SUCCESS {
-            IOObjectRelease(parent)
-            return nil
-        }
+    for _ in 1...level where IORegistryEntryGetParentEntry(parent, kIOServicePlane, &parent) != KERN_SUCCESS {
+        IOObjectRelease(parent)
+        return nil
     }
     
     return parent