DRAFT: feature: enable manual rollback

.pipelines/templates/stages/testing_rollback/testing-template.yml

@@ -43,7 +43,7 @@ parameters:
   - name: skipManualRollbackTesting
     displayName: "Skip manual rollback testing"
     type: string
-    default: "true"
+    default: "false"
 
 jobs:
   - job: RollbackTesting_${{ replace(parameters.flavor, '-', '_') }}
@@ -139,7 +139,7 @@ jobs:
           SKIP_FLAGS=""
 
           if [ "${{ parameters.skipManualRollbackTesting }}" == "true" ]; then
-            # TODO: enable manual rollback testing when it is implemented
+            # skip enable manual rollback testing
             SKIP_FLAGS="$SKIP_FLAGS --skip-manual-rollbacks"
           fi
 

crates/osutils/src/efivar.rs

@@ -15,6 +15,7 @@ const SECURE_BOOT: &str = "SecureBoot";
 const LOADER_ENTRY_ONESHOT: &str = "LoaderEntryOneShot";
 const LOADER_ENTRY_DEFAULT: &str = "LoaderEntryDefault";
 pub const LOADER_ENTRY_SELECTED: &str = "LoaderEntrySelected";
+const LOADER_ENTRIES_DEFAULT: &str = "LoaderEntries";
 
 /// Converts a UTF‑8 Rust string to a UTF-16LE byte array.
 pub fn encode_utf16le(data: &str) -> Vec<u8> {
@@ -42,6 +43,52 @@ fn decode_utf16le(mut data: &[u8]) -> String {
     String::from_utf16_lossy(&utf16_data)
 }
 
+/// Converts a UTF-16LE byte array to a UTF‑8 Rust string.
+fn decode_utf16le_to_strings(data: &[u8]) -> Vec<String> {
+    let mut result = Vec::new();
+    if data.len() <= 2 {
+        return result;
+    }
+
+    let mut start = 0;
+    let u16_null = u16::from_le_bytes([0, 0]);
+
+    // Iterate through the byte slice
+    for (i, &byte) in data.iter().enumerate() {
+        // Combine 2 u8 bytes into a u16
+        if i % 2 == 0 {
+            // Only judge on u16 boundaries
+            continue;
+        }
+        // We are at the second byte of a u16
+        let u16_byte = u16::from_le_bytes([data[i - 1], byte]);
+        if u16_byte == u16_null {
+            // Skip the null-terminating character
+            let end = i - 1;
+            let current_bytes = &data[start..end];
+
+            // If we encounter an empty string (two consecutive nulls, or a null at the very beginning/end)
+            // this usually signifies the end of the list itself.
+            if current_bytes.is_empty() {
+                // Check if this is the final, extra null terminator for the list
+                if i == data.len() - 1 && data.ends_with(b"\0\0") {
+                    break; // End of list found
+                }
+            } else {
+                let utf16_data: Vec<u16> = current_bytes
+                    .chunks(2)
+                    .map(|chunk| u16::from_le_bytes([chunk[0], chunk[1]]))
+                    .collect();
+                let decoded_string = String::from_utf16_lossy(&utf16_data);
+                result.push(decoded_string);
+            }
+            start = i + 1; // Move the start position past the null terminator
+        }
+    }
+
+    result
+}
+
 /// Sets an EFI variable using the efivar command-line tool.
 /// - `name` should include the GUID, e.g. "BootNext-8be4df61-93ca-11d2-aa0d-00e098032b8c"
 /// - `data` should be a hex string, e.g. "0100" for BootNext=0001 (little-endian)
@@ -149,6 +196,32 @@ pub fn set_default_to_current() -> Result<(), TridentError> {
     )
 }
 
