From ce3a2ef9f940309721624533cde4c8bd6d781fa4 Mon Sep 17 00:00:00 2001
From: Sayan Paul <paul.sayan@gmail.com>
Date: Thu, 14 Sep 2023 11:13:07 +0000
Subject: [PATCH] chore: some code refactoring

Signed-off-by: Sayan Paul <paul.sayan@gmail.com>
---
 src/handler/mod.rs | 25 ++++++++++---------------
 src/main.rs        |  4 ++--
 2 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/src/handler/mod.rs b/src/handler/mod.rs
index 08fdea2..f162b8c 100644
--- a/src/handler/mod.rs
+++ b/src/handler/mod.rs
@@ -25,8 +25,8 @@ pub fn handle_reboot(force: bool) -> Result<(), Error> {
 
 /// rollback to previous ostree deployment if boot counter is les than 0
 pub fn handle_rollback() -> Result<(), Error> {
-    match get_boot_counter() {
-        Some(t) if t <= 0 => {
+    if let Some(t) = get_boot_counter() {
+        if t <= 0 {
             log::info!("Greenboot will now attempt rollback");
             let status = Command::new("rpm-ostree").arg("rollback").status()?;
             if status.success() {
@@ -34,26 +34,21 @@ pub fn handle_rollback() -> Result<(), Error> {
             }
             bail!(status.to_string());
         }
-        _ => log::info!("Rollback not initiated as boot_counter is either unset or not equal to 0"),
     }
+    log::info!("Rollback not initiated as boot_counter is either unset or not equal to 0");
     Ok(())
 }
 
 /// sets grub variable boot_counter if not set
 pub fn set_boot_counter(reboot_count: i32) -> Result<()> {
-    match get_boot_counter() {
-        Some(counter) => {
-            log::info!("boot_counter={counter}");
-            Ok(())
-        }
-        None => {
-            if set_grub_var("boot_counter", reboot_count) {
-                log::info!("boot_counter={reboot_count}");
-                return Ok(());
-            }
-            bail!("grub returned error");
-        }
+    if let Some(current_counter) = get_boot_counter() {
+        log::info!("boot_counter={current_counter}");
+        
+    } else if set_grub_var("boot_counter", reboot_count) {
+        log::info!("boot_counter={}", reboot_count);
+        return Ok(());
     }
+    bail!("Failed to set GRUB variable: boot_counter");
 }
 
 /// resets grub variable boot_counter
diff --git a/src/main.rs b/src/main.rs
index 19d948a..41c3388 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -181,7 +181,7 @@ fn run_green() -> Result<(), Error> {
 fn health_check() -> Result<()> {
     let config = GreenbootConfig::get_config();
     log::info!("{config:?}");
-    handle_motd("healthcheck is in progress").ok();
+    handle_motd("healthcheck is in progress")?;
     let run_status = run_diagnostics();
     match run_status {
         Ok(()) => {
@@ -205,7 +205,7 @@ fn health_check() -> Result<()> {
                 .unwrap_or_else(|e| log::error!("cannot set boot_counter as: {}", e.to_string()));
             handle_reboot(false)
                 .unwrap_or_else(|e| log::error!("cannot reboot as: {}", e.to_string()));
-            bail!(e);
+            Err(e)
         }
     }
 }