Add support for unified kernel image (#1519)

4 files changed, 148 insertions, 46 deletions

diff --git a/archinstall/lib/global_menu.py b/archinstall/lib/global_menu.py
index 86c341a7..e4aa1235 100644
--- a/archinstall/lib/global_menu.py
+++ b/archinstall/lib/global_menu.py
@@ -4,6 +4,7 @@ from typing import Any, List, Optional, Dict, TYPE_CHECKING
 
 from . import disk
 from .general import secret
+from .hardware import SysInfo
 from .locale.locale_menu import LocaleConfiguration, LocaleMenu
 from .menu import Selector, AbstractMenu
 from .mirrors import MirrorConfiguration, MirrorMenu
@@ -20,6 +21,7 @@ from .interactions import ask_additional_packages_to_install
 from .interactions import ask_for_additional_users
 from .interactions import ask_for_audio_selection
 from .interactions import ask_for_bootloader
+from .interactions import ask_for_uki
 from .interactions import ask_for_swap
 from .interactions import ask_hostname
 from .interactions import ask_to_configure_network
@@ -85,6 +87,11 @@ class GlobalMenu(AbstractMenu):
 				lambda preset: ask_for_bootloader(preset),
 				display_func=lambda x: x.value,
 				default=Bootloader.get_default())
+		self._menu_options['uki'] = \
+			Selector(
+				_('Unified kernel images'),
+				lambda preset: ask_for_uki(preset),
+				default=False)
 		self._menu_options['hostname'] = \
 			Selector(
 				_('Hostname'),
@@ -216,6 +223,9 @@ class GlobalMenu(AbstractMenu):
 		self._menu_options['install'].update_description(text)
 
 	def post_callback(self, name: Optional[str] = None, value: Any = None):
+		if not SysInfo.has_uefi():
+			self._menu_options['uki'].set_enabled(False)
+
 		self._update_install_text(name, value)
 
 	def _install_text(self):
diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index 39298204..4d6c65b3 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -542,18 +542,13 @@ class Installer:
 
 		return True
 
-	def mkinitcpio(self, flags: List[str], locale_config: LocaleConfiguration) -> bool:
+	def mkinitcpio(self, flags: List[str]) -> bool:
 		for plugin in plugins.values():
 			if hasattr(plugin, 'on_mkinitcpio'):
 				# Allow plugins to override the usage of mkinitcpio altogether.
 				if plugin.on_mkinitcpio(self):
 					return True
 
-		# mkinitcpio will error out if there's no vconsole.
-		if (vconsole := Path(f"{self.target}/etc/vconsole.conf")).exists() is False:
-			with vconsole.open('w') as fh:
-				fh.write(f"KEYMAP={locale_config.kb_layout}\n")
-
 		with open(f'{self.target}/etc/mkinitcpio.conf', 'w') as mkinit:
 			mkinit.write(f"MODULES=({' '.join(self.modules)})\n")
 			mkinit.write(f"BINARIES=({' '.join(self._binaries)})\n")
@@ -587,6 +582,7 @@ class Installer:
 		self,
 		testing: bool = False,
 		multilib: bool = False,
+		mkinitcpio: bool = True,
 		hostname: str = 'archinstall',
 		locale_config: LocaleConfiguration = LocaleConfiguration.default()
 	):
@@ -674,7 +670,7 @@ class Installer:
 		# TODO: Use python functions for this
 		SysCommand(f'/usr/bin/arch-chroot {self.target} chmod 700 /root')
 
-		if not self.mkinitcpio(['-P'], locale_config):
+		if mkinitcpio and not self.mkinitcpio(['-P']):
 			error(f"Error generating initramfs (continuing anyway)")
 
 		self.helper_flags['base'] = True
@@ -783,7 +779,8 @@ class Installer:
 		self,
 		boot_partition: disk.PartitionModification,
 		root_partition: disk.PartitionModification,
-		efi_partition: Optional[disk.PartitionModification]
+		efi_partition: Optional[disk.PartitionModification],
+		uki_enabled: bool = False
 	):
 		self.pacman.strap('efibootmgr')
 
@@ -815,11 +812,18 @@ class Installer:
 		loader_dir = self.target / 'boot/loader'
 		loader_dir.mkdir(parents=True, exist_ok=True)
 
+		default_kernel = self.kernels[0]
+		if uki_enabled:
+			default_entry = f'arch-{default_kernel}.efi'
+		else:
+			entry_name = self.init_time + '_{kernel}{variant}.conf'
+			default_entry = entry_name.format(kernel=default_kernel, variant='')
+
+		default = f'default {default_entry}'
+
 		# Modify or create a loader.conf
 		loader_conf = loader_dir / 'loader.conf'
 
-		default = f'default {self.init_time}_{self.kernels[0]}.conf'
-
 		try:
 			loader_data = loader_conf.read_text().splitlines()
 		except FileNotFoundError:
