DiskMan: Add ability to reload partition table dynamically

2 files changed, 42 insertions, 13 deletions

diff --git a/diskman/disk.cc b/diskman/disk.cc
index 329edaa..5fd526a 100644
--- a/diskman/disk.cc
+++ b/diskman/disk.cc
@@ -79,7 +79,10 @@ Disk::Disk(void *creation, int type, bool partition) {
             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) {
+            if(fdisk_has_label(ctxt) != 1) {
+                /* Disk has no label, so consider it empty */
+                free_mb = contiguous_mb = total_mb;
+            } else 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". */
@@ -101,18 +104,8 @@ Disk::Disk(void *creation, int type, bool partition) {
     }
 
     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));
-                }
-                fdisk_unref_table(parts);
-            }
+        if(reload_partitions()) {
+            /* We're good */
         } else if(type == 0) {
             /* fallback to udev, if available */
             std::cerr << "Falling back to udev partition probing" << std::endl;
@@ -153,6 +146,37 @@ Disk::Disk(void *creation, int type, bool partition) {
     }
 }
 
+bool Disk::reload_partitions() {
+    bool success = false;
+    struct fdisk_context *ctxt = fdisk_new_context();
+    struct fdisk_table *parts = nullptr;
+
+    if(ctxt == nullptr) {
+        return false;
+    }
+
+    /* 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) {
+        goto destroy_context;
+    }
+
+    if(fdisk_get_partitions(ctxt, &parts) == 0) {
+        _partitions.clear();
+        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));
+        }
+        success = true;
+        fdisk_unref_table(parts);
+    }
+
+destroy_context:
+    fdisk_unref_context(ctxt);
+    return success;
+}
+
 const std::vector<Partition> Disk::partitions() const {
     if(!this->has_label()) {
         throw std::logic_error{ "attempt to retrieve partitions for non-labelled disk" };
diff --git a/diskman/disk.hh b/diskman/disk.hh
index dc149e7..668f209 100644
--- a/diskman/disk.hh
+++ b/diskman/disk.hh
@@ -80,6 +80,11 @@ public:
     /*! 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;
+
+    /*! Attempt to (re)load the partition table from the disk.
+     * @returns true if the partition table was loaded successfully.
+     * @note The partition array will not be modified if a failure occurs. */
+    bool reload_partitions();
 private:
     /*! The name of the disk ("sda") */
     std::string _name;