From 5384ab187f43c8672a1807473bb00ff67d233e0d Mon Sep 17 00:00:00 2001 From: Jo <10510431+j178@users.noreply.github.com> Date: Thu, 21 Nov 2024 21:57:30 +0800 Subject: [PATCH] Stash working tree before running hooks (#96) * Add test * Fix test * add signal test * Fix test * Fix tests * Fix tests --- src/cleanup.rs | 17 +++ src/cli/mod.rs | 2 +- src/cli/run.rs | 41 ++++--- src/git.rs | 17 ++- src/main.rs | 4 + src/run.rs | 260 +++++++++++++++++++++++++++++++++++++++-- tests/common/mod.rs | 22 +--- tests/files/Dockerfile | 0 tests/run.rs | 114 ++++++++++++++++++ 9 files changed, 423 insertions(+), 54 deletions(-) create mode 100644 src/cleanup.rs delete mode 100644 tests/files/Dockerfile diff --git a/src/cleanup.rs b/src/cleanup.rs new file mode 100644 index 0000000..03db74b --- /dev/null +++ b/src/cleanup.rs @@ -0,0 +1,17 @@ +use std::sync::Mutex; + +static CLEANUP_HOOKS: Mutex>> = Mutex::new(Vec::new()); + +/// Run all cleanup functions. +pub fn cleanup() { + let mut cleanup = CLEANUP_HOOKS.lock().unwrap(); + for f in cleanup.drain(..) { + f(); + } +} + +/// Add a cleanup function to be run when the program is interrupted. +pub fn add_cleanup(f: F) { + let mut cleanup = CLEANUP_HOOKS.lock().unwrap(); + cleanup.push(Box::new(f)); +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs index aedff20..1efae69 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -219,7 +219,7 @@ pub(crate) struct RunExtraArgs { pub(crate) remote_branch: Option, #[arg(long, hide = true)] pub(crate) local_branch: Option, - #[arg(long, hide = true)] + #[arg(long, hide = true, required_if_eq("hook_stage", "pre-rebase"))] pub(crate) pre_rebase_upstream: Option, #[arg(long, hide = true)] pub(crate) pre_rebase_branch: Option, diff --git a/src/cli/run.rs b/src/cli/run.rs index 53b1515..e84d67f 100644 --- a/src/cli/run.rs +++ b/src/cli/run.rs @@ -13,11 +13,10 @@ use tracing::{debug, trace}; use crate::cli::{ExitStatus, RunExtraArgs}; use crate::config::Stage; use crate::fs::{normalize_path, Simplified}; -use crate::git::{get_all_files, get_changed_files, get_staged_files, has_unmerged_paths, GIT}; +use crate::git; use crate::hook::{Hook, Project}; use crate::printer::Printer; -use crate::process::Cmd; -use crate::run::{run_hooks, FilenameFilter}; +use crate::run::{run_hooks, FilenameFilter, WorkTreeKeeper}; use crate::store::Store; #[allow(clippy::too_many_arguments)] @@ -34,10 +33,17 @@ pub(crate) async fn run( verbose: bool, printer: Printer, ) -> Result { + // Prevent recursive post-checkout hooks. + if matches!(hook_stage, Some(Stage::PostCheckout)) + && std::env::var_os("_PRE_COMMIT_SKIP_POST_CHECKOUT").is_some() + { + return Ok(ExitStatus::Success); + } + let should_stash = !all_files && files.is_empty(); // Check if we have unresolved merge conflict files and fail fast. - if should_stash && has_unmerged_paths().await? { + if should_stash && git::has_unmerged_paths().await? { writeln!( printer.stderr(), "You have unmerged paths. Resolve them before running pre-commit." @@ -55,21 +61,12 @@ pub(crate) async fn run( return Ok(ExitStatus::Failure); } - // Prevent recursive post-checkout hooks. - if matches!(hook_stage, Some(Stage::PostCheckout)) - && std::env::var_os("_PRE_COMMIT_SKIP_POST_CHECKOUT").is_some() - { - return Ok(ExitStatus::Success); - } - + // Set env vars for hooks. let env_vars = fill_envs(from_ref.as_ref(), to_ref.as_ref(), &extra_args); let mut project = Project::new(config_file)?; let store = Store::from_settings()?.init()?; - // TODO: fill env vars - // TODO: impl staged_files_only - let lock = store.lock_async().await?; let hooks = project.init_hooks(&store, printer).await?; @@ -123,6 +120,12 @@ pub(crate) async fn run( install_hooks(&to_run, printer).await?; drop(lock); + // Clear any unstaged changes from the git working directory. + let mut _guard = None; + if should_stash { + _guard = Some(WorkTreeKeeper::clean(&store).await?); + } + let mut filenames = all_filenames( hook_stage, from_ref, @@ -167,7 +170,7 @@ pub(crate) async fn run( } async fn config_not_staged(config: &Path) -> Result { - let status = Cmd::new(GIT.as_ref()?, "git diff") + let status = git::git_cmd("git diff")? .arg("diff") .arg("--quiet") // Implies --exit-code .arg("--no-ext-diff") @@ -266,7 +269,7 @@ async fn all_filenames( .to_string()]); } if let (Some(from_ref), Some(to_ref)) = (from_ref, to_ref) { - let files = get_changed_files(&from_ref, &to_ref).await?; + let files = git::get_changed_files(&from_ref, &to_ref).await?; debug!( "Files changed between {} and {}: {}", from_ref, @@ -285,7 +288,7 @@ async fn all_filenames( return Ok(files); } if all_files { - let files = get_all_files().await?; + let files = git::get_all_files().await?; debug!("All files in the repo: {}", files.len()); return Ok(files); } @@ -293,7 +296,7 @@ async fn all_filenames( // if is_in_merge_conflict() { // return get_conflicted_files(); // } - let files = get_staged_files().await?; + let files = git::get_staged_files().await?; debug!("Staged files: {}", files.len()); Ok(files) } @@ -304,7 +307,7 @@ async fn install_hook(hook: &Hook, env_dir: PathBuf, printer: Printer) -> Result "Installing environment for {}", hook.repo(), )?; - debug!("Install environment for {} to {}", hook, env_dir.display()); + debug!(%hook, target = %env_dir.display(), "Install environment"); if env_dir.try_exists()? { debug!( diff --git a/src/git.rs b/src/git.rs index 73a1d04..db471c6 100644 --- a/src/git.rs +++ b/src/git.rs @@ -40,7 +40,7 @@ static GIT_ENV: LazyLock> = LazyLock::new(|| { .collect() }); -fn git_cmd(summary: &str) -> Result { +pub fn git_cmd(summary: &str) -> Result { let mut cmd = Cmd::new(GIT.as_ref().map_err(|&e| Error::GitNotFound(e))?, summary); cmd.arg("-c").arg("core.useBuiltinFSMonitor=false"); cmd.envs(GIT_ENV.iter().cloned()); @@ -60,7 +60,20 @@ fn zsplit(s: &[u8]) -> Vec { } } -// TODO: improve error display +pub async fn intent_to_add_files() -> Result, Error> { + let output = git_cmd("get intent to add files")? + .arg("diff") + .arg("--no-ext-diff") + .arg("--ignore-submodules") + .arg("--diff-filter=A") + .arg("--name-only") + .arg("-z") + .check(true) + .output() + .await?; + Ok(zsplit(&output.stdout)) +} + pub async fn get_changed_files(old: &str, new: &str) -> Result, Error> { let output = git_cmd("get changed files")? .arg("diff") diff --git a/src/main.rs b/src/main.rs index a108252..28e6a9f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,10 +11,12 @@ use tracing::{debug, error}; use tracing_subscriber::filter::Directive; use tracing_subscriber::EnvFilter; +use crate::cleanup::cleanup; use crate::cli::{Cli, Command, ExitStatus, SelfCommand, SelfNamespace, SelfUpdateArgs}; use crate::git::get_root; use crate::printer::Printer; +mod cleanup; mod cli; mod config; mod fs; @@ -248,6 +250,8 @@ async fn run(mut cli: Cli) -> Result { fn main() -> ExitCode { ctrlc::set_handler(move || { + cleanup(); + #[allow(clippy::exit, clippy::cast_possible_wrap)] std::process::exit(if cfg!(windows) { 0xC000_013A_u32 as i32 diff --git a/src/run.rs b/src/run.rs index 3006879..3a70c0c 100644 --- a/src/run.rs +++ b/src/run.rs @@ -3,10 +3,11 @@ use std::collections::HashMap; use std::fmt::Write as _; use std::future::Future; use std::io::Write as _; -use std::path::Path; -use std::sync::Arc; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::sync::{Arc, Mutex}; -use anstream::ColorChoice; +use anstream::{eprintln, ColorChoice}; use anyhow::Result; use fancy_regex::{self as regex, Regex}; use owo_colors::{OwoColorize, Style}; @@ -17,12 +18,15 @@ use tokio::task::JoinSet; use tracing::{error, trace}; use unicode_width::UnicodeWidthStr; +use crate::cleanup::add_cleanup; use crate::cli::ExitStatus; -use crate::git::{get_diff, GIT}; +use crate::fs::Simplified; +use crate::git; +use crate::git::{get_diff, git_cmd, GIT}; use crate::hook::Hook; use crate::identify::tags_from_path; use crate::printer::Printer; -use crate::process::Cmd; +use crate::store::Store; const SKIPPED: &str = "Skipped"; const NO_FILES: &str = "(no files to check)"; @@ -164,7 +168,7 @@ pub async fn run_hooks( ColorChoice::Always | ColorChoice::AlwaysAnsi => "--color=always", ColorChoice::Never => "--color=never", }; - Cmd::new(GIT.as_ref()?, "run git diff") + git_cmd("git diff")? .arg("--no-pager") .arg("diff") .arg("--no-ext-diff") @@ -385,15 +389,11 @@ fn partitions<'a>( partitions } -pub async fn run_by_batch( - hook: &Hook, - filenames: &[&String], - run: F, -) -> anyhow::Result> +pub async fn run_by_batch(hook: &Hook, filenames: &[&String], run: F) -> Result> where F: Fn(Vec) -> Fut, F: Clone + Send + Sync + 'static, - Fut: Future> + Send + 'static, + Fut: Future> + Send + 'static, T: Send + 'static, { let mut concurrency = target_concurrency(hook.require_serial); @@ -438,3 +438,239 @@ where Ok(results) } + +static RESTORE_WORKTREE: Mutex> = Mutex::new(None); + +struct IntentToAddKeeper(Vec); +struct WorkingTreeKeeper(Option); + +impl IntentToAddKeeper { + async fn clean() -> Result { + let files = git::intent_to_add_files().await?; + if files.is_empty() { + return Ok(Self(vec![])); + } + + // TODO: xargs + git_cmd("git rm")? + .arg("rm") + .arg("--cached") + .arg("--") + .args(&files) + .check(true) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .await?; + + Ok(Self(files.into_iter().map(PathBuf::from).collect())) + } + + fn restore(&self) -> Result<()> { + // Restore the intent-to-add changes. + if !self.0.is_empty() { + Command::new(GIT.as_ref()?) + .arg("add") + .arg("--intent-to-add") + .arg("--") + // TODO: xargs + .args(&self.0) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status()?; + } + Ok(()) + } +} + +impl Drop for IntentToAddKeeper { + fn drop(&mut self) { + if let Err(err) = self.restore() { + eprintln!( + "{}", + format!("Failed to restore intent-to-add changes: {err}").red() + ); + } + } +} + +impl WorkingTreeKeeper { + async fn clean(patch_dir: &Path) -> Result { + let tree = git::write_tree().await?; + + let mut cmd = git_cmd("git diff-index")?; + let output = cmd + .arg("diff-index") + .arg("--ignore-submodules") + .arg("--binary") + .arg("--exit-code") + .arg("--no-color") + .arg("--no-ext-diff") + .arg(tree) + .arg("--") + .check(false) + .output() + .await?; + + if output.status.success() { + trace!("No non-staged changes detected"); + // No non-staged changes + Ok(Self(None)) + } else if output.status.code() == Some(1) { + if output.stdout.trim_ascii().is_empty() { + trace!("diff-index status code 1 with empty stdout"); + // probably git auto crlf behavior quirks + Ok(Self(None)) + } else { + let now = std::time::SystemTime::now(); + let pid = std::process::id(); + let patch_name = format!( + "{}-{}.patch", + now.duration_since(std::time::UNIX_EPOCH)?.as_millis(), + pid + ); + let patch_path = patch_dir.join(&patch_name); + + eprintln!( + "{}", + format!( + "Non-staged changes detected, saving to `{}`", + patch_path.user_display() + ) + .yellow() + ); + fs_err::create_dir_all(patch_dir)?; + fs_err::write(&patch_path, output.stdout)?; + + // Clean the working tree + Self::checkout_working_tree()?; + + Ok(Self(Some(patch_path))) + } + } else { + Err(cmd.check_status(output.status).unwrap_err().into()) + } + } + + fn checkout_working_tree() -> Result<()> { + let status = Command::new(GIT.as_ref()?) + .arg("-c") + .arg("submodule.recurse=0") + .arg("checkout") + .arg("--") + .arg(".") + // prevent recursive post-checkout hooks + .env("_PRE_COMMIT_SKIP_POST_CHECKOUT", "1") + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status()?; + if status.success() { + Ok(()) + } else { + Err(anyhow::anyhow!("Failed to checkout working tree")) + } + } + + fn git_apply(patch: &Path) -> Result<()> { + let status = Command::new(GIT.as_ref()?) + .arg("apply") + .arg("--whitespace=nowarn") + .arg(patch) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status()?; + if status.success() { + Ok(()) + } else { + Err(anyhow::anyhow!("Failed to apply the patch")) + } + } + + fn restore(&self) -> Result<()> { + let Some(patch) = self.0.as_ref() else { + return Ok(()); + }; + + // Try to apply the patch + if Self::git_apply(patch).is_err() { + error!("Failed to apply the patch, rolling back changes"); + eprintln!( + "{}", + "Failed to apply the patch, rolling back changes".red() + ); + + Self::checkout_working_tree()?; + Self::git_apply(patch)?; + }; + + eprintln!( + "{}", + format!( + "\nRestored working tree changes from `{}`", + patch.user_display() + ) + .yellow() + ); + + Ok(()) + } +} + +impl Drop for WorkingTreeKeeper { + fn drop(&mut self) { + if let Err(err) = self.restore() { + eprintln!( + "{}", + format!("Failed to restore working tree changes: {err}").red() + ); + } + } +} + +/// Clean Git intent-to-add files and working tree changes, and restore them when dropped. +pub struct WorkTreeKeeper { + intent_to_add: Option, + working_tree: Option, +} + +#[derive(Default)] +pub struct RestoreGuard { + _guard: (), +} + +impl Drop for RestoreGuard { + fn drop(&mut self) { + if let Some(mut keeper) = RESTORE_WORKTREE.lock().unwrap().take() { + keeper.restore(); + } + } +} + +impl WorkTreeKeeper { + /// Clear intent-to-add changes from the index and clear the non-staged changes from the working directory. + /// Restore them when the instance is dropped. + pub async fn clean(store: &Store) -> Result { + let cleaner = Self { + intent_to_add: Some(IntentToAddKeeper::clean().await?), + working_tree: Some(WorkingTreeKeeper::clean(store.path()).await?), + }; + + // Set to the global for the cleanup hook. + *RESTORE_WORKTREE.lock().unwrap() = Some(cleaner); + + // Make sure restoration when ctrl-c is pressed. + add_cleanup(|| { + if let Some(guard) = &mut *RESTORE_WORKTREE.lock().unwrap() { + guard.restore(); + } + }); + + Ok(RestoreGuard::default()) + } + + /// Restore the intent-to-add changes and non-staged changes. + fn restore(&mut self) { + self.intent_to_add.take(); + self.working_tree.take(); + } +} diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 550f3c4..a0cccb3 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -179,26 +179,6 @@ impl TestContext { /// Initialize a sample project for pre-commit. pub fn init_project(&self) { - // TODO: Clone some repositories used in the tests. - // if !self.cache_dir.join("pre-commit-hooks") { - // Command::new("git") - // .arg("clone") - // .arg("https://github.com/pre-commit/pre-commit-hooks") - // .arg("--depth=1") - // .arg("--branch=v5.0.0") - // .current_dir(&self.cache_dir) - // .assert() - // .success(); - // } - // - // fs_extra::dir::copy( - // self.cache_dir.join("pre-commit-hooks"), - // self.home_dir.join("pre-commit-hooks"), - // &fs_extra::dir::CopyOptions::new(), - // )?; - - // Write some common files. - Command::new("git") .arg("init") .current_dir(&self.temp_dir) @@ -260,6 +240,8 @@ pub const INSTA_FILTERS: &[(&str, &str)] = &[ r"Caused by: .* \(os error 2\)", "Caused by: No such file or directory (os error 2)", ), + // Time seconds + (r"(\d+\.)?\d+s", "[TIME]"), ]; #[allow(unused_macros)] diff --git a/tests/files/Dockerfile b/tests/files/Dockerfile deleted file mode 100644 index e69de29..0000000 diff --git a/tests/run.rs b/tests/run.rs index 52fa80d..9ca19d5 100644 --- a/tests/run.rs +++ b/tests/run.rs @@ -1,5 +1,6 @@ use anyhow::Result; use assert_fs::prelude::*; +use insta::assert_snapshot; use crate::common::{cmd_snapshot, TestContext}; @@ -52,6 +53,8 @@ fn run_basic() -> Result<()> { ----- stderr ----- "#); + context.git_add("."); + cmd_snapshot!(context.filters(), context.run().arg("trailing-whitespace"), @r#" success: true exit_code: 0 @@ -61,6 +64,8 @@ fn run_basic() -> Result<()> { ----- stderr ----- "#); + context.git_add("."); + cmd_snapshot!(context.filters(), context.run().arg("typos").arg("--hook-stage").arg("pre-push"), @r#" success: false exit_code: 1 @@ -600,3 +605,112 @@ fn pass_env_vars() { let env = context.read("env.txt"); assert_eq!(env, "1\n"); } + +#[test] +fn staged_files_only() -> Result<()> { + let context = TestContext::new(); + context.init_project(); + context.write_pre_commit_config(indoc::indoc! {r#" + repos: + - repo: local + hooks: + - id: trailing-whitespace + name: trailing-whitespace + language: system + entry: python3 -c 'print(open("file.txt", "rt").read())' + verbose: true + types: [text] + "#}); + + context + .workdir() + .child("file.txt") + .write_str("Hello, world!")?; + context.git_add("."); + + // Non-staged files should be stashed and restored. + context + .workdir() + .child("file.txt") + .write_str("Hello world again!")?; + + let filters: Vec<_> = context + .filters() + .into_iter() + .chain([(r"/\d+-\d+.patch", "/[TIME]-[PID].patch")]) + .collect(); + + cmd_snapshot!(filters, context.run(), @r#" + success: true + exit_code: 0 + ----- stdout ----- + trailing-whitespace......................................................Passed + - hook id: trailing-whitespace + - duration: [TIME] + Hello, world! + + ----- stderr ----- + Non-staged changes detected, saving to `[HOME]/[TIME]-[PID].patch` + + Restored working tree changes from `[HOME]/[TIME]-[PID].patch` + "#); + + let content = context.read("file.txt"); + assert_snapshot!(content, @"Hello world again!"); + + Ok(()) +} + +#[cfg(unix)] +#[test] +fn restore_on_interrupt() -> Result<()> { + let context = TestContext::new(); + context.init_project(); + // The hook will sleep for 3 seconds. + context.write_pre_commit_config(indoc::indoc! {r#" + repos: + - repo: local + hooks: + - id: trailing-whitespace + name: trailing-whitespace + language: system + entry: python3 -c 'import time; open("out.txt", "wt").write(open("file.txt", "rt").read()); time.sleep(10)' + verbose: true + types: [text] + "#}); + + context + .workdir() + .child("file.txt") + .write_str("Hello, world!")?; + context.git_add("."); + + // Non-staged files should be stashed and restored. + context + .workdir() + .child("file.txt") + .write_str("Hello world again!")?; + + let mut child = context.run().spawn()?; + let child_id = child.id(); + + // Send an interrupt signal to the process. + let handle = std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_secs(1)); + #[allow(clippy::cast_possible_wrap)] + unsafe { + libc::kill(child_id as i32, libc::SIGINT) + }; + }); + + handle.join().unwrap(); + child.wait()?; + + let content = context.read("out.txt"); + assert_snapshot!(content, @"Hello, world!"); + + let content = context.read("file.txt"); + assert_snapshot!(content, @"Hello world again!"); + + Ok(()) +}