+/// Sets the LoaderEntryDefault EFI variable to the previous boot entry
+pub fn set_default_to_previous() -> Result<(), TridentError> {
+    let current = read_efi_variable(BOOTLOADER_INTERFACE_GUID, LOADER_ENTRY_SELECTED)?;
+    let current_decoded = decode_utf16le(&current);
+    let boot_entries = read_efi_variable(BOOTLOADER_INTERFACE_GUID, LOADER_ENTRIES_DEFAULT)?;
+    let boot_entries_decoded = decode_utf16le_to_strings(&boot_entries);
+    if boot_entries_decoded.len() < 2 {
+        return Err(TridentError::new(ServicingError::SetEfiVariable {
+            name: LOADER_ENTRIES_DEFAULT.to_string(),
+        }))
+        .message("Not enough boot entries to determine previous entry");
+    }
+    if boot_entries_decoded[0] != current_decoded {
+        return Err(TridentError::new(ServicingError::SetEfiVariable {
+            name: LOADER_ENTRIES_DEFAULT.to_string(),
+        }))
+        .message("Current boot entry does not match first entry in boot entries list");
+    }
+    let previous = &boot_entries_decoded[1];
+
+    set_efi_variable(
+        &format!("{BOOTLOADER_INTERFACE_GUID}-{LOADER_ENTRY_DEFAULT}"),
+        &encode_utf16le(previous),
+    )
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/trident/src/cli.rs

@@ -39,6 +39,14 @@ pub fn to_operations(allowed_operations: &[AllowedOperation]) -> Operations {
     ops
 }
 
+/// The operations that Trident is allowed to perform
+#[derive(clap::ValueEnum, Copy, Clone, Debug, Eq, PartialEq)]
+pub enum RollbackShowOperation {
+    Validation,
+    Target,
+    Chain,
+}
+
 #[derive(Subcommand, Debug)]
 pub enum Commands {
     /// Initiate an install of Azure Linux
@@ -176,6 +184,33 @@ pub enum Commands {
         history_path: Option<PathBuf>,
     },
 
+    /// Manually rollback to previous state
+    Rollback {
+        /// Declare expectation that rollback undoes a runtime update
+        #[arg(long, conflicts_with = "ab")]
+        runtime: bool,
+
+        /// Declare expectation that rollback undoes an A/B update
+        #[arg(long, conflicts_with = "runtime")]
+        ab: bool,
+
+        /// Comma-separated list of operations that Trident will be allowed to perform
+        #[clap(long, value_delimiter = ',', num_args = 0.., default_value = "stage,finalize")]
+        allowed_operations: Vec<AllowedOperation>,
+
+        /// Show available rollback points
+        #[clap(long)]
+        show: Option<RollbackShowOperation>,
+
+        /// Path to save the resulting Host Status
+        #[clap(short, long)]
+        status: Option<PathBuf>,
+
+        /// Path to save an eventual fatal error
+        #[clap(short, long)]
+        error: Option<PathBuf>,
+    },
+
     #[cfg(feature = "dangerous-options")]
     StreamImage {
         /// URL of the image to stream
@@ -212,6 +247,7 @@ impl Commands {
             Commands::OfflineInitialize { .. } => "offline-initialize",
             #[cfg(feature = "dangerous-options")]
             Commands::StreamImage { .. } => "stream-image",
+            Commands::Rollback { .. } => "rollback",
         }
     }
 }

crates/trident/src/datastore.rs

@@ -2,6 +2,7 @@ use std::{fs, path::Path};
 
 use log::debug;
 
+use sqlite::State;
 use trident_api::{
     error::{
         DatastoreError, InternalError, ReportError, ServicingError, TridentError, TridentResultExt,
@@ -75,6 +76,37 @@ impl DataStore {
         })
     }
 
+    pub(crate) fn get_host_statuses(&self) -> Result<Vec<HostStatus>, TridentError> {
+        let mut all_rows_data: Vec<HostStatus> = Vec::new();
+
+        // Read all HostStatus entries from the datastore, parse them into
+        // HostStatus structs, and return a slice of them.
+        let mut query_statement = self
+            .db
+            .as_ref()
+            .unwrap()
+            .prepare("SELECT contents FROM hoststatus ORDER BY id DESC")
+            .structured(ServicingError::Datastore {
+                inner: DatastoreError::InitializeDatastore,
+            })?;
+
+        while let Ok(State::Row) = query_statement.next() {
+            let host_status_yaml =
+                query_statement
+                    .read::<String, _>(0)
+                    .structured(ServicingError::Datastore {
+                        inner: DatastoreError::InitializeDatastore,
+                    })?;
+            let host_status =
+                serde_yaml::from_str(&host_status_yaml).structured(ServicingError::Datastore {
+                    inner: DatastoreError::InitializeDatastore,
+                })?;
+            all_rows_data.insert(0, host_status);
+        }
+
+        Ok(all_rows_data)
+    }
+
     pub(crate) fn is_persistent(&self) -> bool {
         !self.temporary
     }

crates/trident/src/engine/bootentries.rs

@@ -59,6 +59,11 @@ pub fn create_and_update_boot_variables(
     // Get the label and path for the EFI bootloader of the inactive A/B update volume.
     let (entry_label_new, bootloader_path_new) =
         get_label_and_path(ctx, BOOT_EFI).structured(ServicingError::GetLabelAndPath)?;
+    debug!(
+        "Creating boot entry with label '{}' and bootloader path '{:?}'",
+        entry_label_new.as_str(),
+        bootloader_path_new,
+    );
 
     // Check if the boot entry already exists, if so, delete the entry and
     // remove it from the `BootOrder`.
@@ -166,8 +171,12 @@ pub fn set_boot_next_and_update_boot_order(
             // have boot entries start disappearing again.
             update_boot_order(entry_numbers, &BootOrderPosition::Last)
                 .structured(ServicingError::UpdateBootOrder)?;
-        } else if ctx.servicing_type == ServicingType::CleanInstall && !use_virtdeploy_workaround {
-            // During clean install, immediately set the bootorder to use the new entry.
+        } else if matches!(
+            ctx.servicing_type,
+            ServicingType::CleanInstall | ServicingType::ManualRollback
+        ) && !use_virtdeploy_workaround
+        {
+            // During clean install or manual rollback, immediately set the bootorder to use the new entry.
             update_boot_order(entry_numbers, &BootOrderPosition::First)
                 .structured(ServicingError::UpdateBootOrder)?;
         }

crates/trident/src/engine/context/mod.rs

@@ -85,14 +85,19 @@ impl EngineContext {
             ServicingType::NoActiveServicing => None,
             // If host is executing a runtime update, active and update volumes are the same.
             ServicingType::RuntimeUpdate => self.ab_active_volume,
+
+            // If host is executing a manual rollback and this is executed, an
+            // A/B update is being undone.
+            ServicingType::ManualRollback
             // If host is executing an A/B update, update volume is the opposite of active volume.
-            ServicingType::AbUpdate => {
+            | ServicingType::AbUpdate => {
                 if self.ab_active_volume == Some(AbVolumeSelection::VolumeA) {
                     Some(AbVolumeSelection::VolumeB)
                 } else {
                     Some(AbVolumeSelection::VolumeA)
                 }
             }
+
             // If host is executing a clean install, update volume is always A.
             ServicingType::CleanInstall => Some(AbVolumeSelection::VolumeA),
         }

crates/trident/src/engine/mod.rs

@@ -167,7 +167,7 @@ lazy_static::lazy_static! {
 ///
 /// In case of clean install, the files are persisted to the datastore path in the new root, so
 /// newroot_path is provided.
-fn persist_background_log_and_metrics(
+pub fn persist_background_log_and_metrics(
     datastore_path: &Path,
     newroot_path: Option<&Path>,
     servicing_state: ServicingState,

crates/trident/src/engine/rollback.rs

@@ -48,9 +48,9 @@ pub fn validate_boot(datastore: &mut DataStore) -> Result<BootValidationResult,
     let current_servicing_state = datastore.host_status().servicing_state;
     let ab_active_volume = match current_servicing_state {
         // For *Finalized, use the active volume set in Host Status
-        ServicingState::AbUpdateFinalized | ServicingState::CleanInstallFinalized => {
-            datastore.host_status().ab_active_volume
-        }
+        ServicingState::AbUpdateFinalized
+        | ServicingState::CleanInstallFinalized
+        | ServicingState::ManualRollbackFinalized => datastore.host_status().ab_active_volume,
         // For AbUpdateHealthCheckFailed, use the opposite active volume of the one
         // set in Host Status
         ServicingState::AbUpdateHealthCheckFailed => {
@@ -73,6 +73,7 @@ pub fn validate_boot(datastore: &mut DataStore) -> Result<BootValidationResult,
             ServicingType::AbUpdate
         }
         ServicingState::CleanInstallFinalized => ServicingType::CleanInstall,
+        ServicingState::ManualRollbackFinalized => ServicingType::ManualRollback,
         _ => ServicingType::NoActiveServicing,
     };
 
@@ -117,6 +118,17 @@ pub fn validate_boot(datastore: &mut DataStore) -> Result<BootValidationResult,
                 servicing_type,
             );
         }
+        (true, ServicingState::ManualRollbackFinalized) => {
+            datastore.with_host_status(|host_status| {
+                host_status.servicing_state = ServicingState::Provisioned;
+                host_status.spec_old = Default::default();
+                host_status.ab_active_volume = match host_status.ab_active_volume {
+                    None | Some(AbVolumeSelection::VolumeB) => Some(AbVolumeSelection::VolumeA),
+                    Some(AbVolumeSelection::VolumeA) => Some(AbVolumeSelection::VolumeB),
+                };
+            })?;
+            return Ok(BootValidationResult::ValidBootProvisioned);
+        }
         //
         // Every case below will return an error.
         //

crates/trident/src/engine/storage/rebuild.rs

@@ -176,16 +176,19 @@ pub(crate) fn validate_rebuild_raid(
     match host_status.servicing_state {
         ServicingState::NotProvisioned
         | ServicingState::CleanInstallStaged
-        | ServicingState::CleanInstallFinalized => {
+        | ServicingState::CleanInstallFinalized
+        | ServicingState::AbUpdateHealthCheckFailed
+        | ServicingState::ManualRollbackStaged
+        | ServicingState::ManualRollbackFinalized
+        | ServicingState::RuntimeUpdateStaged => {
             bail!(
                 "rebuild-raid command is not allowed when servicing state is {:?}",
                 host_status.servicing_state
             );
         }
         ServicingState::Provisioned
         | ServicingState::AbUpdateStaged
-        | ServicingState::AbUpdateFinalized
-        | ServicingState::AbUpdateHealthCheckFailed => {}
+        | ServicingState::AbUpdateFinalized => {}
     }
 
     validate_raid_recovery(host_config, disks_to_rebuild)

crates/trident/src/engine/update.rs

@@ -169,6 +169,9 @@ pub(crate) fn update(
         ServicingType::CleanInstall => Err(TridentError::new(
             InvalidInputError::CleanInstallOnProvisionedHost,
         )),
+        ServicingType::ManualRollback => Err(TridentError::internal(
+            "Cannot update during manual rollback",
+        )),
         ServicingType::NoActiveServicing => Err(TridentError::internal("No active servicing type")),
     }
 }