diff --git a/src/handler/mod.rs b/src/handler/mod.rs index b65562b..91e6b87 100644 --- a/src/handler/mod.rs +++ b/src/handler/mod.rs @@ -1,19 +1,17 @@ /// This module contains most of the low-level commands /// and grub variable modifications use anyhow::{bail, Result}; + use std::process::Command; use std::str; /// reboots the system if boot_counter is greater than 0 or can be forced too pub fn handle_reboot(force: bool) -> Result<()> { if !force { - match get_boot_counter() { - Some(t) => { - if t <= 0 { - bail!("countdown ended, check greenboot-rollback status") - } - } - None => bail!("boot_counter is not set"), + if let Some(t) = get_boot_counter()? { + if t <= 0 { + bail!("countdown ended, check greenboot-rollback status") + }; } } log::info!("restarting the system"); @@ -23,26 +21,29 @@ pub fn handle_reboot(force: bool) -> Result<()> { /// rollback to previous ostree deployment if boot counter is less than 0 pub fn handle_rollback() -> Result<()> { - if let Some(t) = get_boot_counter() { - if t <= 0 { + if let Some(boot_counter) = get_boot_counter()? { + if boot_counter <= 0 { log::info!("Greenboot will now attempt to rollback"); Command::new("rpm-ostree").arg("rollback").status()?; return Ok(()); - } else { - bail!("Rollback not initiated, boot_counter {t}"); } } - bail!("Rollback not initiated, boot_counter unset"); + bail!("Rollback not initiated"); } /// sets grub variable boot_counter if not set pub fn set_boot_counter(reboot_count: i32) -> Result<()> { - if let Some(current_counter) = get_boot_counter() { - log::info!("boot_counter={current_counter}"); - return Ok(()); + match get_boot_counter() { + Ok(Some(current_counter)) => { + log::info!("boot_counter={current_counter}"); + bail!("counter already set"); + } + Ok(None) => { + log::info!("setting boot counter"); + set_grub_var("boot_counter", reboot_count)?; + } + Err(e) => bail!(e), } - log::info!("setting boot counter"); - set_grub_var("boot_counter", reboot_count)?; Ok(()) } @@ -77,22 +78,13 @@ pub fn handle_motd(state: &str) -> Result<()> { } /// fetches boot_counter value, none if not set -pub fn get_boot_counter() -> Option { - let grub_vars = Command::new("grub2-editenv").arg("-").arg("list").output(); - if grub_vars.is_err() { - log::error!("{}", grub_vars.unwrap_err()); - return None; - } - let grub_vars = grub_vars.unwrap(); - let grub_vars = match str::from_utf8(&grub_vars.stdout[..]) { - Ok(vars) => vars.lines(), - Err(e) => { - log::error!("{e},unable to fetch grub variables"); - return None; - } - }; - - for var in grub_vars { +pub fn get_boot_counter() -> Result> { + let grub_vars = Command::new("grub2-editenv") + .arg("-") + .arg("list") + .output()?; + let grub_vars = str::from_utf8(&grub_vars.stdout[..])?; + for var in grub_vars.lines() { let (k, v) = if let Some(kv) = var.split_once('=') { kv } else { @@ -101,15 +93,11 @@ pub fn get_boot_counter() -> Option { if k != "boot_counter" { continue; } - match v.parse::() { - Ok(count) => return Some(count), - Err(_) => { - log::error!("boot_counter {v} , should be a valid integer"); - return None; - } - } + + let counter_value = v.parse::()?; + return Ok(Some(counter_value)); } - None + Ok(None) } /// helper function to set any grub variable diff --git a/src/main.rs b/src/main.rs index 4f0a208..cfa7d7f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -318,7 +318,7 @@ mod tests { fn test_boot_counter_set() { unset_boot_counter().ok(); set_boot_counter(10).ok(); - assert_eq!(get_boot_counter(), Some(10)); + assert_eq!(get_boot_counter().unwrap(), Some(10)); unset_boot_counter().ok(); } @@ -327,7 +327,7 @@ mod tests { unset_boot_counter().ok(); set_boot_counter(10).ok(); set_boot_counter(20).ok(); - assert_eq!(get_boot_counter(), Some(10)); + assert_eq!(get_boot_counter().unwrap(), Some(10)); unset_boot_counter().ok(); } @@ -341,7 +341,7 @@ mod tests { .spawn() .context("Cannot create grub variable boot_counter"); set_boot_counter(13).ok(); - assert_eq!(get_boot_counter(), Some(13)); + assert_eq!(get_boot_counter().unwrap(), Some(13)); unset_boot_counter().ok(); }