🧩 Implement PhinkEnv enum for environment variables

Refactored environment variable management by introducing the `PhinkEnv` enum, simplifying and centralizing the handling of key environment variables. Additionally, moved the error formatting function to a more appropriate module and streamlined various function calls to use the `PhinkEnv` enum.

Signed-off-by: Kevin Valerio <kevin@srlabs.de>

rust-toolchain.toml

@@ -1,5 +1,5 @@
 [toolchain]
-channel = "nightly-2024-08-13"
+channel = "nightly"
 components = ["rustfmt", "cargo", "rustc", "rust-std", "clippy"]
 targets = [
     "aarch64-apple-darwin", # macOS on ARM architecture

src/cli/env.rs

@@ -0,0 +1,52 @@
+use std::fmt::Display;
+
+#[derive(Debug, Clone, Copy)]
+pub enum PhinkEnv {
+    FuzzingWithConfig,
+    FromZiggy,
+    CargoManifestDir,
+    AflForkServerTimeout,
+    AflDebug,
+    AflLLVMAllowList,
+}
+
+impl Display for PhinkEnv {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(match self {
+            PhinkEnv::FuzzingWithConfig => "PHINK_START_FUZZING_WITH_CONFIG",
+            PhinkEnv::FromZiggy => "PHINK_FROM_ZIGGY",
+            PhinkEnv::CargoManifestDir => "CARGO_MANIFEST_DIR",
+            PhinkEnv::AflForkServerTimeout => "AFL_FORKSRV_INIT_TMOUT",
+            PhinkEnv::AflDebug => "AFL_DEBUG",
+            PhinkEnv::AflLLVMAllowList => "AFL_LLVM_ALLOWLIST",
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_phink_env_display() {
+        assert_eq!(
+            PhinkEnv::FuzzingWithConfig.to_string(),
+            "PHINK_START_FUZZING_WITH_CONFIG"
+        );
+        assert_eq!(PhinkEnv::FromZiggy.to_string(), "PHINK_FROM_ZIGGY");
+        assert_eq!(PhinkEnv::CargoManifestDir.to_string(), "CARGO_MANIFEST_DIR");
+    }
+
+    #[test]
+    fn test_phink_env_debug() {
+        assert_eq!(
+            format!("{:?}", PhinkEnv::FuzzingWithConfig),
+            "StartFuzzingWithConfig"
+        );
+        assert_eq!(format!("{:?}", PhinkEnv::FromZiggy), "FromZiggy");
+        assert_eq!(
+            format!("{:?}", PhinkEnv::CargoManifestDir),
+            "CargoManifestDir"
+        );
+    }
+}

src/cli/mod.rs

@@ -1,3 +1,37 @@
+use colored::Colorize;
+use std::backtrace::BacktraceStatus;
+
 pub mod config;
+pub mod env;
 pub mod ui;
 pub mod ziggy;
+
+/// Prints `anyhow` errors correctly in the Phink style
+///
+/// # Argument
+/// * `e`: The `anyhow::Error` to pretty-print
+///
+/// returns: `String`
+///
+/// # Example
+/// `eprintln!("{}", format_error(e));`
+pub fn format_error(e: anyhow::Error) -> String {
+    let mut message = format!("\n{}: {}\n", "Phink got an error...".red().bold(), e);
+
+    if e.backtrace().status() == BacktraceStatus::Captured {
+        message = format!(
+            "{}\n{}\n{}\n",
+            message,
+            "Backtrace ->".yellow(),
+            e.backtrace()
+        )
+    }
+
+    let mut source = e.source();
+    while let Some(cause) = source {
+        message = format!("{}\n\n{}: {}", message, "Caused by".cyan().bold(), cause);
+        source = cause.source();
+    }
+
+    message
+}

src/cli/ziggy.rs

@@ -23,6 +23,13 @@ use crate::cli::{
         },
         PhinkFiles,
     },
+    env::{
+        PhinkEnv,
+        PhinkEnv::{
+            AflDebug,
+            AflForkServerTimeout,
+        },
+    },
     ui,
 };
 use std::{
@@ -34,6 +41,11 @@ use std::{
         Stdio,
     },
 };
+use PhinkEnv::{
+    AflLLVMAllowList,
+    FromZiggy,
+    FuzzingWithConfig,
+};
 
 pub const AFL_DEBUG: &str = "1";
 pub const AFL_FORKSRV_INIT_TMOUT: &str = "10000000";
@@ -70,7 +82,7 @@ impl ZiggyConfig {
     pub fn parse(config_str: String) -> Self {
         let config: Self = serde_json::from_str(&config_str).expect("❌ Failed to parse config");
         if config.config.verbose {
-            println!("🖨️ Using PHINK_START_FUZZING_WITH_CONFIG = {config_str}",);
+            println!("🖨️ Using {} = {config_str}", FuzzingWithConfig);
         }
         config
     }
@@ -102,9 +114,9 @@ impl ZiggyConfig {
         let command_builder = binding
             .arg("ziggy")
             .arg(ziggy_command)
-            .env("PHINK_FROM_ZIGGY", "1")
-            .env("AFL_FORKSRV_INIT_TMOUT", AFL_FORKSRV_INIT_TMOUT)
-            .env("AFL_DEBUG", AFL_DEBUG)
+            .env(FromZiggy.to_string(), "1")
+            .env(AflForkServerTimeout.to_string(), AFL_FORKSRV_INIT_TMOUT)
+            .env(AflDebug.to_string(), AFL_DEBUG)
             .stdout(Stdio::null())
             .stderr(Stdio::null());
 
@@ -132,7 +144,7 @@ impl ZiggyConfig {
         if cfg!(not(target_os = "macos")) {
             let allowlist = PhinkFiles::new(self.config.fuzz_output.clone().unwrap_or_default())
                 .path(AllowlistPath);
-            command_builder.env("AFL_LLVM_ALLOWLIST", allowlist.canonicalize()?);
+            command_builder.env(AflLLVMAllowList.to_string(), allowlist.canonicalize()?);
         }
         Ok(())
     }
@@ -167,10 +179,7 @@ impl ZiggyConfig {
             ))
         }
 
-        let fuzz_config = vec![(
-            "PHINK_START_FUZZING_WITH_CONFIG".to_string(),
-            serde_json::to_string(self)?,
-        )];
+        let fuzz_config = vec![(FuzzingWithConfig.to_string(), serde_json::to_string(self)?)];
 
         self.start(
             ZiggyCommand::Fuzz(self.config.show_ui),
@@ -183,10 +192,7 @@ impl ZiggyConfig {
         self.start(
             ZiggyCommand::Cover,
             vec![],
-            vec![(
-                "PHINK_START_FUZZING_WITH_CONFIG".into(),
-                serde_json::to_string(self)?,
-            )],
+            vec![(FuzzingWithConfig.to_string(), serde_json::to_string(self)?)],
         )?;
         Ok(())
     }
@@ -203,10 +209,7 @@ impl ZiggyConfig {
         self.start(
             ZiggyCommand::Run,
             vec![],
-            vec![(
-                "PHINK_START_FUZZING_WITH_CONFIG".into(),
-                serde_json::to_string(self)?,
-            )],
+            vec![(FuzzingWithConfig.to_string(), serde_json::to_string(self)?)],
         )?;
         Ok(())
     }
@@ -217,7 +220,7 @@ impl ZiggyConfig {
             .path(AllowlistPath);
 
         if allowlist_path.exists() {
-            println!("❗ AFL_LLVM_ALLOWLIST already exists... skipping");
+            println!("❗ {} already exists... skipping", AflLLVMAllowList);
             return Ok(());
         }
 
@@ -229,7 +232,7 @@ impl ZiggyConfig {
             writeln!(allowlist_file, "fun: {}", func)?;
         }
 
-        println!("✅ AFL_LLVM_ALLOWLIST created successfully");
+        println!("✅ {} created successfully", AflLLVMAllowList);
         Ok(())
     }
 }
@@ -353,7 +356,7 @@ mod tests {
                 .collect();
 
             assert!(env_vars.contains(&(
-                "AFL_LLVM_ALLOWLIST".to_string(),
+                AflLLVMAllowList.to_string(),
                 allowlist_path.to_str().unwrap().to_string()
             )));
         }

src/instrumenter/instrumentation.rs

@@ -51,6 +51,7 @@ pub struct InkFilesPath {
     pub wasm_path: PathBuf,
     pub specs_path: PathBuf,
 }
+/// Invokes `println!` only if `verbose` is `true`
 macro_rules! phink_log {
     ($self:expr, $($arg:tt)*) => {
         if $self.to_owned().verbose() {
@@ -68,7 +69,8 @@ impl Instrumenter {
     }
 
     pub fn find(&self) -> anyhow::Result<InkFilesPath> {
-        let wasm_path = fs::read_dir(&self.z_config.contract_path.join("target/ink/"))
+        let c_path = self.to_owned().z_config.contract_path;
+        let wasm_path = fs::read_dir(c_path.join("target/ink/"))
             .with_context(|| {
                 format!(
                     "It seems that your contract is not compiled into `target/ink`. \
@@ -79,7 +81,7 @@ impl Instrumenter {
                         .instrumented_path()
                         .to_str()
                         .unwrap(),
-                    self.z_config.contract_path.to_str().unwrap(),
+                    c_path.to_str().unwrap(),
                 )
             })?
             .filter_map(|entry| {
@@ -111,25 +113,35 @@ impl Instrumenter {
             bail!("There was probably a fork issue, as {p_display} doesn't exist.")
         }
 
-        let status = Command::new("cargo")
-            .current_dir(path.to_owned())
+        let mut build = Command::new("cargo")
+            .current_dir(path.as_path())
             .args(["contract", "build", "--features=phink"])
-            .status()?;
+            .spawn()?;
+
+        phink_log!(self, "✂️ Compiling {}...", p_display);
+
+        let status = build.wait()?;
 
         if !status.success() {
             bail!(
                 "🙅 It seems that your instrumented smart contract did not compile properly. \
-                Please go to {p_display}, edit the source code, and run cargo contract build again\
-                ({status}).",
+        Please go to {p_display}, edit the source code, and run `cargo contract build` again \
+        ({status}).",
             )
         }
+
+        println!(
+            "\n🤞 Contract '{}' has been instrumented and compiled.\n🤞 You can find the instrumented contract in '{p_display}./'",
+            self.z_config.contract_path.display(),
+        );
+
         Ok(())
     }
 
     fn fork(self) -> anyhow::Result<PathBuf> {
         let new_dir = &self.to_owned().z_config.instrumented_path();
 
-        println!("🏗️ Creating new directory: {:?}", new_dir);
+        phink_log!(self, "🏗️ Creating new directory: {:?}", new_dir);
 
         fs::create_dir_all(new_dir)
             .with_context(|| format!("🙅 Failed to create directory: {}", new_dir.display()))?;
@@ -150,7 +162,7 @@ impl Instrumenter {
             );
 
             if entry.path().is_dir() {
-                println!("📂 Creating subdirectory: {:?}", target_path);
+                phink_log!(self, "📂 Creating subdirectory: {:?}", target_path);
                 fs::create_dir_all(&target_path)?;
             } else {
                 phink_log!(
@@ -310,7 +322,6 @@ mod tests {
             self,
             File,
         },
-        io::Read,
     };
     use tempfile::{
         tempdir,

src/instrumenter/instrumented_path.rs

@@ -44,6 +44,7 @@ impl InstrumentedPath {
     }
 
     #[deprecated]
+
     pub fn clean() -> anyhow::Result<()> {
         let dirs_to_remove =
             Self::get_dirs_to_remove(Path::new("/tmp"), DEFAULT_PATH_PATTERN_INSTRUMENTEDPATH)?;