DiskMan: Add partition support, codify the source of creation objects

5 files changed, 267 insertions, 35 deletions

diff --git a/diskman/disk.cc b/diskman/disk.cc
index ddb357f..8006052 100644
--- a/diskman/disk.cc
+++ b/diskman/disk.cc
@@ -13,54 +13,148 @@
 #include "disk.hh"
 
 #include <cstring>
+#include <iostream>
+#include <libfdisk/libfdisk.h>
 #include <libudev.h>
+#include <stdexcept>
 
 namespace Horizon {
 namespace DiskMan {
 
-/*! The full serial number of the disk */
-std::string _full_serial;
-
 #define SAFE_SET(ivar, udev_call) \
     value = udev_call;\
     if(value != nullptr) {\
         ivar = std::string(value);\
     }
 
-Disk::Disk(void *creation, bool partition) {
-    struct udev_device *device = static_cast<struct udev_device *>(creation);
-    const char *value;
-
-    SAFE_SET(_name, udev_device_get_sysname(device));
-    SAFE_SET(_model, udev_device_get_property_value(device, "ID_MODEL"));
-    SAFE_SET(_node, udev_device_get_devnode(device));
-    SAFE_SET(_devpath, udev_device_get_devpath(device));
-
-    value = udev_device_get_property_value(device, "ID_PART_TABLE_TYPE");
-    if(value == nullptr) {
-        _has_label = false;
-        _label = Unknown;
-    } else {
-        _has_label = true;
-        if(::strcmp(value, "apm") == 0) {
-            _label = APM;
-        } else if(::strcmp(value, "dos") == 0) {
-            _label = MBR;
-        } else if(::strcmp(value, "gpt") == 0) {
-            _label = GPT;
-        } else {
+Disk::Disk(void *creation, int type, bool partition) {
+    switch(type) {
+    case 0: { /* udev */
+        struct udev_device *device = static_cast<struct udev_device *>(creation);
+        const char *value;
+
+        SAFE_SET(_name, udev_device_get_sysname(device));
+        SAFE_SET(_model, udev_device_get_property_value(device, "ID_MODEL"));
+        SAFE_SET(_full_serial, udev_device_get_property_value(device, "ID_SERIAL"));
+        SAFE_SET(_node, udev_device_get_devnode(device));
+        SAFE_SET(_devpath, udev_device_get_devpath(device));
+
+        value = udev_device_get_property_value(device, "ID_PART_TABLE_TYPE");
+        if(value == nullptr) {
+            _has_label = false;
             _label = Unknown;
+        } else {
+            _has_label = true;
+            if(::strcmp(value, "apm") == 0) {
+                _label = APM;
+            } else if(::strcmp(value, "dos") == 0) {
+                _label = MBR;
+            } else if(::strcmp(value, "gpt") == 0) {
+                _label = GPT;
+            } else {
+                _label = Unknown;
+            }
         }
+
+        value = udev_device_get_property_value(device, "ID_FS_TYPE");
+        if(value == nullptr) {
+            _has_fs = false;
+        } else {
+            _has_fs = true;
+            _fs_type = std::string(value);
+            SAFE_SET(_fs_label, udev_device_get_property_value(device, "ID_FS_LABEL"));
+        }
+        break;
+    }
+    default:
+        throw new std::invalid_argument{ "invalid type code" };
     }
 
-    value = udev_device_get_property_value(device, "ID_FS_TYPE");
-    if(value == nullptr) {
-        _has_fs = false;
-    } else {
-        _has_fs = true;
-        _fs_type = std::string(value);
-        SAFE_SET(_fs_label, udev_device_get_property_value(device, "ID_FS_LABEL"));
+    total_mb = free_mb = contiguous_mb = 0;
+
+    struct fdisk_context *ctxt = fdisk_new_context();
+    if(ctxt != nullptr) {
+        /* Open the device in read-only mode.  We don't need to write to it */
+        if(fdisk_assign_device(ctxt, _node.c_str(), 1) == 0) {
+            unsigned long ssize = fdisk_get_sector_size(ctxt);
+            total_mb = (fdisk_get_nsectors(ctxt) * ssize) / 1048576;
+            struct fdisk_table *frees = nullptr;
+            if(fdisk_get_freespaces(ctxt, &frees) == 0) {
+                for(size_t next = 0; next < fdisk_table_get_nents(frees);
+                    next++) {
+                    /* Each entry in frees is a "free space partition". */
+                    struct fdisk_partition *part =
+                            fdisk_table_get_partition(frees, next);
+                    fdisk_sector_t size;
+                    if(!fdisk_partition_has_size(part)) continue;
+                    size = (fdisk_partition_get_size(part) * ssize) / 1048576;
+                    free_mb += size;
+                    if(size > contiguous_mb) contiguous_mb = size;
+                }
+                fdisk_unref_table(frees);
+            }
+            /* We used to free ctxt here, but it's useful for partition probing */
+        } else {
+            fdisk_unref_context(ctxt);
+            ctxt = nullptr;
+        }
     }
+
+    if(partition) {
+        if(ctxt != nullptr) {
+            /* retrieve partitions using libfdisk */
+            struct fdisk_table *parts = nullptr;
+            if(fdisk_get_partitions(ctxt, &parts) == 0) {
+                for(size_t next = 0; next < fdisk_table_get_nents(parts);
+                    next++) {
+                    struct fdisk_partition *part =
+                            fdisk_table_get_partition(parts, next);
+                    _partitions.push_back(Partition(*this, part, 0));
+                }
+            }
+        } else if(type == 0) {
+            /* fallback to udev, if available */
+            std::cerr << "Falling back to udev partition probing" << std::endl;
+
+            struct udev_device *device = static_cast<struct udev_device *>(creation);
+            struct udev *udev = udev_device_get_udev(device);
+            struct udev_enumerate *part_enum = udev_enumerate_new(udev);
+            if(part_enum != NULL) {
+                struct udev_list_entry *first, *item;
+                struct udev_device *part_device = NULL;
+
+                udev_enumerate_add_match_subsystem(part_enum, "block");
+                udev_enumerate_add_match_property(part_enum, "DEVTYPE",
+                                                  "partition");
+                udev_enumerate_add_match_parent(part_enum, device);
+                udev_enumerate_scan_devices(part_enum);
+
+                first = udev_enumerate_get_list_entry(part_enum);
+                if(first != NULL) {
+                    udev_list_entry_foreach(item, first) {
+                        const char *path = udev_list_entry_get_name(item);
+                        if(part_device != NULL) udev_device_unref(part_device);
+                        part_device = udev_device_new_from_syspath(udev, path);
+                        _partitions.push_back(Partition(*this, part_device, 1));
+                    }
+                }
+            }
+        } else {
+            std::cerr << "Cannot load partitions for " << _name << std::endl;
+        }
+    }
+
+    if(ctxt != nullptr) {
+        fdisk_unref_context(ctxt);
+    }
+}
+
+const std::vector<Partition> Disk::partitions() const {
+    if(!this->has_label()) {
+        throw std::logic_error{ "attempt to retrieve partitions for non-labelled disk" };
+    }
+
+    return this->_partitions;
 }
 
 }
diff --git a/diskman/disk.hh b/diskman/disk.hh
index dcb9793..dc149e7 100644
--- a/diskman/disk.hh
+++ b/diskman/disk.hh
@@ -13,8 +13,10 @@
 #ifndef DISKMAN__DISK_HH
 #define DISKMAN__DISK_HH
 
-#include <memory>
 #include <string>
+#include <vector>
+
+#include "partition.hh"
 
 namespace Horizon {
 namespace DiskMan {
@@ -62,6 +64,22 @@ public:
     /*! Retrieve the label of the file system written on this disk.
      *  Only valid if has_fs() is true. */
     const std::string fs_label() const { return this->_fs_label; }
+
+    /*! Retrieve the total size of the disk.
+     * @returns The size of the disk, in mebibytes (MiB). */
+    uint32_t total_size() const { return this->total_mb; }
+    /*! Retrieve the amount of free space available on the disk.
+     * @returns The size of free space on the disk, in mebibytes (MiB). */
+    uint32_t free_space() const { return this->free_mb; }
+    /*! Retrieve the size of the largest contiguous block of free space on
+     *  the disk.
+     * @returns The size of the largest contiguous free space block on the
+     *          disk, in mebibytes (MiB). */
+    uint32_t contiguous_block() const { return this->contiguous_mb; }
+
+    /*! Retrieve the partitions contained in the label, if any.
+     * @note You may only call this method if *has_label* is true. */
+    const std::vector<Partition> partitions() const;
 private:
     /*! The name of the disk ("sda") */
     std::string _name;
@@ -78,6 +96,8 @@ private:
     bool _has_label;
     /*! The type of disk label used, if +has_label+ is true */
     enum Label _label;
+    /*! Partitions inside the disklabel */
+    std::vector<Partition> _partitions;
 
     /*! Whether this disk has a direct filesystem (non-labelled) */
     bool _has_fs;
@@ -86,7 +106,14 @@ private:
     /*! The label of the file system, if +has_fs+ is true */
     std::string _fs_label;
 
-    Disk(void *creation, bool partition);
+    /*! Total size of this disk, in mebibytes */
+    uint32_t total_mb;
+    /*! Free space on this disk, in mebibytes */
+    uint32_t free_mb;
+    /*! Largest contiguous block of free space on this disk, in mebibytes */
+    uint32_t contiguous_mb;
+
+    Disk(void *creation, int type, bool partition);
     friend class DiskMan;
 };
 
diff --git a/diskman/diskman.cc b/diskman/diskman.cc
index 32a2c40..a7d4a26 100644
--- a/diskman/diskman.cc
+++ b/diskman/diskman.cc
@@ -71,7 +71,11 @@ std::vector<Disk> DiskMan::find_disks(bool include_part, bool include_vg,
             /* Skip LVM volumes if requested. */
             continue;
         }
-        disks.push_back(Disk(device, include_part));
+        if(udev_device_get_property_value(device, "ID_CDROM") != nullptr) {
+            /* REQ: UI.Partition.Install.Ignore */
+            continue;
+        }
+        disks.push_back(Disk(device, 0, include_part));
     }
 
     return disks;
diff --git a/diskman/partition.cc b/diskman/partition.cc
index e69de29..308d002 100644
--- a/diskman/partition.cc
+++ b/diskman/partition.cc
@@ -0,0 +1,61 @@
+/*
+ * partition.cc - Implementation of the Partition class
+ * diskman, the Disk Manipulation library for
+ * Project Horizon
+ *
+ * Copyright (c) 2020 Adélie Linux and contributors.  All rights reserved.
+ * This code is licensed under the AGPL 3.0 license, as noted in the
+ * LICENSE-code file in the root directory of this repository.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+#include "partition.hh"
+#include "disk.hh"
+
+#include <blkid/blkid.h>
+#include <iostream>
+#include <libfdisk/libfdisk.h>
+#include <libudev.h>
+#include <stdexcept>
+
+namespace Horizon {
+namespace DiskMan {
+
+Partition::Partition(Disk &d, void *creation, int type) {
+    switch(type) {
+    case 0: { /* libfdisk */
+        struct fdisk_partition *part = static_cast<struct fdisk_partition *>(creation);
+        if(fdisk_partition_has_size(part)) {
+            /* XXX BUG FIXME TODO sector size */
+            this->_size = fdisk_partition_get_size(part) * 512;
+        } else {
+            this->_size = 0;
+        }
+        const char *name = fdisk_partname(d.node().c_str(),
+                                          fdisk_partition_get_partno(part) + 1);
+        const char *value;
+        value = blkid_get_tag_value(nullptr, "TYPE", name);
+        if(value != nullptr) this->_fs_type = std::string(value);
+        value = blkid_get_tag_value(nullptr, "LABEL", name);
+        if(value != nullptr) this->_label = std::string(value);
+        break;
+    }
+    case 1: { /* udev */
+        struct udev_device *dev = static_cast<struct udev_device *>(creation);
+        const char *value;
+        value = udev_device_get_property_value(dev, "ID_FS_TYPE");
+        if(value != nullptr) this->_fs_type = std::string(value);
+        value = udev_device_get_property_value(dev, "ID_FS_LABEL");
+        if(value != nullptr) this->_label = std::string(value);
+        value = udev_device_get_property_value(dev, "ID_PART_ENTRY_SIZE");
+        if(value != nullptr) this->_size = strtoull(value, nullptr, 10) * 512;
+        break;
+    }
+    default:
+        throw std::invalid_argument{ "invalid type code" };
+    }
+}
+
+}
+}
diff --git a/diskman/partition.hh b/diskman/partition.hh
index e69de29..93090d8 100644
--- a/diskman/partition.hh
+++ b/diskman/partition.hh
@@ -0,0 +1,46 @@
+/*
+ * partition.hh - Definition of the Partition class
+ * diskman, the Disk Manipulation library for
+ * Project Horizon
+ *
+ * Copyright (c) 2020 Adélie Linux and contributors.  All rights reserved.
+ * This code is licensed under the AGPL 3.0 license, as noted in the
+ * LICENSE-code file in the root directory of this repository.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+#ifndef DISKMAN__PARTITION_HH
+#define DISKMAN__PARTITION_HH
+
+#include <cstdint>
+#include <string>
+
+namespace Horizon {
+namespace DiskMan {
+
+class Disk;
+
+class Partition {
+public:
+    /*! Retrieve the size of the partition in bytes. */
+    uint64_t size() const { return this->_size; };
+    /*! Retrieve the file system type of this partition. */
+    const std::string fstype() const { return this->_fs_type; }
+    /*! Retrieve the label of the file system on this partition. */
+    const std::string label() const { return this->_label; }
+private:
+    /*! The size of the partition, in bytes. */
+    uint64_t _size;
+    /*! The type of the file system on this partition (if any) */
+    std::string _fs_type;
+    /*! The label of the file system on this partition (if any) */
+    std::string _label;
+    Partition(Disk &d, void *creation, int type);
+    friend class Disk;
+};
+
+}
+}
+
+#endif  /* !DISKMAN__PARTITION_HH */