From bd9f7594390a3a07af6e9da50ee68e174b0ecd0e Mon Sep 17 00:00:00 2001
From: NotAShelf <raf@notashelf.dev>
Date: Tue, 22 Jul 2025 02:02:31 +0300
Subject: [PATCH] refactor redundant code

---
 eh/src/build.rs |  86 +--------------------------------------
 eh/src/run.rs   |  86 +--------------------------------------
 eh/src/shell.rs |  93 +-----------------------------------------
 eh/src/util.rs  | 106 ++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 112 insertions(+), 259 deletions(-)

diff --git a/eh/src/build.rs b/eh/src/build.rs
index 016e23e..137b111 100644
--- a/eh/src/build.rs
+++ b/eh/src/build.rs
@@ -1,8 +1,4 @@
-use crate::command::{NixCommand, StdIoInterceptor};
-use crate::util::{HashExtractor, NixErrorClassifier, NixFileFixer};
-use std::io::Write;
-use tracing::{info, warn};
-use yansi::Paint;
+use crate::util::{HashExtractor, NixErrorClassifier, NixFileFixer, handle_nix_with_retry};
 
 pub fn handle_nix_build(
     args: &[String],
@@ -10,83 +6,5 @@ pub fn handle_nix_build(
     fixer: &dyn NixFileFixer,
     classifier: &dyn NixErrorClassifier,
 ) {
-    let mut cmd = NixCommand::new("build").print_build_logs(true);
-    for arg in args {
-        cmd = cmd.arg(arg);
-    }
-    let status = cmd
-        .run_with_logs(StdIoInterceptor)
-        .expect("failed to run nix build");
-    if status.success() {
-        return;
-    }
-
-    let output = NixCommand::new("build")
-        .print_build_logs(true)
-        .args(args.iter().cloned())
-        .output()
-        .expect("failed to capture output");
-    let stderr = String::from_utf8_lossy(&output.stderr);
-
-    if let Some(new_hash) = hash_extractor.extract_hash(&stderr) {
-        if fixer.fix_hash_in_files(&new_hash) {
-            info!("{}", Paint::green("✔ Fixed hash mismatch, retrying..."));
-            let retry_status = NixCommand::new("build")
-                .print_build_logs(true)
-                .args(args.iter().cloned())
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-    }
-
-    if classifier.should_retry(&stderr) {
-        if stderr.contains("unfree") {
-            warn!(
-                "{}",
-                Paint::yellow("⚠ Unfree package detected, retrying with NIXPKGS_ALLOW_UNFREE=1...")
-            );
-            let retry_status = NixCommand::new("build")
-                .print_build_logs(true)
-                .args(args.iter().cloned())
-                .env("NIXPKGS_ALLOW_UNFREE", "1")
-                .impure(true)
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-        if stderr.contains("insecure") {
-            warn!(
-                "{}",
-                Paint::yellow(
-                    "⚠ Insecure package detected, retrying with NIXPKGS_ALLOW_INSECURE=1..."
-                )
-            );
-            let retry_status = NixCommand::new("build")
-                .print_build_logs(true)
-                .args(args.iter().cloned())
-                .env("NIXPKGS_ALLOW_INSECURE", "1")
-                .impure(true)
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-        if stderr.contains("broken") {
-            warn!(
-                "{}",
-                Paint::yellow("⚠ Broken package detected, retrying with NIXPKGS_ALLOW_BROKEN=1...")
-            );
-            let retry_status = NixCommand::new("build")
-                .print_build_logs(true)
-                .args(args.iter().cloned())
-                .env("NIXPKGS_ALLOW_BROKEN", "1")
-                .impure(true)
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-    }
-
-    std::io::stderr().write_all(output.stderr.as_ref()).unwrap();
-    std::process::exit(status.code().unwrap_or(1));
+    handle_nix_with_retry("build", args, hash_extractor, fixer, classifier, false);
 }
diff --git a/eh/src/run.rs b/eh/src/run.rs
index e488c96..dc1b1e7 100644
--- a/eh/src/run.rs
+++ b/eh/src/run.rs
@@ -1,8 +1,4 @@
-use crate::command::{NixCommand, StdIoInterceptor};
-use crate::util::{HashExtractor, NixErrorClassifier, NixFileFixer};
-use std::io::Write;
-use tracing::{info, warn};
-use yansi::Paint;
+use crate::util::{HashExtractor, NixErrorClassifier, NixFileFixer, handle_nix_with_retry};
 
 pub fn handle_nix_run(
     args: &[String],
@@ -10,83 +6,5 @@ pub fn handle_nix_run(
     fixer: &dyn NixFileFixer,
     classifier: &dyn NixErrorClassifier,
 ) {
-    let mut cmd = NixCommand::new("run").print_build_logs(true);
-    for arg in args {
-        cmd = cmd.arg(arg);
-    }
-    let status = cmd
-        .run_with_logs(StdIoInterceptor)
-        .expect("failed to run nix run");
-    if status.success() {
-        return;
-    }
-
-    let output = NixCommand::new("run")
-        .print_build_logs(true)
-        .args(args.iter().cloned())
-        .output()
-        .expect("failed to capture output");
-    let stderr = String::from_utf8_lossy(&output.stderr);
-
-    if let Some(new_hash) = hash_extractor.extract_hash(&stderr) {
-        if fixer.fix_hash_in_files(&new_hash) {
-            info!("{}", Paint::green("✔ Fixed hash mismatch, retrying..."));
-            let retry_status = NixCommand::new("run")
-                .print_build_logs(true)
-                .args(args.iter().cloned())
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-    }
-
-    if classifier.should_retry(&stderr) {
-        if stderr.contains("unfree") {
-            warn!(
-                "{}",
-                Paint::yellow("⚠ Unfree package detected, retrying with NIXPKGS_ALLOW_UNFREE=1...")
-            );
-            let retry_status = NixCommand::new("run")
-                .print_build_logs(true)
-                .args(args.iter().cloned())
-                .env("NIXPKGS_ALLOW_UNFREE", "1")
-                .impure(true)
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-        if stderr.contains("insecure") {
-            warn!(
-                "{}",
-                Paint::yellow(
-                    "⚠ Insecure package detected, retrying with NIXPKGS_ALLOW_INSECURE=1..."
-                )
-            );
-            let retry_status = NixCommand::new("run")
-                .print_build_logs(true)
-                .args(args.iter().cloned())
-                .env("NIXPKGS_ALLOW_INSECURE", "1")
-                .impure(true)
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-        if stderr.contains("broken") {
-            warn!(
-                "{}",
-                Paint::yellow("⚠ Broken package detected, retrying with NIXPKGS_ALLOW_BROKEN=1...")
-            );
-            let retry_status = NixCommand::new("run")
-                .print_build_logs(true)
-                .args(args.iter().cloned())
-                .env("NIXPKGS_ALLOW_BROKEN", "1")
-                .impure(true)
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-    }
-
-    std::io::stderr().write_all(output.stderr.as_ref()).unwrap();
-    std::process::exit(status.code().unwrap_or(1));
+    handle_nix_with_retry("run", args, hash_extractor, fixer, classifier, false);
 }
diff --git a/eh/src/shell.rs b/eh/src/shell.rs
index f6053a3..523049b 100644
--- a/eh/src/shell.rs
+++ b/eh/src/shell.rs
@@ -1,8 +1,4 @@
-use crate::command::{NixCommand, StdIoInterceptor};
-use crate::util::{HashExtractor, NixErrorClassifier, NixFileFixer};
-use std::io::Write;
-use tracing::{info, warn};
-use yansi::Paint;
+use crate::util::{HashExtractor, NixErrorClassifier, NixFileFixer, handle_nix_with_retry};
 
 pub fn handle_nix_shell(
     args: &[String],
@@ -10,90 +6,5 @@ pub fn handle_nix_shell(
     fixer: &dyn NixFileFixer,
     classifier: &dyn NixErrorClassifier,
 ) {
-    let mut cmd = NixCommand::new("shell")
-        .print_build_logs(true)
-        .interactive(true);
-    for arg in args {
-        cmd = cmd.arg(arg);
-    }
-    let status = cmd
-        .run_with_logs(StdIoInterceptor)
-        .expect("failed to run nix shell");
-    if status.success() {
-        std::process::exit(0);
-    }
-
-    // Try to capture error output for retry logic
-    let output = NixCommand::new("shell")
-        .print_build_logs(true)
-        .args(args.iter().cloned())
-        .output()
-        .expect("failed to capture output");
-    let stderr = String::from_utf8_lossy(&output.stderr);
-
-    if let Some(new_hash) = hash_extractor.extract_hash(&stderr) {
-        if fixer.fix_hash_in_files(&new_hash) {
-            info!("{}", Paint::green("✔ Fixed hash mismatch, retrying..."));
-            let retry_status = NixCommand::new("shell")
-                .print_build_logs(true)
-                .interactive(true)
-                .args(args.iter().cloned())
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-    }
-
-    if classifier.should_retry(&stderr) {
-        if stderr.contains("unfree") {
-            warn!(
-                "{}",
-                Paint::yellow("⚠ Unfree package detected, retrying with NIXPKGS_ALLOW_UNFREE=1...")
-            );
-            let retry_status = NixCommand::new("shell")
-                .print_build_logs(true)
-                .interactive(true)
-                .args(args.iter().cloned())
-                .env("NIXPKGS_ALLOW_UNFREE", "1")
-                .impure(true)
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-        if stderr.contains("insecure") {
-            warn!(
-                "{}",
-                Paint::yellow(
-                    "⚠ Insecure package detected, retrying with NIXPKGS_ALLOW_INSECURE=1..."
-                )
-            );
-            let retry_status = NixCommand::new("shell")
-                .print_build_logs(true)
-                .interactive(true)
-                .args(args.iter().cloned())
-                .env("NIXPKGS_ALLOW_INSECURE", "1")
-                .impure(true)
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-        if stderr.contains("broken") {
-            warn!(
-                "{}",
-                Paint::yellow("⚠ Broken package detected, retrying with NIXPKGS_ALLOW_BROKEN=1...")
-            );
-            let retry_status = NixCommand::new("shell")
-                .print_build_logs(true)
-                .interactive(true)
-                .args(args.iter().cloned())
-                .env("NIXPKGS_ALLOW_BROKEN", "1")
-                .impure(true)
-                .run_with_logs(StdIoInterceptor)
-                .unwrap();
-            std::process::exit(retry_status.code().unwrap_or(1));
-        }
-    }
-
-    std::io::stderr().write_all(output.stderr.as_ref()).unwrap();
-    std::process::exit(status.code().unwrap_or(1));
+    handle_nix_with_retry("shell", args, hash_extractor, fixer, classifier, true);
 }
diff --git a/eh/src/util.rs b/eh/src/util.rs
index d45c77c..0cad88b 100644
--- a/eh/src/util.rs
+++ b/eh/src/util.rs
@@ -1,6 +1,10 @@
+use crate::command::{NixCommand, StdIoInterceptor};
 use regex::Regex;
 use std::fs;
+use std::io::Write;
 use std::path::{Path, PathBuf};
+use tracing::{info, warn};
+use yansi::Paint;
 
 pub trait HashExtractor {
     fn extract_hash(&self, stderr: &str) -> Option<String>;
@@ -111,6 +115,108 @@ pub trait NixErrorClassifier {
     fn should_retry(&self, stderr: &str) -> bool;
 }
 
+/// Shared retry logic for nix commands (build/run/shell).
+pub fn handle_nix_with_retry(
+    subcommand: &str,
+    args: &[String],
+    hash_extractor: &dyn HashExtractor,
+    fixer: &dyn NixFileFixer,
+    classifier: &dyn NixErrorClassifier,
+    interactive: bool,
+) -> ! {
+    let mut cmd = NixCommand::new(subcommand).print_build_logs(true);
+    if interactive {
+        cmd = cmd.interactive(true);
+    }
+    for arg in args {
+        cmd = cmd.arg(arg);
+    }
+    let status = cmd
+        .run_with_logs(StdIoInterceptor)
+        .expect("failed to run nix command");
+    if status.success() {
+        std::process::exit(0);
+    }
+
+    let mut output_cmd = NixCommand::new(subcommand)
+        .print_build_logs(true)
+        .args(args.iter().cloned());
+    if interactive {
+        output_cmd = output_cmd.interactive(true);
+    }
+    let output = output_cmd.output().expect("failed to capture output");
+    let stderr = String::from_utf8_lossy(&output.stderr);
+
+    if let Some(new_hash) = hash_extractor.extract_hash(&stderr) {
+        if fixer.fix_hash_in_files(&new_hash) {
+            info!("{}", Paint::green("✔ Fixed hash mismatch, retrying..."));
+            let mut retry_cmd = NixCommand::new(subcommand)
+                .print_build_logs(true)
+                .args(args.iter().cloned());
+            if interactive {
+                retry_cmd = retry_cmd.interactive(true);
+            }
+            let retry_status = retry_cmd.run_with_logs(StdIoInterceptor).unwrap();
+            std::process::exit(retry_status.code().unwrap_or(1));
+        }
+    }
+
+    if classifier.should_retry(&stderr) {
+        if stderr.contains("unfree") {
+            warn!(
+                "{}",
+                Paint::yellow("⚠ Unfree package detected, retrying with NIXPKGS_ALLOW_UNFREE=1...")
+            );
+            let mut retry_cmd = NixCommand::new(subcommand)
+                .print_build_logs(true)
+                .args(args.iter().cloned())
+                .env("NIXPKGS_ALLOW_UNFREE", "1")
+                .impure(true);
+            if interactive {
+                retry_cmd = retry_cmd.interactive(true);
+            }
+            let retry_status = retry_cmd.run_with_logs(StdIoInterceptor).unwrap();
+            std::process::exit(retry_status.code().unwrap_or(1));
+        }
+        if stderr.contains("insecure") {
+            warn!(
+                "{}",
+                Paint::yellow(
+                    "⚠ Insecure package detected, retrying with NIXPKGS_ALLOW_INSECURE=1..."
+                )
+            );
+            let mut retry_cmd = NixCommand::new(subcommand)
+                .print_build_logs(true)
+                .args(args.iter().cloned())
+                .env("NIXPKGS_ALLOW_INSECURE", "1")
+                .impure(true);
+            if interactive {
+                retry_cmd = retry_cmd.interactive(true);
+            }
+            let retry_status = retry_cmd.run_with_logs(StdIoInterceptor).unwrap();
+            std::process::exit(retry_status.code().unwrap_or(1));
+        }
+        if stderr.contains("broken") {
+            warn!(
+                "{}",
+                Paint::yellow("⚠ Broken package detected, retrying with NIXPKGS_ALLOW_BROKEN=1...")
+            );
+            let mut retry_cmd = NixCommand::new(subcommand)
+                .print_build_logs(true)
+                .args(args.iter().cloned())
+                .env("NIXPKGS_ALLOW_BROKEN", "1")
+                .impure(true);
+            if interactive {
+                retry_cmd = retry_cmd.interactive(true);
+            }
+            let retry_status = retry_cmd.run_with_logs(StdIoInterceptor).unwrap();
+            std::process::exit(retry_status.code().unwrap_or(1));
+        }
+    }
+
+    std::io::stderr().write_all(output.stderr.as_ref()).unwrap();
+    std::process::exit(status.code().unwrap_or(1));
+}
 pub struct DefaultNixErrorClassifier;
 
 impl NixErrorClassifier for DefaultNixErrorClassifier {