fix: prevent database loss when disk is full

Cargo.toml

@@ -56,6 +56,7 @@ enum_dispatch = "0.3.13"
 derive_builder = "0.20.2"
 globset = "0.4.19"
 prost = "0.14.4"
+tempfile = "3.27.0"
 
 [dev-dependencies]
 assert_cmd = "2.2.2"

src/database_file.rs

@@ -0,0 +1,130 @@
+use std::fs::{self, File};
+use std::io::{self, Write};
+use std::path::Path;
+
+pub fn write(path: &Path, content: &[u8]) -> io::Result<()> {
+    replace_with(path, |file| {
+        file.write_all(content)?;
+        file.sync_all()
+    })
+}
+
+fn replace_with(
+    path: &Path,
+    write_and_sync: impl FnOnce(&mut File) -> io::Result<()>,
+) -> io::Result<()> {
+    // Follow existing symlinks, as opening the database directly did before.
+    let path = match path.canonicalize() {
+        Ok(path) => path,
+        Err(e) if e.kind() == io::ErrorKind::NotFound => path.to_path_buf(),
+        Err(e) => return Err(e),
+    };
+    let parent = path
+        .parent()
+        .filter(|parent| !parent.as_os_str().is_empty())
+        .unwrap_or_else(|| Path::new("."));
+    // Keeping the temporary file beside the database allows an atomic replacement.
+    let mut temporary = tempfile::NamedTempFile::new_in(parent)?;
+    match fs::metadata(&path) {
+        Ok(metadata) => temporary
+            .as_file()
+            .set_permissions(metadata.permissions())?,
+        Err(e) if e.kind() == io::ErrorKind::NotFound => {}
+        Err(e) => return Err(e),
+    }
+    write_and_sync(temporary.as_file_mut())?;
+    temporary.persist(&path).map_err(|e| e.error)?;
+    // Persist the directory entry as well on platforms that support directory sync.
+    #[cfg(unix)]
+    File::open(parent)?.sync_all()?;
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn replaces_existing_database() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("db.cotp");
+        fs::write(&path, b"old database with a longer payload").unwrap();
+
+        write(&path, b"new database").unwrap();
+
+        assert_eq!(fs::read(&path).unwrap(), b"new database");
+        assert_eq!(fs::read_dir(dir.path()).unwrap().count(), 1);
+    }
+
+    #[test]
+    fn creates_new_database() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("db.cotp");
+
+        write(&path, b"new database").unwrap();
+
+        assert_eq!(fs::read(&path).unwrap(), b"new database");
+    }
+
+    #[test]
+    fn failed_write_preserves_database_and_removes_temporary_file() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("db.cotp");
+        fs::write(&path, b"original database").unwrap();
+
+        let error = replace_with(&path, |file| {
+            file.write_all(b"partial replacement")?;
+            Err(io::Error::from(io::ErrorKind::StorageFull))
+        })
+        .unwrap_err();
+
+        assert_eq!(error.kind(), io::ErrorKind::StorageFull);
+        assert_eq!(fs::read(&path).unwrap(), b"original database");
+        assert_eq!(fs::read_dir(dir.path()).unwrap().count(), 1);
+    }
+
+    #[test]
+    fn failed_persist_removes_temporary_file() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("db.cotp");
+        fs::create_dir(&path).unwrap();
+
+        assert!(write(&path, b"new database").is_err());
+
+        assert!(path.is_dir());
+        assert_eq!(fs::read_dir(dir.path()).unwrap().count(), 1);
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn preserves_permissions() {
+        use std::os::unix::fs::PermissionsExt;
+
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("db.cotp");
+        fs::write(&path, b"old database").unwrap();
+        fs::set_permissions(&path, fs::Permissions::from_mode(0o640)).unwrap();
+
+        write(&path, b"new database").unwrap();
+
+        assert_eq!(
+            fs::metadata(&path).unwrap().permissions().mode() & 0o777,
+            0o640
+        );
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn updates_symlink_target() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("db.cotp");
+        let link = dir.path().join("link.cotp");
+        fs::write(&path, b"old database").unwrap();
+        std::os::unix::fs::symlink(&path, &link).unwrap();
+
+        write(&link, b"new database").unwrap();
+
+        assert!(fs::symlink_metadata(&link).unwrap().is_symlink());
+        assert_eq!(fs::read(&path).unwrap(), b"new database");
+    }
+}

src/main.rs

@@ -17,6 +17,7 @@ use zeroize::Zeroize;
 mod arguments;
 mod clipboard;
 mod crypto;
+mod database_file;
 mod exporters;
 mod importers;
 mod interface;
@@ -78,8 +79,8 @@ fn main() -> AppResult<()> {
                 println!("Modifications have been persisted");
                 0
             }
-            _ => {
-                eprintln!("An error occurred during database overwriting");
+            Err(e) => {
+                eprintln!("An error occurred during database overwriting: {e}");
                 -1
             }
         }

