From 7601509e6e860f2a0903ba71df1692d17076ffb1 Mon Sep 17 00:00:00 2001
From: Ruize Tang <1466040111@qq.com>
Date: Mon, 9 Dec 2024 16:56:25 +0800
Subject: [PATCH] Fix OSDK original directory not restored if bundle validation
 fails

---
 osdk/src/bundle/mod.rs         |  6 ++----
 osdk/src/commands/build/mod.rs |  1 +
 osdk/src/util.rs               | 39 ++++++++++++++++++++++++++++++++++
 3 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/osdk/src/bundle/mod.rs b/osdk/src/bundle/mod.rs
index f66a7b268..f88ea15b5 100644
--- a/osdk/src/bundle/mod.rs
+++ b/osdk/src/bundle/mod.rs
@@ -22,6 +22,7 @@ use crate::{
     },
     error::Errno,
     error_msg,
+    util::DirGuard,
 };
 
 /// The osdk bundle artifact that stores as `bundle` directory.
@@ -84,8 +85,7 @@ impl Bundle {
         let manifest_file_content = std::fs::read_to_string(manifest_file_path).ok()?;
         let manifest: BundleManifest = toml::from_str(&manifest_file_content).ok()?;
 
-        let original_dir = std::env::current_dir().unwrap();
-        std::env::set_current_dir(&path).unwrap();
+        let _dir_guard = DirGuard::change_dir(&path);
 
         if let Some(aster_bin) = &manifest.aster_bin {
             if !aster_bin.validate() {
@@ -103,8 +103,6 @@ impl Bundle {
             }
         }
 
-        std::env::set_current_dir(original_dir).unwrap();
-
         Some(Self {
             manifest,
             path: path.as_ref().to_path_buf(),
diff --git a/osdk/src/commands/build/mod.rs b/osdk/src/commands/build/mod.rs
index 6e88e7461..85b887f55 100644
--- a/osdk/src/commands/build/mod.rs
+++ b/osdk/src/commands/build/mod.rs
@@ -239,6 +239,7 @@ fn build_kernel_elf(
     }
 
     info!("Building kernel ELF using command: {:#?}", command);
+    info!("Building directory: {:?}", std::env::current_dir().unwrap());
 
     let status = command.status().unwrap();
     if !status.success() {
diff --git a/osdk/src/util.rs b/osdk/src/util.rs
index 59cc678e1..4b32d2891 100644
--- a/osdk/src/util.rs
+++ b/osdk/src/util.rs
@@ -251,3 +251,42 @@ pub fn trace_panic_from_log(qemu_log: File, bin_path: PathBuf) {
     addr2line_proc.kill().unwrap();
     addr2line_proc.wait().unwrap();
 }
+
+/// A guard that ensures the current working directory is restored
+/// to its original state when the guard goes out of scope.
+pub struct DirGuard(PathBuf);
+
+impl DirGuard {
+    /// Creates a new `DirGuard` that restores the provided directory
+    /// when it goes out of scope.
+    ///
+    /// # Arguments
+    ///
+    /// * `original_dir` - The directory to restore when the guard is dropped.
+    pub fn new(original_dir: PathBuf) -> Self {
+        Self(original_dir)
+    }
+
+    /// Creates a new `DirGuard` using the current working directory as the original directory.
+    pub fn from_current_dir() -> Self {
+        Self::new(std::env::current_dir().unwrap())
+    }
+
+    /// Stores the current directory as the original directory and
+    /// changes the working directory to the specified `new_dir`.
+    ///
+    /// # Arguments
+    ///
+    /// * `new_dir` - The directory to switch to.
+    pub fn change_dir(new_dir: impl AsRef<Path>) -> Self {
+        let original_dir_guard = DirGuard::from_current_dir();
+        std::env::set_current_dir(new_dir.as_ref()).unwrap();
+        original_dir_guard
+    }
+}
+
+impl Drop for DirGuard {
+    fn drop(&mut self) {
+        std::env::set_current_dir(&self.0).unwrap();
+    }
+}