From 604ad8fadeb524d8be346b6b494b986b530daea7 Mon Sep 17 00:00:00 2001
From: Dan Vail <46790194+thedanvail@users.noreply.github.com>
Date: Tue, 7 Jun 2022 01:11:15 -0700
Subject: [PATCH] Adds a more graceful exit when the directory does not exist.
 Fix #53 (#111)

Adds a more graceful exit when the directory does not exist.
Decided to utilize `anyhow` at the recommendation of @aviramha - went with the `with_context` option as that seemed the most elegant solution.
---
 CHANGELOG.md            |  1 +
 mirrord-cli/src/main.rs | 18 ++++++++++--------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 14a70ef1dc2..db094f19a8b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ Check [Keep a Changelog](http://keepachangelog.com/) for recommendations on how
 
 ## [Unreleased]
 ### Changed
+- Added graceful exit for library extraction logic in case of error
 - Refactor the CI by splitting the building of mirrord-agent in a separate job and caching the agent image for E2E tests.
 
 ### Fixed
diff --git a/mirrord-cli/src/main.rs b/mirrord-cli/src/main.rs
index 2459fad8328..38d853416bf 100644
--- a/mirrord-cli/src/main.rs
+++ b/mirrord-cli/src/main.rs
@@ -1,6 +1,6 @@
 use std::{env::temp_dir, fs::File, io::Write, time::Duration};
 
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, Context, Result};
 use clap::{Args, Parser, Subcommand};
 use exec::execvp;
 use semver::Version;
@@ -60,7 +60,7 @@ const INJECTION_ENV_VAR: &str = "LD_PRELOAD";
 #[cfg(target_os = "macos")]
 const INJECTION_ENV_VAR: &str = "DYLD_INSERT_LIBRARIES";
 
-fn extract_library(dest_dir: Option<String>) -> String {
+fn extract_library(dest_dir: Option<String>) -> Result<String> {
     let library_file = env!("CARGO_CDYLIB_FILE_MIRRORD_LAYER");
     let library_path = std::path::Path::new(library_file);
     let file_name = library_path.components().last().unwrap();
@@ -68,11 +68,12 @@ fn extract_library(dest_dir: Option<String>) -> String {
         Some(dest_dir) => std::path::Path::new(&dest_dir).join(file_name),
         None => temp_dir().as_path().join(file_name),
     };
-    let mut file = File::create(file_path.clone()).unwrap();
+    let mut file = File::create(&file_path)
+        .with_context(|| format!("Path \"{}\" creation failed", file_path.display()))?;
     let bytes = include_bytes!(env!("CARGO_CDYLIB_FILE_MIRRORD_LAYER"));
     file.write_all(bytes).unwrap();
     debug!("Extracted library file to {:?}", &file_path);
-    file_path.to_str().unwrap().to_string()
+    Ok(file_path.to_str().unwrap().to_string())
 }
 
 fn add_to_preload(path: &str) -> Result<()> {
@@ -117,7 +118,7 @@ fn exec(args: &ExecArgs) -> Result<()> {
     if args.accept_invalid_certificates {
         std::env::set_var("MIRRORD_ACCEPT_INVALID_CERTIFICATES", "true");
     }
-    let library_path = extract_library(None);
+    let library_path = extract_library(None)?;
     add_to_preload(&library_path).unwrap();
     let mut binary_args = args.binary_args.clone();
     binary_args.insert(0, args.binary.clone());
@@ -127,7 +128,7 @@ fn exec(args: &ExecArgs) -> Result<()> {
 }
 
 const CURRENT_VERSION: &str = env!("CARGO_PKG_VERSION");
-fn main() {
+fn main() -> Result<()> {
     registry()
         .with(fmt::layer())
         .with(EnvFilter::from_default_env())
@@ -136,11 +137,12 @@ fn main() {
 
     let cli = Cli::parse();
     match cli.commands {
-        Commands::Exec(args) => exec(&args).unwrap(),
+        Commands::Exec(args) => exec(&args)?,
         Commands::Extract { path } => {
-            extract_library(Some(path));
+            extract_library(Some(path))?;
         }
     }
+    Ok(())
 }
 
 fn prompt_outdated_version() {