@@ -837,6 +841,9 @@ class Installer:
 
 		loader_conf.write_text('\n'.join(loader_data) + '\n')
 
+		if uki_enabled:
+			return
+
 		# Ensure that the $BOOT/loader/entries/ directory exists before we try to create files in it
 		entries_dir = loader_dir / 'entries'
 		entries_dir.mkdir(parents=True, exist_ok=True)
@@ -867,7 +874,8 @@ class Installer:
 					options,
 				]
 
-				entry_conf = entries_dir / f'{self.init_time}_{kernel}{variant}.conf'
+				name = entry_name.format(kernel=kernel, variant=variant)
+				entry_conf = entries_dir / name
 				entry_conf.write_text('\n'.join(entry) + '\n')
 
 		self.helper_flags['bootloader'] = 'systemd'
@@ -876,17 +884,19 @@ class Installer:
 		self,
 		boot_partition: disk.PartitionModification,
 		root_partition: disk.PartitionModification,
-		efi_partition: Optional[disk.PartitionModification]
+		efi_partition: Optional[disk.PartitionModification],
+		uki_enabled: bool = False
 	):
 		self.pacman.strap('grub')  # no need?
 
-		grub_default = self.target / 'etc/default/grub'
-		config = grub_default.read_text()
+		if not uki_enabled:
+			grub_default = self.target / 'etc/default/grub'
+			config = grub_default.read_text()
 
-		kernel_parameters = ' '.join(self._get_kernel_params(root_partition, False, False))
-		config = re.sub(r'(GRUB_CMDLINE_LINUX=")("\n)', rf'\1{kernel_parameters}\2', config, 1)
+			kernel_parameters = ' '.join(self._get_kernel_params(root_partition, False, False))
+			config = re.sub(r'(GRUB_CMDLINE_LINUX=")("\n)', rf'\1{kernel_parameters}\2', config, 1)
 
-		grub_default.write_text(config)
+			grub_default.write_text(config)
 
 		info(f"GRUB boot partition: {boot_partition.dev_path}")
 
@@ -1067,7 +1077,8 @@ TIMEOUT=5
 	def _add_efistub_bootloader(
 		self,
 		boot_partition: disk.PartitionModification,
-		root_partition: disk.PartitionModification
+		root_partition: disk.PartitionModification,
+		uki_enabled: bool = False
 	):
 		self.pacman.strap('efibootmgr')
 
@@ -1078,41 +1089,103 @@ TIMEOUT=5
 		# points towards the same disk and/or partition.
 		# And in which case we should do some clean up.
 
-		microcode = []
+		if not uki_enabled:
+			loader = '/vmlinuz-{kernel}'
 
-		if ucode := self._get_microcode():
-			microcode.append(f'initrd=\\{ucode}')
-		else:
-			debug('Archinstall will not add any ucode to firmware boot entry.')
+			microcode = []
+
+			if ucode := self._get_microcode():
+				microcode.append(f'initrd=/{ucode}')
+			else:
+				debug('Archinstall will not add any ucode to firmware boot entry.')
 
-		kernel_parameters = self._get_kernel_params(root_partition)
+			entries = (
+				*microcode,
+				'initrd=/initramfs-{kernel}.img',
+				*self._get_kernel_params(root_partition)
+			)
+
+			cmdline = tuple(' '.join(entries))
+		else:
+			loader = '/EFI/Linux/arch-{kernel}.efi'
+			cmdline = tuple()
 
 		parent_dev_path = disk.device_handler.get_parent_device_path(boot_partition.safe_dev_path)
 
+		cmd_template = (
+			'efibootmgr',
+			'--create',
+			'--disk', str(parent_dev_path),
+			'--part', str(boot_partition.partn),
+			'--label', 'Arch Linux ({kernel})',
+			'--loader', loader,
+			'--unicode', *cmdline,
+			'--verbose'
+		)
+
 		for kernel in self.kernels:
 			# Setup the firmware entry
-			cmdline = [
-				*microcode,
-				f"initrd=\\initramfs-{kernel}.img",
-				*kernel_parameters,
-			]
-
-			cmd = [
-				'efibootmgr',
-				'--disk', str(parent_dev_path),
-				'--part', str(boot_partition.partn),
-				'--create',
-				'--label', f'Arch Linux ({kernel})',
-				'--loader', f"/vmlinuz-{kernel}",
-				'--unicode', ' '.join(cmdline),
-				'--verbose'
-			]
-
+			cmd = [arg.format(kernel=kernel) for arg in cmd_template]
 			SysCommand(cmd)
 
 		self.helper_flags['bootloader'] = "efistub"
 
