From e96d4b53676bb0ebceb7ca67d2650a03006c3121 Mon Sep 17 00:00:00 2001
From: Leandro Motta Barros <leandro@balena.io>
Date: Wed, 1 May 2024 11:33:15 -0300
Subject: [PATCH] Refactor early checks to a separate module

This is in preparation for a bigger change that will come next.

Signed-off-by: Leandro Motta Barros <leandro@balena.io>
Change-type: patch
---
 src/stage1.rs        | 18 ++++++++++++++----
 src/stage1/checks.rs | 15 +++++++++++++++
 2 files changed, 29 insertions(+), 4 deletions(-)
 create mode 100644 src/stage1/checks.rs

diff --git a/src/stage1.rs b/src/stage1.rs
index 6f173bd..cbb1483 100644
--- a/src/stage1.rs
+++ b/src/stage1.rs
@@ -33,6 +33,7 @@ mod device_impl;
 
 mod exe_copy;
 
+mod checks;
 mod image_retrieval;
 mod utils;
 mod wifi_config;
@@ -45,7 +46,7 @@ use crate::{
             SYSTEM_CONNECTIONS_DIR, SYSTEM_PROXY_DIR, SYS_EFIVARS_DIR, SYS_EFI_DIR, TELINIT_CMD,
         },
         error::{Error, ErrorKind, Result, ToError},
-        file_exists, format_size_with_unit, get_mem_info, get_os_name, is_admin,
+        file_exists, format_size_with_unit, get_mem_info, get_os_name,
         options::Options,
         path_append,
         stage2_config::{Stage2Config, UmountPart},
@@ -63,6 +64,8 @@ use crate::common::stage2_config::LogDevice;
 use crate::common::system::{is_dir, mkdir, stat};
 use mod_logger::{LogDestination, Logger, NO_STREAM};
 
+use self::checks::do_early_checks;
+
 const S1_XTRA_FS_SIZE: u64 = 10 * 1024 * 1024; // const XTRA_MEM_FREE: u64 = 10 * 1024 * 1024; // 10 MB
 
 fn prepare_configs<P1: AsRef<Path>>(
@@ -582,9 +585,16 @@ pub fn stage1(opts: &Options) -> Result<()> {
         }
     };
 
-    if !is_admin()? {
-        error!("please run this program as root");
-        return Err(Error::displayed());
+    match do_early_checks(opts) {
+        Ok(_) => {
+            info!("Early checks passed");
+        }
+        Err(why) => {
+            return Err(Error::from_upstream(
+                Box::new(why),
+                "Failed early checks, exiting",
+            ));
+        }
     }
 
     if !opts.no_ack() {
diff --git a/src/stage1/checks.rs b/src/stage1/checks.rs
new file mode 100644
index 0000000..e07897e
--- /dev/null
+++ b/src/stage1/checks.rs
@@ -0,0 +1,15 @@
+use log::error;
+
+use crate::common::{is_admin, Error, Options, Result};
+
+/// Performs checks to ensure that the program can run properly with the
+/// provided command-line options. Returns an error if the program cannot run
+/// for some reason.
+pub(crate) fn do_early_checks(_opts: &Options) -> Result<()> {
+    if !is_admin()? {
+        error!("please run this program as root");
+        return Err(Error::displayed());
+    }
+
+    Ok(())
+}