From 7b4d657a2d0cd95059a13290a6b6c37dbe57fb63 Mon Sep 17 00:00:00 2001
From: Aiden McClelland <me@drbonez.dev>
Date: Mon, 16 May 2022 16:36:44 -0600
Subject: [PATCH] actually address warning instead of muting it like a
 sociopath

---
 backend/src/bin/embassy-init.rs |  4 ++--
 backend/src/setup.rs            | 23 +++++++++++++++++++++--
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/backend/src/bin/embassy-init.rs b/backend/src/bin/embassy-init.rs
index 840de505a..b437833f2 100644
--- a/backend/src/bin/embassy-init.rs
+++ b/backend/src/bin/embassy-init.rs
@@ -79,7 +79,7 @@ async fn setup_or_init(cfg_path: Option<&str>) -> Result<(), Error> {
         let guid_string = tokio::fs::read_to_string("/embassy-os/disk.guid") // unique identifier for volume group - keeps track of the disk that goes with your embassy
             .await?;
         let guid = guid_string.trim();
-        let reboot = embassy::disk::main::import(
+        let requires_reboot = embassy::disk::main::import(
             guid,
             cfg.datadir(),
             if tokio::fs::metadata(REPAIR_DISK_PATH).await.is_ok() {
@@ -95,7 +95,7 @@ async fn setup_or_init(cfg_path: Option<&str>) -> Result<(), Error> {
                 .await
                 .with_ctx(|_| (embassy::ErrorKind::Filesystem, REPAIR_DISK_PATH))?;
         }
-        if reboot.0 {
+        if requires_reboot.0 {
             embassy::disk::main::export(guid, cfg.datadir()).await?;
             Command::new("reboot")
                 .invoke(embassy::ErrorKind::Unknown)
diff --git a/backend/src/setup.rs b/backend/src/setup.rs
index 1a8b67598..8cb1cc3f4 100644
--- a/backend/src/setup.rs
+++ b/backend/src/setup.rs
@@ -35,6 +35,7 @@ use crate::disk::mount::filesystem::cifs::Cifs;
 use crate::disk::mount::filesystem::ReadOnly;
 use crate::disk::mount::guard::TmpMountGuard;
 use crate::disk::util::{pvscan, recovery_info, DiskListResponse, EmbassyOsRecoveryInfo};
+use crate::disk::REPAIR_DISK_PATH;
 use crate::hostname::PRODUCT_KEY_PATH;
 use crate::id::Id;
 use crate::init::init;
@@ -95,13 +96,31 @@ pub async fn attach(
     #[context] ctx: SetupContext,
     #[arg] guid: Arc<String>,
 ) -> Result<SetupResult, Error> {
-    let _ = crate::disk::main::import(
+    let requires_reboot = crate::disk::main::import(
         &*guid,
         &ctx.datadir,
-        RepairStrategy::Preen,
+        if tokio::fs::metadata(REPAIR_DISK_PATH).await.is_ok() {
+            RepairStrategy::Aggressive
+        } else {
+            RepairStrategy::Preen
+        },
         DEFAULT_PASSWORD,
     )
     .await?;
+    if tokio::fs::metadata(REPAIR_DISK_PATH).await.is_ok() {
+        tokio::fs::remove_file(REPAIR_DISK_PATH)
+            .await
+            .with_ctx(|_| (ErrorKind::Filesystem, REPAIR_DISK_PATH))?;
+    }
+    if requires_reboot.0 {
+        crate::disk::main::export(&*guid, &ctx.datadir).await?;
+        return Err(Error::new(
+            eyre!(
+                "Errors were corrected with your disk, but the Embassy must be restarted in order to proceed"
+            ),
+            ErrorKind::DiskManagement,
+        ));
+    }
     let product_key = ctx.product_key().await?;
     let product_key_path = Path::new("/embassy-data/main/product_key.txt");
     if tokio::fs::metadata(product_key_path).await.is_ok() {