From 923cd5917d76e5e222b7d367e9e92fb54d14cdd3 Mon Sep 17 00:00:00 2001 From: AlexKaravaev Date: Tue, 4 Jun 2024 21:22:03 +0200 Subject: [PATCH] orb-slot-ctrl: Prettify handling of arguments from user --- orb-slot-ctrl/src/lib.rs | 106 +++++++++++++++++++++++- orb-slot-ctrl/src/main.rs | 167 +++++++++++++++----------------------- 2 files changed, 168 insertions(+), 105 deletions(-) diff --git a/orb-slot-ctrl/src/lib.rs b/orb-slot-ctrl/src/lib.rs index 330ba3df..f40dd717 100644 --- a/orb-slot-ctrl/src/lib.rs +++ b/orb-slot-ctrl/src/lib.rs @@ -6,6 +6,7 @@ use std::{ fmt, io, path::{Path, PathBuf}, + str::FromStr, }; mod efivar; @@ -46,8 +47,15 @@ pub enum Error { RemoveEfiVar { path: PathBuf, source: io::Error }, #[error("failed reading efivar, invalid data length")] InvalidEfiVarLen, + #[error("invalid slot provided {slot}. Use one of the available slot aliases: \n{help_message}")] + InvalidSlotProvided { slot: String, help_message: String }, #[error("invalid slot configuration")] InvalidSlotData, + #[error("invalid status provided {status}. Use one of the available status variant aliases: \n{help_message}")] + InvalidRootFsStatusProvided { + status: String, + help_message: String, + }, #[error("invalid rootfs status")] InvalidRootFsStatusData, #[error("invalid retry counter({counter}), exceeding the maximum ({max})")] @@ -101,7 +109,7 @@ impl Error { } /// Representation of the slot. -#[derive(Clone, Copy)] +#[derive(Clone, Copy, Debug)] #[repr(u8)] pub enum Slot { /// The Slot A is represented as 0. @@ -110,6 +118,49 @@ pub enum Slot { B = SLOT_B, } +impl Slot { + fn variants() -> Vec<(Slot, &'static str, Vec<&'static str>)> { + vec![ + (Self::A, "A", vec!["a", "0"]), + (Self::B, "B", vec!["b", "1"]), + ] + } + + /// Retrieves a help message listing each slot variant along with its corresponding aliases. + #[must_use] + pub fn help_message() -> String { + let variants = Self::variants(); + let message_parts: Vec = variants + .iter() + .map(|(_, desc, aliases)| format!("{}({})", desc, aliases.join(", "))) + .collect(); + + let message = message_parts.join(";\n"); + + message.to_string() + } +} + +impl FromStr for Slot { + type Err = Error; + + fn from_str(s: &str) -> Result { + Self::variants() + .iter() + .find_map(|(variant, _, aliases)| { + if aliases.contains(&s.to_lowercase().as_str()) { + Some(*variant) + } else { + None + } + }) + .ok_or_else(|| Error::InvalidSlotProvided { + slot: s.to_string(), + help_message: Self::help_message(), + }) + } +} + /// Format slot as lowercase to match Nvidia standard in file system. impl fmt::Display for Slot { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { @@ -121,7 +172,7 @@ impl fmt::Display for Slot { } /// Representation of the rootfs status. -#[derive(Debug)] +#[derive(Clone, Copy, Debug)] #[repr(u8)] pub enum RootFsStatus { /// Default status of the rootfs. @@ -135,6 +186,37 @@ pub enum RootFsStatus { } impl RootFsStatus { + fn variants() -> Vec<(RootFsStatus, &'static str, Vec<&'static str>)> { + vec![ + (Self::Normal, "Normal", vec!["normal", "0"]), + ( + Self::UpdateInProcess, + "Update in Process", + vec!["updateinprocess", "updinprocess", "1"], + ), + ( + Self::UpdateDone, + "Update Done", + vec!["updatedone", "upddone", "2"], + ), + (Self::Unbootable, "Unbootable", vec!["unbootable", "3"]), + ] + } + + /// Retrieves a help message listing each status variant along with its corresponding aliases. + #[must_use] + pub fn help_message() -> String { + let variants = Self::variants(); + let message_parts: Vec = variants + .iter() + .map(|(_, desc, aliases)| format!("{}({})", desc, aliases.join(", "))) + .collect(); + + let message = message_parts.join(";\n"); + + message.to_string() + } + /// Checks if current status is `RootFsStats::Normal`. #[must_use] pub fn is_normal(self) -> bool { @@ -174,6 +256,26 @@ impl TryFrom for RootFsStatus { } } +impl FromStr for RootFsStatus { + type Err = Error; + + fn from_str(s: &str) -> Result { + Self::variants() + .iter() + .find_map(|(variant, _, aliases)| { + if aliases.contains(&s.to_lowercase().as_str()) { + Some(*variant) + } else { + None + } + }) + .ok_or_else(|| Error::InvalidRootFsStatusProvided { + status: s.to_string(), + help_message: Self::help_message(), + }) + } +} + /// Get the current active slot. pub fn get_current_slot() -> Result { match efivar::bootchain::get_current_boot_slot()? { diff --git a/orb-slot-ctrl/src/main.rs b/orb-slot-ctrl/src/main.rs index bbba3e7b..87128149 100644 --- a/orb-slot-ctrl/src/main.rs +++ b/orb-slot-ctrl/src/main.rs @@ -1,4 +1,5 @@ use clap::{Parser, Subcommand}; +use orb_slot_ctrl::{RootFsStatus, Slot}; use std::{env, process::exit}; use orb_build_info::{make_build_info, BuildInfo}; @@ -26,7 +27,13 @@ enum Commands { GetNextSlot, /// Set slot for the next boot. #[command(name = "set", short_flag = 's')] - SetNextSlot { slot: String }, + SetNextSlot { + #[arg( + value_parser = clap::value_parser!(Slot), + help = Slot::help_message() + )] + slot: Slot, + }, /// Rootfs status controls. Status { /// Control the inactive slot instead of the active. @@ -47,7 +54,13 @@ enum StatusCommands { GetRootfsStatus, /// Set the rootfs status. #[command(name = "set", short_flag = 's')] - SetRootfsStatus { status: String }, + SetRootfsStatus { + #[arg( + value_parser = clap::value_parser!(RootFsStatus), + help = RootFsStatus::help_message() + )] + status: RootFsStatus, + }, /// Get the retry counter. #[command(name = "retries", short_flag = 'c')] GetRetryCounter, @@ -82,120 +95,68 @@ fn main() -> eyre::Result<()> { println!("{}", orb_slot_ctrl::get_next_boot_slot()?); } Commands::SetNextSlot { slot } => { - let slot = match slot.as_str() { - // Slot A alias. - "A" => orb_slot_ctrl::Slot::A, - "a" => orb_slot_ctrl::Slot::A, - "0" => orb_slot_ctrl::Slot::A, - // Slot B alias. - "B" => orb_slot_ctrl::Slot::B, - "b" => orb_slot_ctrl::Slot::B, - "1" => orb_slot_ctrl::Slot::B, - _ => { - println!( - "Invalid slot provided, please use either A/a/0 or B/b/1." - ); - exit(1) - } - }; if let Err(e) = orb_slot_ctrl::set_next_boot_slot(slot) { check_running_as_root(e); }; } - Commands::Status { inactive, subcmd } => { - match subcmd { - StatusCommands::GetRootfsStatus => { - if inactive { - println!( - "{:?}", - orb_slot_ctrl::get_rootfs_status( - orb_slot_ctrl::get_inactive_slot()? - )? - ); - } else { - println!("{:?}", orb_slot_ctrl::get_current_rootfs_status()?); - } + Commands::Status { inactive, subcmd } => match subcmd { + StatusCommands::GetRootfsStatus => { + if inactive { + println!( + "{:?}", + orb_slot_ctrl::get_rootfs_status( + orb_slot_ctrl::get_inactive_slot()? + )? + ); + } else { + println!("{:?}", orb_slot_ctrl::get_current_rootfs_status()?); } - StatusCommands::SetRootfsStatus { status } => { - let status = match status.as_str() { - // Status Normal alias. - "Normal" => orb_slot_ctrl::RootFsStatus::Normal, - "normal" => orb_slot_ctrl::RootFsStatus::Normal, - "0" => orb_slot_ctrl::RootFsStatus::Normal, - // Status UpdateInProcess alias. - "UpdateInProcess" => { - orb_slot_ctrl::RootFsStatus::UpdateInProcess - } - "updateinprocess" => { - orb_slot_ctrl::RootFsStatus::UpdateInProcess - } - "updinprocess" => orb_slot_ctrl::RootFsStatus::UpdateInProcess, - "1" => orb_slot_ctrl::RootFsStatus::UpdateInProcess, - // Status UpdateDone alias. - "UpdateDone" => orb_slot_ctrl::RootFsStatus::UpdateDone, - "updatedone" => orb_slot_ctrl::RootFsStatus::UpdateDone, - "upddone" => orb_slot_ctrl::RootFsStatus::UpdateDone, - "2" => orb_slot_ctrl::RootFsStatus::UpdateDone, - // Status Unbootable alias. - "Unbootable" => orb_slot_ctrl::RootFsStatus::Unbootable, - "unbootable" => orb_slot_ctrl::RootFsStatus::Unbootable, - "3" => orb_slot_ctrl::RootFsStatus::Unbootable, - _ => { - println!("Invalid status provided. For a full list of available rootfs status run:"); - println!("slot-ctrl status --list"); - exit(1) - } - }; - if inactive { - if let Err(e) = orb_slot_ctrl::set_rootfs_status( - status, - orb_slot_ctrl::get_inactive_slot()?, - ) { - check_running_as_root(e); - } - } else if let Err(e) = - orb_slot_ctrl::set_current_rootfs_status(status) - { + } + StatusCommands::SetRootfsStatus { status } => { + if inactive { + if let Err(e) = orb_slot_ctrl::set_rootfs_status( + status, + orb_slot_ctrl::get_inactive_slot()?, + ) { check_running_as_root(e); } + } else if let Err(e) = orb_slot_ctrl::set_current_rootfs_status(status) + { + check_running_as_root(e); } - StatusCommands::GetRetryCounter => { - if inactive { - println!( - "{}", - orb_slot_ctrl::get_retry_count( - orb_slot_ctrl::get_inactive_slot()? - )? - ); - } else { - println!("{}", orb_slot_ctrl::get_current_retry_count()?); - } - } - StatusCommands::GetMaxRetryCounter => { - println!("{}", orb_slot_ctrl::get_max_retry_count()?); + } + StatusCommands::GetRetryCounter => { + if inactive { + println!( + "{}", + orb_slot_ctrl::get_retry_count( + orb_slot_ctrl::get_inactive_slot()? + )? + ); + } else { + println!("{}", orb_slot_ctrl::get_current_retry_count()?); } - StatusCommands::ResetRetryCounter => { - if inactive { - if let Err(e) = orb_slot_ctrl::reset_retry_count_to_max( - orb_slot_ctrl::get_inactive_slot()?, - ) { - check_running_as_root(e) - } - } else if let Err(e) = - orb_slot_ctrl::reset_current_retry_count_to_max() - { + } + StatusCommands::GetMaxRetryCounter => { + println!("{}", orb_slot_ctrl::get_max_retry_count()?); + } + StatusCommands::ResetRetryCounter => { + if inactive { + if let Err(e) = orb_slot_ctrl::reset_retry_count_to_max( + orb_slot_ctrl::get_inactive_slot()?, + ) { check_running_as_root(e) } - } - StatusCommands::ListStatusVariants => { - println!("Available Rootfs status variants with their aliases):"); - println!(" Normal (normal, 0)"); - println!(" UpdateInProcess (updateinprocess, updinprocess, 1)"); - println!(" UpdateDone (updatedone, upddone, 2)"); - println!(" Unbootable (unbootable, 3)"); + } else if let Err(e) = orb_slot_ctrl::reset_current_retry_count_to_max() + { + check_running_as_root(e) } } - } + StatusCommands::ListStatusVariants => { + println!("Available Rootfs status variants with their aliases:"); + println!("{}", RootFsStatus::help_message()); + } + }, Commands::GitDescribe => { println!("{}", BUILD_INFO.git.describe); }