Implement audited findings

crates/storage/src/repository.rs

@@ -11,6 +11,8 @@ use std::{
 
 #[cfg(test)]
 use std::collections::BTreeSet;
+#[cfg(test)]
+use std::sync::atomic::{AtomicUsize, Ordering};
 
 pub use crate::secret::SecretBytes;
 use cap_std::{ambient_authority, fs::Dir};
@@ -316,6 +318,8 @@ impl RepositorySnapshot {
 pub struct Repository {
     root_path: PathBuf,
     root: Dir,
+    #[cfg(test)]
+    snapshot_count: AtomicUsize,
 }
 
 impl fmt::Debug for Repository {
@@ -346,7 +350,12 @@ impl Repository {
             .map_err(|error| io_error("canonicalize repository root", requested, error))?;
         let root = Dir::open_ambient_dir(&root_path, ambient_authority())
             .map_err(|error| io_error("open repository root", &root_path, error))?;
-        Ok(Self { root_path, root })
+        Ok(Self {
+            root_path,
+            root,
+            #[cfg(test)]
+            snapshot_count: AtomicUsize::new(0),
+        })
     }
 
     pub fn root_path(&self) -> &Path {
@@ -354,11 +363,18 @@ impl Repository {
     }
 
     pub fn snapshot(&self) -> Result<RepositorySnapshot, RepositoryError> {
+        #[cfg(test)]
+        self.snapshot_count.fetch_add(1, Ordering::Relaxed);
         let mut snapshot = RepositorySnapshot::default();
         scan_directory(&self.root, &DirectoryPath::root(), &mut snapshot)?;
         Ok(snapshot)
     }
 
+    #[cfg(test)]
+    pub(crate) fn snapshot_count(&self) -> usize {
+        self.snapshot_count.load(Ordering::Relaxed)
+    }
+
     pub fn read_entry(&self, path: &EntryPath) -> Result<EncryptedEntry, RepositoryError> {
         let (parent, file_name) = self.open_entry_parent(path)?;
         let metadata = child_metadata(&parent, &file_name, &path.encrypted_relative_path())?
@@ -757,27 +773,41 @@ impl Repository {
     }
 
     fn rollback_created<T>(
+        &self,
+        created: Vec<PathBuf>,
+        original: RepositoryError,
+    ) -> Result<T, RepositoryError> {
+        self.rollback_created_with_checkpoint(created, original, |_| Ok(()))
+    }
+
+    fn rollback_created_with_checkpoint<T>(
         &self,
         mut created: Vec<PathBuf>,
         original: RepositoryError,
+        mut checkpoint: impl FnMut(RollbackStage) -> Result<(), RepositoryError>,
     ) -> Result<T, RepositoryError> {
         while let Some(path) = created.pop() {
             let parent = path.parent().unwrap_or_else(|| Path::new(""));
-            let parent_dir = self.open_directory(parent)?;
-            let name = path
-                .file_name()
-                .expect("created directory path has a file name");
-            match parent_dir.remove_dir(name) {
-                Ok(()) => {
-                    sync_directory(&parent_dir, parent)?;
-                }
-                Err(error) if error.kind() == io::ErrorKind::NotFound => {}
-                Err(error) => {
-                    return Err(RepositoryError::RollbackFailed {
-                        path,
-                        source: error.kind(),
-                    });
+            let rollback = (|| {
+                checkpoint(RollbackStage::BeforeReopen)?;
+                let parent_dir = self.open_directory(parent)?;
+                let name = path
+                    .file_name()
+                    .expect("created directory path has a file name");
+                match parent_dir.remove_dir(name) {
+                    Ok(()) => {
+                        checkpoint(RollbackStage::BeforeSync)?;
+                        sync_directory(&parent_dir, parent)
+                    }
+                    Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
+                    Err(error) => Err(io_error("remove rollback directory", &path, error)),
                 }
+            })();
+            if let Err(rollback) = rollback {
+                return Err(RepositoryError::RollbackFailed {
+                    operation: Box::new(original),
+                    rollback: Box::new(rollback),
+                });
             }
         }
         Err(original)
@@ -816,6 +846,12 @@ enum WriteStage {
     AfterRename,
 }
 
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+enum RollbackStage {
+    BeforeReopen,
+    BeforeSync,
+}
+
 #[cfg(test)]
 impl WriteStage {
     fn fixture_id(self) -> &'static str {
@@ -856,8 +892,8 @@ pub enum RepositoryError {
         source: io::ErrorKind,
     },
     RollbackFailed {
-        path: PathBuf,
-        source: io::ErrorKind,
+        operation: Box<RepositoryError>,
+        rollback: Box<RepositoryError>,
     },
     DurabilityUncertain {
         path: PathBuf,
@@ -919,10 +955,12 @@ impl fmt::Display for RepositoryError {
                 path,
                 source,
             } => write!(formatter, "cannot {operation} {}: {source}", path.display()),
-            Self::RollbackFailed { path, source } => write!(
+            Self::RollbackFailed {
+                operation,
+                rollback,
+            } => write!(
                 formatter,
-                "failed to roll back empty password-store directory {}: {source}",
-                path.display()
+                "repository operation failed ({operation}) and directory rollback failed ({rollback})"
             ),
             Self::DurabilityUncertain { path } => write!(
                 formatter,
@@ -1330,4 +1368,47 @@ mod tests {
         assert!(!temporary.path().join("vault/one").exists());
         Ok(())
     }
+
+    #[test]
+    fn rollback_retains_operation_and_directory_failure() -> Result<(), Box<dyn Error>> {
+        for failed_stage in [RollbackStage::BeforeReopen, RollbackStage::BeforeSync] {
+            let temporary = tempfile::tempdir()?;
+            let vault = temporary.path().join("vault");
+            fs::create_dir_all(vault.join("created"))?;
+            let repository = Repository::open(&vault)?;
+            let operation = RepositoryError::Io {
+                operation: "simulated operation",
+                path: PathBuf::from("entry"),
+                source: io::ErrorKind::Interrupted,
+            };
+            let rollback = RepositoryError::Io {
+                operation: "simulated rollback",
+                path: PathBuf::from("created"),
+                source: io::ErrorKind::Other,
+            };
+
+            let error = repository
+                .rollback_created_with_checkpoint::<()>(
+                    vec![PathBuf::from("created")],
+                    operation.clone(),
+                    |stage| {
+                        if stage == failed_stage {
+                            Err(rollback.clone())
+                        } else {
+                            Ok(())
+                        }
+                    },
+                )
+                .expect_err("rollback checkpoint must fail");
+
+            assert_eq!(
+                error,
+                RepositoryError::RollbackFailed {
+                    operation: Box::new(operation),
+                    rollback: Box::new(rollback),
+                }
+            );
+        }
+        Ok(())
+    }
 }