src/otp/otp_element.rs

@@ -1,6 +1,6 @@
 use color_eyre::eyre::{ErrReport, eyre};
 use derive_builder::Builder;
-use std::{fs::File, io::Write, vec};
+use std::vec;
 
 use crate::crypto::cryptography::{argon_derive_key, encrypt_string_with_key, gen_salt};
 use crate::otp::otp_error::OtpError;
@@ -58,26 +58,18 @@ impl OTPDatabase {
     }
 
     pub fn save(&mut self, key: &Vec<u8>, salt: &[u8]) -> color_eyre::Result<()> {
-        self.needs_modification = false;
         migrate(self)?;
-        match self.overwrite_database_key(key, salt) {
-            Ok(()) => Ok(()),
-            Err(e) => Err(ErrReport::from(e)),
-        }
+        self.overwrite_database_key(key, salt)?;
+        self.needs_modification = false;
+        Ok(())
     }
 
-    fn overwrite_database_key(&self, key: &Vec<u8>, salt: &[u8]) -> Result<(), std::io::Error> {
+    fn overwrite_database_key(&self, key: &Vec<u8>, salt: &[u8]) -> color_eyre::Result<()> {
         let json: &str = &serde_json::to_string(&self)?;
-        let encrypted = encrypt_string_with_key(json, key, salt).unwrap();
-        let mut file = File::create(DATABASE_PATH.get().unwrap())?;
-        match serde_json::to_string(&encrypted) {
-            Ok(content) => {
-                file.write_all(content.as_bytes())?;
-                file.sync_all()?;
-                Ok(())
-            }
-            Err(e) => Err(std::io::Error::from(e)),
-        }
+        let encrypted = encrypt_string_with_key(json, key, salt)?;
+        let content = serde_json::to_vec(&encrypted)?;
+        crate::database_file::write(DATABASE_PATH.get().unwrap(), &content)?;
+        Ok(())
     }
 
     pub fn save_with_pw(&mut self, password: &str) -> color_eyre::Result<(Vec<u8>, [u8; 16])> {

tests/add_integration_tests.rs

@@ -31,11 +31,18 @@ For more information, try '--help'.
     #[test_case("--label" ; "Long subcommand")]
     fn add_with_label_should_work(label_arg: &str) {
         // Arrange / Act
+        let dir = assert_fs::TempDir::new().unwrap();
+        let database = dir.path().join("db.cotp");
+        std::fs::copy(
+            "test_samples/cli_integration_test/empty_database",
+            &database,
+        )
+        .unwrap();
         let mut command = cargo_bin_cmd!("cotp");
         let assertion = command
             .arg("--password-stdin")
             .arg("--database-path")
-            .arg("test_samples/cli_integration_test/empty_database")
+            .arg(&database)
             .arg("add")
             .arg(label_arg)
             .arg("test")
@@ -53,4 +60,56 @@ AA
             .stderr(is_empty())
             .stdout(eq("Modifications have been persisted\n"));
     }
+
+    #[cfg(unix)]
+    #[test]
+    fn failed_database_write_preserves_existing_otps() {
+        let dir = assert_fs::TempDir::new().unwrap();
+        let database = dir.path().join("db.cotp");
+        std::fs::copy(
+            "test_samples/cli_integration_test/empty_database",
+            &database,
+        )
+        .unwrap();
+
+        // Store an OTP first so that the regression checks preservation of real data.
+        cargo_bin_cmd!("cotp")
+            .arg("--password-stdin")
+            .arg("--database-path")
+            .arg(&database)
+            .args(["add", "--label", "existing", "--secret-stdin"])
+            .write_stdin("12345678\nAA\n")
+            .assert()
+            .success();
+        let original = std::fs::read(&database).unwrap();
+
+        // A file-size limit forces a real partial write failure without filling the disk.
+        // Ignoring SIGXFSZ makes the write return an error instead of killing cotp.
+        let mut limited = assert_cmd::Command::new("/bin/sh");
+        limited
+            .args(["-c", "trap '' XFSZ; ulimit -f 1; exec \"$@\"", "sh"])
+            .arg(assert_cmd::cargo::cargo_bin!("cotp"))
+            .arg("--password-stdin")
+            .arg("--database-path")
+            .arg(&database)
+            .args(["add", "--label", &"x".repeat(8192), "--secret-stdin"])
+            .write_stdin("12345678\nAA\n")
+            .assert()
+            .failure()
+            .stderr(predicates::str::contains(
+                "An error occurred during database overwriting:",
+            ));
+
+        assert_eq!(std::fs::read(&database).unwrap(), original);
+        assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 1);
+        cargo_bin_cmd!("cotp")
+            .arg("--password-stdin")
+            .arg("--database-path")
+            .arg(&database)
+            .arg("list")
+            .write_stdin("12345678\n")
+            .assert()
+            .success()
+            .stdout(predicates::str::contains("existing"));
+    }
 }