Fixing the 'device_instance' being None in some partition places (#902)

* Added a new return value from add_partition. Also added an exception to make sure `add_partition` can't continue silently

* Added a log of debugging to add_partition

* Removed a blank line (flake8)

* Misconfigured variable

* Added some more debugging information to partprobe

* FIX: disk layout: partprobe should be called and checked only for target device (#896)

* disk layout: partprobe should be called and checked only for target device

* disk layout: partprobe: removed unnecessary bash subprocess

* Properly defined BlockDevice() on Partition() creation. Also made sure mount-checks got some rrro handling and non-block devices should no longer attempt to return a size

Co-authored-by: Anton Hvornum <anton.feeds@gmail.com>
Co-authored-by: Victor Gavro <vgavro@gmail.com>

4 files changed, 62 insertions, 20 deletions

diff --git a/archinstall/lib/disk/blockdevice.py b/archinstall/lib/disk/blockdevice.py
index 5288f92b..5ffa06a8 100644
--- a/archinstall/lib/disk/blockdevice.py
+++ b/archinstall/lib/disk/blockdevice.py
@@ -241,7 +241,7 @@ class BlockDevice:
 		count = 0
 		while count < 5:
 			for partition_uuid, partition in self.partitions.items():
-				if partition.uuid == uuid:
+				if partition.uuid.lower() == uuid.lower():
 					return partition
 			else:
 				log(f"uuid {uuid} not found. Waiting for {count +1} time",level=logging.DEBUG)
diff --git a/archinstall/lib/disk/filesystem.py b/archinstall/lib/disk/filesystem.py
index a4fe2cff..b5dcdf85 100644
--- a/archinstall/lib/disk/filesystem.py
+++ b/archinstall/lib/disk/filesystem.py
@@ -89,12 +89,16 @@ class Filesystem:
 				print("Re-using partition_instance:", partition_instance)
 				partition['device_instance'] = partition_instance
 			else:
-				raise ValueError(f"{self}.load_layout() doesn't know how to continue without a new partition definition or a UUID ({partition.get('PARTUUID')}) on the device ({self.blockdevice.get_partition(uuid=partition_uuid)}).")
+				raise ValueError(f"{self}.load_layout() doesn't know how to continue without a new partition definition or a UUID ({partition.get('PARTUUID')}) on the device ({self.blockdevice.get_partition(uuid=partition.get('PARTUUID'))}).")
 
 			if partition.get('filesystem', {}).get('format', False):
+
 				# needed for backward compatibility with the introduction of the new "format_options"
 				format_options = partition.get('options',[]) + partition.get('filesystem',{}).get('format_options',[])
 				if partition.get('encrypted', False):
+					if not partition['device_instance']:
+						raise DiskError(f"Internal error caused us to loose the partition. Please report this issue upstream!")
+
 					if not partition.get('!password'):
 						if not storage['arguments'].get('!encryption-password'):
 							if storage['arguments'] == 'silent':
@@ -109,6 +113,7 @@ class Filesystem:
 						loopdev = f"{storage.get('ENC_IDENTIFIER', 'ai')}{pathlib.Path(partition['mountpoint']).name}loop"
 					else:
 						loopdev = f"{storage.get('ENC_IDENTIFIER', 'ai')}{pathlib.Path(partition['device_instance'].path).name}"
+					
 					partition['device_instance'].encrypt(password=partition['!password'])
 					# Immediately unlock the encrypted device to format the inner volume
 					with luks2(partition['device_instance'], loopdev, partition['!password'], auto_unmount=True) as unlocked_device:
@@ -129,6 +134,9 @@ class Filesystem:
 
 						unlocked_device.format(partition['filesystem']['format'], options=format_options)
 				elif partition.get('format', False):
+					if not partition['device_instance']:
+						raise DiskError(f"Internal error caused us to loose the partition. Please report this issue upstream!")
+
 					partition['device_instance'].format(partition['filesystem']['format'], options=format_options)
 
 			if partition.get('boot', False):
@@ -141,7 +149,13 @@ class Filesystem:
 				return partition
 
 	def partprobe(self) -> bool:
-		return SysCommand(f'partprobe {self.blockdevice.device}').exit_code == 0
+		result = SysCommand(f'partprobe {self.blockdevice.device}')
+		
+		if result.exit_code != 0:
+			log(f"Could not execute partprobe: {result!r}", level=logging.ERROR, fg="red")
+			raise DiskError(f"Could not run partprobe: {result!r}")
+
+		return True
 
 	def raw_parted(self, string: str) -> SysCommand:
 		if (cmd_handle := SysCommand(f'/usr/bin/parted -s {string}')).exit_code != 0:
@@ -157,9 +171,7 @@ class Filesystem:
 		:type string: str
 		"""
 		if (parted_handle := self.raw_parted(string)).exit_code == 0:
-			if self.partprobe():
-				return True
-			return False
+			return self.partprobe()
 		else:
 			raise DiskError(f"Parted failed to add a partition: {parted_handle}")
 
@@ -167,7 +179,7 @@ class Filesystem:
 		# TODO: Implement this with declarative profiles instead.
 		raise ValueError("Installation().use_entire_disk() has to be re-worked.")
 
-	def add_partition(self, partition_type :str, start :str, end :str, partition_format :Optional[str] = None) -> None:
+	def add_partition(self, partition_type :str, start :str, end :str, partition_format :Optional[str] = None) -> Partition:
 		log(f'Adding partition to {self.blockdevice}, {start}->{end}', level=logging.INFO)
 
 		previous_partition_uuids = {partition.uuid for partition in self.blockdevice.partitions.values()}
@@ -181,31 +193,41 @@ class Filesystem:
 		else:
 			parted_string = f'{self.blockdevice.device} mkpart {partition_type} {start} {end}'
 
+		log(f"Adding partition using the following parted command: {parted_string}", level=logging.DEBUG)
+
 		if self.parted(parted_string):
 			count = 0
 			while count < 10:
 				new_uuid = None
 				new_uuid_set = (previous_partition_uuids ^ {partition.uuid for partition in self.blockdevice.partitions.values()})
+
 				if len(new_uuid_set) > 0:
 					new_uuid = new_uuid_set.pop()
+
 				if new_uuid:
 					try:
 						return self.blockdevice.get_partition(new_uuid)
 					except Exception as err:
-						print('Blockdevice:', self.blockdevice)
-						print('Partitions:', self.blockdevice.partitions)
-						print('Partition set:', new_uuid_set)
-						print('New UUID:', [new_uuid])
-						print('get_partition():', self.blockdevice.get_partition)
+						log(f'Blockdevice: {self.blockdevice}', level=logging.ERROR, fg="red")
+						log(f'Partitions: {self.blockdevice.partitions}', level=logging.ERROR, fg="red")
+						log(f'Partition set: {new_uuid_set}', level=logging.ERROR, fg="red")
+						log(f'New UUID: {[new_uuid]}', level=logging.ERROR, fg="red")
+						log(f'get_partition(): {self.blockdevice.get_partition}', level=logging.ERROR, fg="red")
 						raise err
 				else:
 					count += 1
 					log(f"Could not get UUID for partition. Waiting before retry attempt {count} of 10 ...",level=logging.DEBUG)
 					time.sleep(float(storage['arguments'].get('disk-sleep', 0.2)))
 			else:
-				log("Add partition is exiting due to excessive wait time",level=logging.INFO)
+				log("Add partition is exiting due to excessive wait time", level=logging.ERROR, fg="red")
 				raise DiskError(f"New partition never showed up after adding new partition on {self}.")
 
+		# TODO: This should never be able to happen
+		log(f"Could not find the new PARTUUID after adding the partition.", level=logging.ERROR, fg="red")
+		log(f"Previous partitions: {previous_partition_uuids}", level=logging.ERROR, fg="red")
+		log(f"New partitions: {(previous_partition_uuids ^ {partition.uuid for partition in self.blockdevice.partitions.values()})}", level=logging.ERROR, fg="red")
+		raise DiskError(f"Could not add partition using: {parted_string}")
+
 	def set_name(self, partition: int, name: str) -> bool:
 		return self.parted(f'{self.blockdevice.device} name {partition + 1} "{name}"') == 0
 
diff --git a/archinstall/lib/disk/helpers.py b/archinstall/lib/disk/helpers.py
index e9f6bc10..26f701d2 100644
--- a/archinstall/lib/disk/helpers.py
+++ b/archinstall/lib/disk/helpers.py
@@ -190,10 +190,26 @@ def get_partitions_in_use(mountpoint :str) -> List[Partition]:
 
 	output = json.loads(output)
 	for target in output.get('filesystems', []):
-		mounts.append(Partition(target['source'], None, filesystem=target.get('fstype', None), mountpoint=target['target']))
+		# We need to create a BlockDevice() instead of 'None' here when creaiting Partition()
+		# Otherwise subsequent calls to .size etc will fail due to BlockDevice being None.
+
+		# So first, we create the partition without a BlockDevice and carefully only use it to get .real_device
+		# Note: doing print(partition) here will break because the above mentioned issue.
+		partition = Partition(target['source'], None, filesystem=target.get('fstype', None), mountpoint=target['target'])
+		partition = Partition(target['source'], partition.real_device, filesystem=target.get('fstype', None), mountpoint=target['target'])
+
+		# Once we have the real device (for instance /dev/nvme0n1p5) we can find the parent block device using
+		# (lsblk pkname lists both the partition and blockdevice, BD being the last entry)
+		result = SysCommand(f'lsblk -no pkname {partition.real_device}').decode().rstrip('\r\n').split('\r\n')[-1]
+		block_device = BlockDevice(f"/dev/{result}")
+
+		# Once we figured the block device out, we can properly create the partition object
+		partition = Partition(target['source'], block_device, filesystem=target.get('fstype', None), mountpoint=target['target'])
+
+		mounts.append(partition)
 
 		for child in target.get('children', []):
-			mounts.append(Partition(child['source'], None, filesystem=child.get('fstype', None), mountpoint=child['target']))
+			mounts.append(Partition(child['source'], block_device, filesystem=child.get('fstype', None), mountpoint=child['target']))
 
 	return mounts
 
diff --git a/archinstall/lib/disk/partition.py b/archinstall/lib/disk/partition.py
index b3db143a..afd34f20 100644
--- a/archinstall/lib/disk/partition.py
+++ b/archinstall/lib/disk/partition.py
@@ -46,11 +46,12 @@ class Partition:
 		except DiskError:
 			mount_information = {}
 
-		if self.mountpoint != mount_information.get('target', None) and mountpoint:
-			raise DiskError(f"{self} was given a mountpoint but the actual mountpoint differs: {mount_information.get('target', None)}")
+		if mount_information.get('target', None):
+			if self.mountpoint != mount_information.get('target', None) and mountpoint:
+				raise DiskError(f"{self} was given a mountpoint but the actual mountpoint differs: {mount_information.get('target', None)}")
 
-		if target := mount_information.get('target', None):
-			self.mountpoint = target
+			if target := mount_information.get('target', None):
+				self.mountpoint = target
 
 		if not self.filesystem and autodetect_filesystem:
 			if fstype := mount_information.get('fstype', get_filesystem_type(path)):
@@ -130,6 +131,9 @@ class Partition:
 
 				for device in lsblk['blockdevices']:
 					return convert_size_to_gb(device['size'])
+			elif handle.exit_code == 8192:
+				# Device is not a block device
+				return None
 
 			time.sleep(storage['DISK_TIMEOUTS'])
 
@@ -225,7 +229,7 @@ class Partition:
 		return bind_name
 
 	def partprobe(self) -> bool:
-		if SysCommand(f'partprobe {self.block_device.device}').exit_code == 0:
+		if self.block_device and SysCommand(f'partprobe {self.block_device.device}').exit_code == 0:
 			time.sleep(1)
 			return True
 		return False