DiskIO: detect physical drive size

fastfetch-cli · Dec 3, 2023 · 2979bfe · 2979bfe
1 parent 42fd1a9
commit 2979bfe
Show file tree

Hide file tree

Showing 6 changed files with 28 additions and 0 deletions.
diff --git a/src/detection/diskio/diskio.h b/src/detection/diskio/diskio.h
@@ -14,6 +14,7 @@ typedef struct FFDiskIOResult
     FFstrbuf name;
     FFstrbuf interconnect;
     FFDiskIOPhysicalType type;
+    uint64_t size;
     FFstrbuf devPath;
     uint64_t bytesRead;
     uint64_t readCount;

diff --git a/src/detection/diskio/diskio_apple.c b/src/detection/diskio/diskio_apple.c
@@ -49,6 +49,7 @@ const char* ffDiskIOGetIoCounters(FFlist* result, FFDiskIOOptions* options)
         ffStrbufInitS(&device->name, deviceName);
         ffStrbufInit(&device->devPath);
         device->type = FF_DISKIO_PHYSICAL_TYPE_UNKNOWN;
+        device->size = 0;
 
         ffCfDictGetInt64(statistics, CFSTR(kIOBlockStorageDriverStatisticsBytesReadKey), (int64_t*) &device->bytesRead);
         ffCfDictGetInt64(statistics, CFSTR(kIOBlockStorageDriverStatisticsBytesWrittenKey), (int64_t*) &device->bytesWritten);

diff --git a/src/detection/diskio/diskio_bsd.c b/src/detection/diskio/diskio_bsd.c
@@ -27,6 +27,7 @@ const char* ffDiskIOGetIoCounters(FFlist* result, FFDiskIOOptions* options)
             continue;
 
         FFDiskIOResult* device = (FFDiskIOResult*) ffListAdd(result);
+        device->size = 0;
         ffStrbufInitS(&device->name, deviceName);
         ffStrbufInitF(&device->devPath, "/dev/%s", deviceName);
         device->type = FF_DISKIO_PHYSICAL_TYPE_UNKNOWN;

diff --git a/src/detection/diskio/diskio_linux.c b/src/detection/diskio/diskio_linux.c
@@ -102,6 +102,12 @@ const char* ffDiskIOGetIoCounters(FFlist* result, FFDiskIOOptions* options)
             device->type = ffStrbufEqualS(&buffer, "1") ? FF_DISKIO_PHYSICAL_TYPE_HDD : FF_DISKIO_PHYSICAL_TYPE_SSD;
         else
             device->type = FF_DISKIO_PHYSICAL_TYPE_UNKNOWN;
+
+        snprintf(pathSysBlock, PATH_MAX, "/sys/block/%s/size", devName);
+        if (ffReadFileBuffer(pathSysBlock, &buffer))
+            device->size = ffStrbufToUInt(&buffer, 0);
+        else
+            device->size = 0;
     }
 
     return NULL;

diff --git a/src/detection/diskio/diskio_windows.c b/src/detection/diskio/diskio_windows.c
@@ -82,6 +82,20 @@ const char* ffDiskIOGetIoCounters(FFlist* result, FFDiskIOOptions* options)
                 device->type = dspd.IncursSeekPenalty ? FF_DISKIO_PHYSICAL_TYPE_HDD : FF_DISKIO_PHYSICAL_TYPE_SSD;
             else
                 device->type = FF_DISKIO_PHYSICAL_TYPE_UNKNOWN;
+
+            DISK_GEOMETRY_EX dge = {};
+            if(DeviceIoControl(
+                hDevice,
+                IOCTL_DISK_GET_DRIVE_GEOMETRY_EX,
+                NULL,
+                0,
+                &dge,
+                sizeof(dge),
+                &retSize,
+                NULL))
+                device->size = (uint64_t) dge.DiskSize.QuadPart;
+            else
+                device->size = 0;
         }
     }
 

diff --git a/src/modules/diskio/diskio.c b/src/modules/diskio/diskio.c
@@ -87,6 +87,8 @@ void ffPrintDiskIO(FFDiskIOOptions* options)
                     physicalType = "Unknown";
                     break;
             }
+            FF_STRBUF_AUTO_DESTROY sizePretty = ffStrbufCreate();
+            ffParseSize(dev->size, &sizePretty);
 
             ffPrintFormatString(key.chars, 0, &options->moduleArgs, FF_PRINT_TYPE_NO_CUSTOM_KEY, FF_DISKIO_NUM_FORMAT_ARGS, (FFformatarg[]){
                 {FF_FORMAT_ARG_TYPE_STRBUF, &buffer},
@@ -99,6 +101,7 @@ void ffPrintDiskIO(FFDiskIOOptions* options)
                 {FF_FORMAT_ARG_TYPE_UINT64, &dev->bytesWritten},
                 {FF_FORMAT_ARG_TYPE_UINT64, &dev->readCount},
                 {FF_FORMAT_ARG_TYPE_UINT64, &dev->writeCount},
+                {FF_FORMAT_ARG_TYPE_STRBUF, &sizePretty},
             });
         }
         ++index;
@@ -198,6 +201,7 @@ void ffGenerateDiskIOJsonResult(FFDiskIOOptions* options, yyjson_mut_doc* doc, y
         yyjson_mut_obj_add_uint(doc, obj, "bytesWritten", dev->bytesWritten);
         yyjson_mut_obj_add_uint(doc, obj, "readCount", dev->readCount);
         yyjson_mut_obj_add_uint(doc, obj, "writeCount", dev->writeCount);
+        yyjson_mut_obj_add_uint(doc, obj, "size", dev->size);
     }
 
     FF_LIST_FOR_EACH(FFDiskIOResult, dev, result)
@@ -221,6 +225,7 @@ void ffPrintDiskIOHelpFormat(void)
         "Size of data written per second (in bytes)",
         "Number of reads",
         "Number of writes",
+        "Device size (formatted)",
     });
 }