Improve UKI detection

osutils/src/efivar.rs

@@ -105,9 +105,13 @@ fn read_efi_variable(guid: &str, variable: &str) -> Result<Vec<u8>, TridentError
     Ok(data[4..].to_vec())
 }
 
-/// Returns whether the LoaderEntrySelected EFI variable is set.
-pub fn current_var_set() -> bool {
-    read_efi_variable(BOOTLOADER_INTERFACE_GUID, LOADER_ENTRY_SELECTED).is_ok()
+/// Returns whether the LoaderEntrySelected EFI variable is set and indicates a UKI boot.
+pub fn current_var_is_uki() -> bool {
+    let Ok(current) = read_efi_variable(BOOTLOADER_INTERFACE_GUID, LOADER_ENTRY_SELECTED) else {
+        return false;
+    };
+
+    decode_utf16le(&current).ends_with(".efi")
 }
 
 /// Returns the value of the LoaderEntrySelected EFI variable. This is the current boot entry.
@@ -177,7 +181,7 @@ mod functional_test {
     #[functional_test(feature = "helpers")]
     fn test_set_default_to_current() {
         // Generate a random current entry
-        let current_entry = format!("CurrentEntry-{}", rand::random::<u32>());
+        let current_entry = format!("CurrentEntry-{}.efi", rand::random::<u32>());
 
         set_efi_variable(
             &format!("{BOOTLOADER_INTERFACE_GUID}-{LOADER_ENTRY_SELECTED}"),
@@ -186,7 +190,7 @@ mod functional_test {
         .unwrap();
 
         // Check that the current entry is set
-        assert!(current_var_set());
+        assert!(current_var_is_uki());
         assert_eq!(read_current_var().unwrap(), current_entry);
 
         // Now set the default to the current entry

src/engine/boot/mod.rs

@@ -27,7 +27,7 @@ impl Subsystem for BootSubsystem {
 
     #[tracing::instrument(name = "boot_configuration", skip_all)]
     fn configure(&mut self, ctx: &EngineContext) -> Result<(), TridentError> {
-        if ctx.is_uki_image()? {
+        if ctx.is_uki()? {
             debug!("Skipping grub configuration because UKI is in use");
             return Ok(());
         }

src/engine/clean_install.rs

@@ -13,8 +13,8 @@ use osutils::{chroot, container, mount, mountpoint, path::join_relative};
 use trident_api::{
     config::{HostConfiguration, Operations},
     constants::{
-        internal_params::NO_TRANSITION, ESP_MOUNT_POINT_PATH, ROOT_MOUNT_POINT_PATH,
-        UPDATE_ROOT_PATH,
+        internal_params::{ENABLE_UKI_SUPPORT, NO_TRANSITION},
+        ESP_MOUNT_POINT_PATH, ROOT_MOUNT_POINT_PATH, UPDATE_ROOT_PATH,
     },
     error::{
         InitializationError, InternalError, InvalidInputError, ReportError, ServicingError,
@@ -197,7 +197,7 @@ fn stage_clean_install(
         partition_paths: Default::default(), // Will be initialized later
         disk_uuids: Default::default(),      // Will be initialized later
         install_index: 0,                    // Will be initialized later
-        is_uki: Some(image.is_uki()),
+        is_uki: Some(image.is_uki() || host_config.internal_params.get_flag(ENABLE_UKI_SUPPORT)),
         image: Some(image),
         storage_graph: engine::build_storage_graph(&host_config.storage)?, // Build storage graph
         filesystems: Vec::new(), // Will be populated after dynamic validation

src/engine/context/mod.rs

@@ -9,8 +9,8 @@ use log::{debug, trace};
 
 use trident_api::{
     config::{HostConfiguration, Partition, VerityDevice},
-    constants::{internal_params::ENABLE_UKI_SUPPORT, ROOT_MOUNT_POINT_PATH},
-    error::TridentError,
+    constants::ROOT_MOUNT_POINT_PATH,
+    error::{InternalError, ReportError, TridentError},
     status::{AbVolumeSelection, ServicingType},
     storage_graph::graph::StorageGraph,
     BlockDeviceId,
@@ -290,18 +290,13 @@ impl EngineContext {
         self.storage_graph.block_device_size(device)
     }
 
-    pub(crate) fn is_uki_image(&self) -> Result<bool, TridentError> {
-        if self.spec.internal_params.get_flag(ENABLE_UKI_SUPPORT) {
-            trace!("internal param {ENABLE_UKI_SUPPORT} specified: UKI image");
-            Ok(true)
-        } else if let Some(is_uki) = self.is_uki {
-            trace!("uki configured as {is_uki}");
-            Ok(is_uki)
-        } else {
-            Err(TridentError::internal(
-                "is_uki() called without it being set",
-            ))
-        }
+    /// Convience method to check if the current context is a UKI context and return a suitable
+    /// error if the flag isn't set.
+    #[track_caller]
+    pub(crate) fn is_uki(&self) -> Result<bool, TridentError> {
+        self.is_uki.structured(InternalError::Internal(
+            "is_uki() called without it being set",
+        ))
     }
 }
 

src/engine/mod.rs

@@ -292,7 +292,7 @@ fn configure(
     let use_overlay = (ctx.servicing_type == ServicingType::CleanInstall
         || ctx.servicing_type == ServicingType::AbUpdate)
         && ctx.storage_graph.root_fs_is_verity()
-        && !ctx.is_uki_image()?;
+        && !ctx.is_uki()?;
 
     info!("Starting step 'Configure'");
     for subsystem in subsystems {

src/engine/rollback.rs

@@ -7,7 +7,7 @@ use osutils::{block_devices, container, efivar, lsblk, pcrlock, veritysetup, vir
 use sysdefs::tpm2::Pcr;
 
 use trident_api::{
-    constants::internal_params::{ENABLE_UKI_SUPPORT, VIRTDEPLOY_BOOT_ORDER_WORKAROUND},
+    constants::internal_params::VIRTDEPLOY_BOOT_ORDER_WORKAROUND,
     error::{InternalError, ReportError, ServicingError, TridentError, TridentResultExt},
     status::{AbVolumeSelection, ServicingState, ServicingType},
     BlockDeviceId,
@@ -43,13 +43,7 @@ pub fn validate_boot(datastore: &mut DataStore) -> Result<(), TridentError> {
         image: None, // Not used for boot validation logic
         storage_graph: engine::build_storage_graph(&datastore.host_status().spec.storage)?, // Build storage graph
         filesystems: Vec::new(), // Left empty since context does not have image
-        is_uki: Some(
-            datastore
-                .host_status()
-                .spec
-                .internal_params
-                .get_flag(ENABLE_UKI_SUPPORT),
-        ),
+        is_uki: Some(efivar::current_var_is_uki()),
     };
 
     // Get the block device path of the current root
@@ -81,9 +75,8 @@ pub fn validate_boot(datastore: &mut DataStore) -> Result<(), TridentError> {
                 .message("Failed to persist boot order after reboot")?;
         }
 
-        // If the bootloader set the LoaderEntrySelected variable, then make its value the default
-        // boot entry. Systemd-boot sets this variable, but GRUB does not.
-        if efivar::current_var_set() {
+        // In UKI mode, set systemd-boot's default boot option to the currently running one.
+        if ctx.is_uki()? {
             efivar::set_default_to_current()
                 .message("Failed to set default boot entry to current")?;
         }
@@ -98,7 +91,7 @@ pub fn validate_boot(datastore: &mut DataStore) -> Result<(), TridentError> {
             .internal_params
             .get_flag("overridePcrlockEncryption")
             || container::is_running_in_container()?;
-        if ctx.is_uki_image()? && ctx.spec.storage.encryption.is_some() {
+        if ctx.is_uki()? && ctx.spec.storage.encryption.is_some() {
             if !override_pcrlock_encryption {
                 debug!("Regenerating pcrlock policy for current boot");
                 // TODO: Add PCR 7 once SecureBoot is enabled in a follow up PR. Related ADO task:

src/engine/storage/raid.rs

@@ -44,7 +44,7 @@ fn create(config: SoftwareRaidArray, ctx: &EngineContext) -> Result<(), Error> {
 
     info!("Initializing '{}': creating RAID array", config.id);
 
-    if ctx.is_uki_image().unstructured("UKI setting unknown")? {
+    if ctx.is_uki().unstructured("UKI setting unknown")? {
         // If UKI support is enabled, we need to create the RAID array with the
         // homehost=any option to ensure that the RAID array can be opened by the
         // runtime OS.

src/engine/update.rs

@@ -7,7 +7,10 @@ use tokio::sync::mpsc;
 use osutils::{chroot, container, path::join_relative};
 use trident_api::{
     config::{HostConfiguration, Operations},
-    constants::{internal_params::NO_TRANSITION, ESP_MOUNT_POINT_PATH},
+    constants::{
+        internal_params::{ENABLE_UKI_SUPPORT, NO_TRANSITION},
+        ESP_MOUNT_POINT_PATH,
+    },
     error::{
         InternalError, InvalidInputError, ReportError, ServicingError, TridentError,
         TridentResultExt,
@@ -61,7 +64,7 @@ pub(crate) fn update(
         ab_active_volume: state.host_status().ab_active_volume,
         disk_uuids: state.host_status().disk_uuids.clone(),
         install_index: state.host_status().install_index,
-        is_uki: Some(image.is_uki()),
+        is_uki: Some(image.is_uki() || host_config.internal_params.get_flag(ENABLE_UKI_SUPPORT)),
         image: Some(image),
         storage_graph: engine::build_storage_graph(&host_config.storage)?, // Build storage graph
         filesystems: Vec::new(), // Will be populated after dynamic validation

src/subsystems/esp.rs

@@ -189,7 +189,7 @@ fn copy_file_artifacts(
             esp_dir_path.display()
         ))?;
 
-    if ctx.is_uki_image().unstructured("UKI setting unknown")? {
+    if ctx.is_uki().unstructured("UKI setting unknown")? {
         // Prepare ESP directory structure for UKI boot
         uki::prepare_esp_for_uki(mount_point)?;
 

src/subsystems/initrd.rs

@@ -18,7 +18,7 @@ impl Subsystem for InitrdSubsystem {
 
     #[tracing::instrument(name = "initrd_regeneration", skip_all)]
     fn configure(&mut self, ctx: &EngineContext) -> Result<(), TridentError> {
-        if ctx.is_uki_image()? {
+        if ctx.is_uki()? {
             debug!("Skipping initrd regeneration because UKI is in use");
             return Ok(());
         }

src/subsystems/osconfig/mod.rs

@@ -123,7 +123,7 @@ impl Subsystem for OsConfigSubsystem {
                 self.name()
             );
             return Ok(());
-        } else if ctx.is_uki_image()? && ctx.storage_graph.root_fs_is_verity() {
+        } else if ctx.is_uki()? && ctx.storage_graph.root_fs_is_verity() {
             error!("Skipping OS configuration changes requested in Host Configuration because UKI root verity is in use.");
             return Ok(());
         }
@@ -176,7 +176,7 @@ impl Subsystem for OsConfigSubsystem {
 
         // If we have a UKI image, update SELinux mode here since it cannot be set via kernel
         // command line.
-        if ctx.is_uki_image()? && ctx.spec.os.selinux.mode.is_some() {
+        if ctx.is_uki()? && ctx.spec.os.selinux.mode.is_some() {
             debug!("Updating SELinux config");
             os_modifier_config.selinux = Some(ctx.spec.os.selinux.clone());
         }

src/subsystems/selinux.rs

@@ -124,7 +124,7 @@ impl Subsystem for SelinuxSubsystem {
 
         // If a verity filesystem is mounted at root, ensure that SELinux is not
         // in enforcing mode and warn if it is in permissive mode
-        if ctx.storage_graph.root_fs_is_verity() && !ctx.is_uki_image()? {
+        if ctx.storage_graph.root_fs_is_verity() && !ctx.is_uki()? {
             match final_selinux_mode {
                 SelinuxMode::Enforcing => {
                     return Err(TridentError::new(InvalidInputError::from(

src/subsystems/storage/mod.rs

@@ -183,7 +183,7 @@ impl Subsystem for StorageSubsystem {
             .internal_params
             .get_flag("overridePcrlockEncryption")
             || container::is_running_in_container()?;
-        if ctx.is_uki_image()? {
+        if ctx.is_uki()? {
             if !override_pcrlock_encryption {
                 debug!("Starting step 'Provision' for subunit '{ENCRYPTION_SUBSYSTEM_NAME}'");
                 encryption::provision(ctx, mount_path).message(format!(
@@ -203,7 +203,7 @@ impl Subsystem for StorageSubsystem {
 
     #[tracing::instrument(name = "storage_configuration", skip_all)]
     fn configure(&mut self, ctx: &EngineContext) -> Result<(), TridentError> {
-        if ctx.is_uki_image()? && ctx.storage_graph.root_fs_is_verity() {
+        if ctx.is_uki()? && ctx.storage_graph.root_fs_is_verity() {
             debug!("Skipping storage configuration because UKI root verity is in use");
             return Ok(());
         }