-	def add_bootloader(self, bootloader: Bootloader):
+	def _config_uki(
+		self,
+		root_partition: disk.PartitionModification,
+		efi_partition: Optional[disk.PartitionModification]
+	):
+		if not efi_partition or not efi_partition.mountpoint:
+			raise ValueError(f'Could not detect ESP at mountpoint {self.target}')
+
+		# Set up kernel command line
+		with open(self.target / 'etc/kernel/cmdline', 'w') as cmdline:
+			kernel_parameters = self._get_kernel_params(root_partition)
+			cmdline.write(' '.join(kernel_parameters) + '\n')
+
+		ucode = self._get_microcode()
+
+		esp = efi_partition.mountpoint
+
+		diff_mountpoint = None
+		if esp != Path('/efi'):
+			diff_mountpoint = str(esp)
+
+		image_re = re.compile('(.+_image="/([^"]+).+\n)')
+		uki_re = re.compile('#((.+_uki=")/[^/]+(.+\n))')
+
+		# Modify .preset files
+		for kernel in self.kernels:
+			preset = self.target / 'etc/mkinitcpio.d' / (kernel + '.preset')
+			config = preset.read_text().splitlines(True)
+
+			for index, line in enumerate(config):
+				if not ucode and line.startswith('ALL_microcode='):
+					config[index] = '#' + line
+				# Avoid storing redundant image file
+				elif m := image_re.match(line):
+					image = self.target / m.group(2)
+					image.unlink(missing_ok=True)
+					config[index] = '#' + m.group(1)
+				elif m := uki_re.match(line):
+					if diff_mountpoint:
+						config[index] = m.group(2) + diff_mountpoint + m.group(3)
+					else:
+						config[index] = m.group(1)
+				elif line.startswith('#default_options='):
+					config[index] = line.removeprefix('#')
+
+			preset.write_text(''.join(config))
+
+		# Directory for the UKIs
+		uki_dir = self.target / esp.relative_to(Path('/')) / 'EFI/Linux'
+		uki_dir.mkdir(parents=True, exist_ok=True)
+
+		# Build the UKIs
+		if not self.mkinitcpio(['-P']):
+			error(f"Error generating initramfs (continuing anyway)")
+
+	def add_bootloader(self, bootloader: Bootloader, uki_enabled: bool = False):
 		"""
 		Adds a bootloader to the installation instance.
 		Archinstall supports one of three types:
@@ -1143,13 +1216,16 @@ TIMEOUT=5
 
 		info(f'Adding bootloader {bootloader.value} to {boot_partition.dev_path}')
 
+		if uki_enabled:
+			self._config_uki(root_partition, efi_partition)
+
 		match bootloader:
 			case Bootloader.Systemd:
-				self._add_systemd_bootloader(boot_partition, root_partition, efi_partition)
+				self._add_systemd_bootloader(boot_partition, root_partition, efi_partition, uki_enabled)
 			case Bootloader.Grub:
-				self._add_grub_bootloader(boot_partition, root_partition, efi_partition)
+				self._add_grub_bootloader(boot_partition, root_partition, efi_partition, uki_enabled)
 			case Bootloader.Efistub:
-				self._add_efistub_bootloader(boot_partition, root_partition)
+				self._add_efistub_bootloader(boot_partition, root_partition, uki_enabled)
 			case Bootloader.Limine:
 				self._add_limine_bootloader(boot_partition, root_partition)
 
diff --git a/archinstall/lib/interactions/__init__.py b/archinstall/lib/interactions/__init__.py
index 50c0012d..4b696a78 100644
--- a/archinstall/lib/interactions/__init__.py
+++ b/archinstall/lib/interactions/__init__.py
@@ -15,5 +15,5 @@ from .general_conf import (
 )
 
 from .system_conf import (
-	select_kernel, ask_for_bootloader, select_driver, ask_for_swap
+	select_kernel, ask_for_bootloader, ask_for_uki, select_driver, ask_for_swap
 )
diff --git a/archinstall/lib/interactions/system_conf.py b/archinstall/lib/interactions/system_conf.py
index 0e5e0f1e..aa72748e 100644
--- a/archinstall/lib/interactions/system_conf.py
+++ b/archinstall/lib/interactions/system_conf.py
@@ -65,6 +65,22 @@ def ask_for_bootloader(preset: Bootloader) -> Bootloader:
 	return preset
 
 
+def ask_for_uki(preset: bool = True) -> bool:
+	if preset:
+		preset_val = Menu.yes()
+	else:
+		preset_val = Menu.no()
+
+	prompt = _('Would you like to use unified kernel images?')
+	choice = Menu(prompt, Menu.yes_no(), default_option=Menu.no(), preset_values=preset_val).run()
+
+	match choice.type_:
+		case MenuSelectionType.Skip: return preset
+		case MenuSelectionType.Selection: return False if choice.value == Menu.no() else True
+
+	return preset
+
+
 def select_driver(options: List[GfxDriver] = [], current_value: Optional[GfxDriver] = None) -> Optional[GfxDriver]:
 	"""
 	Some what convoluted function, whose job is simple.