From 62c9dad01656dca9d8baaaa0abe6d731faca796c Mon Sep 17 00:00:00 2001
From: lbeder <leonid@lbeder.com>
Date: Wed, 11 Dec 2024 18:06:41 +0400
Subject: [PATCH] Add output file lock

---
 Cargo.lock                  | 11 ++++++++++
 Cargo.toml                  |  1 +
 src/main.rs                 | 11 ++++++++++
 src/utils/file_lock.rs      | 43 +++++++++++++++++++++++++++++++++++++
 src/utils/mod.rs            |  1 +
 src/utils/outputs/output.rs |  4 ----
 6 files changed, 67 insertions(+), 4 deletions(-)
 create mode 100644 src/utils/file_lock.rs

diff --git a/Cargo.lock b/Cargo.lock
index 125a544..46e238b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -663,6 +663,16 @@ dependencies = [
  "percent-encoding",
 ]
 
+[[package]]
+name = "fs2"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213"
+dependencies = [
+ "libc",
+ "winapi",
+]
+
 [[package]]
 name = "futures"
 version = "0.3.31"
@@ -1688,6 +1698,7 @@ dependencies = [
  "criterion",
  "crossterm",
  "dialoguer",
+ "fs2",
  "glob",
  "hex",
  "humantime",
diff --git a/Cargo.toml b/Cargo.toml
index 74e59e8..ddb7467 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,6 +17,7 @@ color-backtrace = "0.6.1"
 criterion = "0.5.1"
 crossterm = "0.27.0"
 dialoguer = "0.11.0"
+fs2 = "0.4.3"
 glob = "0.3.1"
 hex = "0.4.3"
 humantime = "2.1.0"
diff --git a/src/main.rs b/src/main.rs
index 33b58af..03f610d 100755
--- a/src/main.rs
+++ b/src/main.rs
@@ -41,6 +41,7 @@ use utils::{
         version::Version,
     },
     color_hash::color_hash,
+    file_lock::FileLock,
     outputs::{
         fingerprint::Fingerprint,
         output::{OpenOutputOptions, Output, OutputOptions},
@@ -668,8 +669,18 @@ fn derive(derive_options: DeriveOptions) {
     let mut output_key = derive_options.output_key;
     let mut checkpoint: Option<Checkpoint> = None;
 
+    let mut _output_lock: Option<FileLock> = None;
     let mut out: Option<Output> = None;
     if let Some(path) = derive_options.output {
+        if path.exists() {
+            panic!("Output file \"{}\" already exists", path.to_string_lossy());
+        }
+
+        _output_lock = match FileLock::try_lock(&path) {
+            Ok(lock) => Some(lock),
+            Err(_) => panic!("Unable to lock {}", path.to_string_lossy()),
+        };
+
         if output_key.is_none() {
             output_key = Some(get_output_key());
         }
diff --git a/src/utils/file_lock.rs b/src/utils/file_lock.rs
new file mode 100644
index 0000000..bf5fff4
--- /dev/null
+++ b/src/utils/file_lock.rs
@@ -0,0 +1,43 @@
+use fs2::FileExt;
+use std::fs::File;
+use std::io::{self};
+use std::path::Path;
+
+pub struct FileLock {
+    file: File,
+}
+
+impl FileLock {
+    #[allow(dead_code)]
+    pub fn lock<P: AsRef<Path>>(path: P) -> io::Result<Self> {
+        let file = File::options()
+            .read(true)
+            .write(true)
+            .create(true)
+            .truncate(true)
+            .open(&path)?;
+
+        file.lock_exclusive()?;
+
+        Ok(FileLock { file })
+    }
+
+    pub fn try_lock<P: AsRef<Path>>(path: P) -> io::Result<Self> {
+        let file = File::options()
+            .read(true)
+            .write(true)
+            .create(true)
+            .truncate(true)
+            .open(&path)?;
+
+        file.try_lock_exclusive()?;
+
+        Ok(FileLock { file })
+    }
+}
+
+impl Drop for FileLock {
+    fn drop(&mut self) {
+        let _ = self.file.unlock();
+    }
+}
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 3891dde..2e25546 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -2,5 +2,6 @@ pub mod algorithms;
 pub mod chacha20poly1305;
 pub mod checkpoints;
 pub mod color_hash;
+pub mod file_lock;
 pub mod outputs;
 pub mod sodium_init;
diff --git a/src/utils/outputs/output.rs b/src/utils/outputs/output.rs
index b67f1f7..a5d938f 100644
--- a/src/utils/outputs/output.rs
+++ b/src/utils/outputs/output.rs
@@ -121,10 +121,6 @@ pub struct Output {
 
 impl Output {
     pub fn new(opts: &OutputOptions) -> Self {
-        if opts.path.exists() {
-            panic!("Output file \"{}\" already exists", opts.path.to_str().unwrap());
-        }
-
         Self {
             path: opts.path.clone(),
             cipher: ChaCha20Poly1305::new(&opts.key),