Add proper support for destruction of disklabels.

With this patch we can destroy a disklabel of one type and then create
a disklabel of a different type in the process of partitioning.

3 files changed, 43 insertions, 16 deletions

diff --git a/storage/devices.py b/storage/devices.py
index ea71190b8..1af4ecf16 100644
--- a/storage/devices.py
+++ b/storage/devices.py
@@ -1107,16 +1107,6 @@ class PartitionDevice(StorageDevice):
         return spec
 
     def _getPartedPartition(self):
-        if not self.exists:
-            return self._partedPartition
-
-        partitionDisk = self._partedPartition.disk.getPedDisk()
-        disklabelDisk = self.disk.format.partedDisk.getPedDisk()
-        if partitionDisk is not disklabelDisk:
-            # The disklabel's parted disk has been reset, reget our partition
-            log.debug("regetting partedPartition %s for %s because of PartedDisk reset" % (self._origPath, self.path))
-            self._setPartedPartition(self.disk.format.partedDisk.getPartitionByPath(self._origPath))
-
         return self._partedPartition
 
     def _setPartedPartition(self, partition):
@@ -1138,6 +1128,31 @@ class PartitionDevice(StorageDevice):
     partedPartition = property(lambda d: d._getPartedPartition(),
                                lambda d,p: d._setPartedPartition(p))
 
+    def resetPartedPartition(self):
+        """ Re-get self.partedPartition from the original disklabel. """
+        log_method_call(self, self.name)
+        if not self.exists:
+            return
+
+        # find the correct partition on the original parted.Disk since the
+        # name/number we're now using may no longer match
+        _disklabel = self.disk.originalFormat
+
+        if self.isExtended:
+            # getPartitionBySector doesn't work on extended partitions
+            _partition = _disklabel.extendedPartition
+            log.debug("extended lookup found partition %s"
+                        % devicePathToName(getattr(_partition, "path", None)))
+        else:
+            # lookup the partition by sector to avoid the renumbering
+            # nonsense entirely
+            _sector = self.partedPartition.geometry.start
+            _partition = _disklabel.partedDisk.getPartitionBySector(_sector)
+            log.debug("sector-based lookup found partition %s"
+                        % devicePathToName(getattr(_partition, "path", None)))
+
+        self.partedPartition = _partition
+
     def _getWeight(self):
         return self.req_base_weight
 
@@ -1341,8 +1356,11 @@ class PartitionDevice(StorageDevice):
             raise DeviceError("Cannot destroy non-leaf device", self.name)
 
         self.setupParents(orig=True)
-        self.disk.format.removePartition(self.partedPartition)
-        self.disk.format.commit()
+
+        # we should have already set self.partedPartition to point to the
+        # partition on the original disklabel
+        self.disk.originalFormat.removePartition(self.partedPartition)
+        self.disk.originalFormat.commit()
 
         self.exists = False
 
diff --git a/storage/devicetree.py b/storage/devicetree.py
index 75dde9c58..c96bfbdd6 100644
--- a/storage/devicetree.py
+++ b/storage/devicetree.py
@@ -627,16 +627,20 @@ class DeviceTree(object):
         for device in self.devices:
             if device.partitioned:
                 device.format.resetPartedDisk()
+                if device.originalFormat.type == "disklabel" and \
+                   device.originalFormat != device.format:
+                    device.originalFormat.resetPartedDisk()
 
         # reget parted.Partition for remaining preexisting devices
         for device in self.devices:
-            if isinstance(device, PartitionDevice):
-                p = device.partedPartition
+            if isinstance(device, PartitionDevice) and device.exists:
+                device.resetPartedPartition()
 
         # reget parted.Partition for existing devices we're removing
         for action in self._actions:
-            if isinstance(action.device, PartitionDevice):
-                p = action.device.partedPartition
+            if isinstance(action.device, PartitionDevice) and \
+               action.device.exists:
+                action.device.resetPartedPartition()
 
         # setup actions to create any extended partitions we added
         #
diff --git a/storage/formats/disklabel.py b/storage/formats/disklabel.py
index d4a706ddb..a76e45283 100644
--- a/storage/formats/disklabel.py
+++ b/storage/formats/disklabel.py
@@ -219,6 +219,11 @@ class DiskLabel(DeviceFormat):
             raise DeviceFormatError("device exists and is active")
 
         DeviceFormat.create(self, *args, **kwargs)
+
+        # We're relying on someone having called resetPartedDisk -- we
+        # could ensure a fresh disklabel by setting self._partedDisk to
+        # None right before calling self.commit(), but that might hide
+        # other problems.
         self.commit()
         self.exists = True