From d01a84925be473c04764ced732c6f2e4c833f08e Mon Sep 17 00:00:00 2001
From: "m.aliarslan" <86300262+Kewe63@users.noreply.github.com>
Date: Sun, 25 Oct 2026 21:33:45 +0300
Subject: [PATCH] fix(upgrade): clean temporary directories on failure

---
 src/commands/upgrade.rs  | 30 +++++++++++++-----
 tests/cli_integration.rs | 66 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+), 8 deletions(-)

diff --git a/src/commands/upgrade.rs b/src/commands/upgrade.rs
index 4740d7f..8e82596 100644
--- a/src/commands/upgrade.rs
+++ b/src/commands/upgrade.rs
@@ -7,6 +7,20 @@ use anyhow::{Context, bail};
 const REPO: &str = "Polymarket/polymarket-cli";
 const BINARY: &str = "polymarket";
 
+struct TemporaryDirectory(String);
+
+impl TemporaryDirectory {
+    fn path(&self) -> &str {
+        &self.0
+    }
+}
+
+impl Drop for TemporaryDirectory {
+    fn drop(&mut self) {
+        let _ = fs::remove_dir_all(&self.0);
+    }
+}
+
 pub fn execute() -> anyhow::Result<()> {
     let current_version = env!("CARGO_PKG_VERSION");
     println!("Current version: v{current_version}");
@@ -30,7 +44,7 @@ pub fn execute() -> anyhow::Result<()> {
     let current_exe = env::current_exe().context("Failed to determine current executable path")?;
 
     let tmpdir = tempdir()?;
-    let tarball = format!("{tmpdir}/{BINARY}.tar.gz");
+    let tarball = format!("{}/{BINARY}.tar.gz", tmpdir.path());
 
     let tarball_name = format!("{BINARY}-{latest_tag}-{target}.tar.gz");
     let checksums_url =
@@ -46,7 +60,7 @@ pub fn execute() -> anyhow::Result<()> {
         bail!("Download failed (HTTP error)");
     }
 
-    let checksums_file = format!("{tmpdir}/checksums.txt");
+    let checksums_file = format!("{}/checksums.txt", tmpdir.path());
     let status = Command::new("curl")
         .args(["-sSfL", "-o", &checksums_file, &checksums_url])
         .status()
@@ -58,14 +72,14 @@ pub fn execute() -> anyhow::Result<()> {
     verify_checksum(&tarball, &checksums_file, &tarball_name)?;
 
     let status = Command::new("tar")
-        .args(["xzf", &tarball, "-C", &tmpdir])
+        .args(["xzf", &tarball, "-C", tmpdir.path()])
         .status()
         .context("Failed to extract archive")?;
     if !status.success() {
         bail!("Failed to extract archive");
     }
 
-    let new_binary = format!("{tmpdir}/{BINARY}");
+    let new_binary = format!("{}/{BINARY}", tmpdir.path());
 
     // Replace the current binary
     let exe_path = current_exe.to_str().context("Non-UTF8 executable path")?;
@@ -90,8 +104,6 @@ pub fn execute() -> anyhow::Result<()> {
     }
 
     let _ = fs::remove_file(&backup);
-    let _ = fs::remove_dir_all(&tmpdir);
-
     println!("Updated to {latest_tag}");
     Ok(())
 }
@@ -132,7 +144,7 @@ fn detect_target() -> anyhow::Result<&'static str> {
     }
 }
 
-fn tempdir() -> anyhow::Result<String> {
+fn tempdir() -> anyhow::Result<TemporaryDirectory> {
     let output = Command::new("mktemp")
         .args(["-d"])
         .output()
@@ -140,7 +152,9 @@ fn tempdir() -> anyhow::Result<String> {
     if !output.status.success() {
         bail!("mktemp failed");
     }
-    Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
+    Ok(TemporaryDirectory(
+        String::from_utf8_lossy(&output.stdout).trim().to_string(),
+    ))
 }
 
 fn verify_checksum(
diff --git a/tests/cli_integration.rs b/tests/cli_integration.rs
index 41d3d11..8f1fc00 100644
--- a/tests/cli_integration.rs
+++ b/tests/cli_integration.rs
@@ -486,3 +486,69 @@ fn wallet_address_succeeds_or_fails_gracefully() {
     // Either succeeds or fails with an error message — not a panic
     assert!(output.status.success() || !output.stderr.is_empty());
 }
+
+#[cfg(unix)]
+#[test]
+fn upgrade_cleans_up_temporary_directory_after_download_failure() {
+    use std::fs;
+    use std::os::unix::fs::PermissionsExt;
+    use std::time::{SystemTime, UNIX_EPOCH};
+
+    struct TestDirectory(std::path::PathBuf);
+
+    impl Drop for TestDirectory {
+        fn drop(&mut self) {
+            let _ = fs::remove_dir_all(&self.0);
+        }
+    }
+
+    let unique = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_nanos();
+    let root = TestDirectory(std::env::temp_dir().join(format!(
+        "polymarket-upgrade-cleanup-{}-{unique}",
+        std::process::id()
+    )));
+    let fake_bin = root.0.join("bin");
+    let upgrade_directory = root.0.join("upgrade-temporary-directory");
+    fs::create_dir_all(&fake_bin).unwrap();
+
+    let fake_curl = fake_bin.join("curl");
+    fs::write(
+        &fake_curl,
+        "#!/bin/sh\ncase \"$*\" in\n  *releases/latest*) printf '%s\\n' '{\"tag_name\":\"v999.0.0\"}'; exit 0 ;;\n  *) exit 22 ;;\nesac\n",
+    )
+    .unwrap();
+    fs::set_permissions(&fake_curl, fs::Permissions::from_mode(0o755)).unwrap();
+
+    let fake_mktemp = fake_bin.join("mktemp");
+    fs::write(
+        &fake_mktemp,
+        format!(
+            "#!/bin/sh\nmkdir -p '{}'\nprintf '%s\\n' '{}'\n",
+            upgrade_directory.display(),
+            upgrade_directory.display()
+        ),
+    )
+    .unwrap();
+    fs::set_permissions(&fake_mktemp, fs::Permissions::from_mode(0o755)).unwrap();
+
+    let path = format!(
+        "{}:{}",
+        fake_bin.display(),
+        std::env::var("PATH").unwrap_or_default()
+    );
+    polymarket()
+        .arg("upgrade")
+        .env("PATH", path)
+        .assert()
+        .failure()
+        .stderr(predicate::str::contains("Download failed"));
+
+    assert!(
+        !upgrade_directory.exists(),
+        "temporary upgrade directory was not removed: {}",
+        upgrade_directory.display()
+    );
+}