[Fix] Stability

iota-util/src/atomic_file.rs

@@ -0,0 +1,158 @@
+use std::fs::{self, File, OpenOptions};
+use std::io::{self, Write};
+use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::time::{SystemTime, UNIX_EPOCH};
+
+static TEMPORARY_ID: AtomicU64 = AtomicU64::new(0);
+
+/* Persist small state files without exposing a partially written version after
+ * a crash. Backups give operators a local recovery point for keys and config. */
+pub fn replace(path: &Path, contents: &[u8], backup_limit: usize) -> io::Result<()> {
+    replace_with_mode(path, contents, backup_limit, false)
+}
+
+pub fn replace_private(path: &Path, contents: &[u8], backup_limit: usize) -> io::Result<()> {
+    replace_with_mode(path, contents, backup_limit, true)
+}
+
+fn replace_with_mode(
+    path: &Path,
+    contents: &[u8],
+    backup_limit: usize,
+    private: bool,
+) -> io::Result<()> {
+    let parent = path.parent().ok_or_else(|| {
+        io::Error::new(
+            io::ErrorKind::InvalidInput,
+            "persistent file has no parent directory",
+        )
+    })?;
+    fs::create_dir_all(parent)?;
+
+    if backup_limit > 0 && path.is_file() {
+        create_backup(path, backup_limit)?;
+    }
+
+    let temporary = temporary_path(path)?;
+    let write_result = (|| {
+        let mut file = OpenOptions::new()
+            .write(true)
+            .create_new(true)
+            .open(&temporary)?;
+        set_private_permissions(&temporary, private)?;
+        file.write_all(contents)?;
+        file.sync_all()?;
+        fs::rename(&temporary, path)?;
+        sync_directory(parent)
+    })();
+    if write_result.is_err() {
+        let _ = fs::remove_file(&temporary);
+    }
+    write_result
+}
+
+#[cfg(unix)]
+fn set_private_permissions(path: &Path, private: bool) -> io::Result<()> {
+    if private {
+        use std::os::unix::fs::PermissionsExt;
+        fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
+    }
+    Ok(())
+}
+
+#[cfg(not(unix))]
+fn set_private_permissions(_path: &Path, _private: bool) -> io::Result<()> {
+    Ok(())
+}
+
+fn create_backup(path: &Path, backup_limit: usize) -> io::Result<()> {
+    let parent = path.parent().ok_or_else(|| {
+        io::Error::new(
+            io::ErrorKind::InvalidInput,
+            "persistent file has no parent directory",
+        )
+    })?;
+    let name = path.file_name().ok_or_else(|| {
+        io::Error::new(io::ErrorKind::InvalidInput, "persistent file has no name")
+    })?;
+    let timestamp = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_millis();
+    let id = TEMPORARY_ID.fetch_add(1, Ordering::Relaxed);
+    let backup = parent.join(format!(
+        ".{}.backup-{timestamp}-{id}",
+        name.to_string_lossy()
+    ));
+    fs::copy(path, &backup)?;
+    File::open(&backup)?.sync_all()?;
+    sync_directory(parent)?;
+
+    let prefix = format!(".{}.backup-", name.to_string_lossy());
+    let mut backups = fs::read_dir(parent)?
+        .filter_map(Result::ok)
+        .filter(|entry| entry.file_name().to_string_lossy().starts_with(&prefix))
+        .collect::<Vec<_>>();
+    backups.sort_by_key(|entry| entry.file_name());
+    let obsolete = backups.len().saturating_sub(backup_limit);
+    for entry in backups.into_iter().take(obsolete) {
+        fs::remove_file(entry.path())?;
+    }
+    Ok(())
+}
+
+fn temporary_path(path: &Path) -> io::Result<PathBuf> {
+    let parent = path.parent().ok_or_else(|| {
+        io::Error::new(
+            io::ErrorKind::InvalidInput,
+            "persistent file has no parent directory",
+        )
+    })?;
+    let name = path.file_name().ok_or_else(|| {
+        io::Error::new(io::ErrorKind::InvalidInput, "persistent file has no name")
+    })?;
+    let id = TEMPORARY_ID.fetch_add(1, Ordering::Relaxed);
+    Ok(parent.join(format!(
+        ".{}.{}.{}.tmp",
+        name.to_string_lossy(),
+        std::process::id(),
+        id
+    )))
+}
+
+#[cfg(unix)]
+fn sync_directory(path: &Path) -> io::Result<()> {
+    File::open(path)?.sync_all()
+}
+
+#[cfg(not(unix))]
+fn sync_directory(_path: &Path) -> io::Result<()> {
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::replace;
+
+    #[test]
+    fn replace_preserves_a_previous_version_as_a_backup() {
+        let directory = tempfile::tempdir().unwrap();
+        let path = directory.path().join("state");
+        replace(&path, b"first", 2).unwrap();
+        replace(&path, b"second", 2).unwrap();
+
+        assert_eq!(std::fs::read(&path).unwrap(), b"second");
+        let backups = std::fs::read_dir(directory.path())
+            .unwrap()
+            .filter_map(Result::ok)
+            .filter(|entry| {
+                entry
+                    .file_name()
+                    .to_string_lossy()
+                    .contains(".state.backup-")
+            })
+            .count();
+        assert_eq!(backups, 1);
+    }
+}

iota-util/src/crypto_helper.rs

@@ -6,7 +6,10 @@ pub fn generate_keyring() -> Keyring {
 }
 
 pub fn keyring_to_base64(keyring: &Keyring) -> String {
-    STANDARD.encode(keyring.to_bytes())
+    keyring
+        .try_to_bytes()
+        .map(|bytes| STANDARD.encode(bytes))
+        .unwrap_or_default()
 }
 
 pub fn keyring_from_base64(s: &str) -> Option<Keyring> {
@@ -15,7 +18,10 @@ pub fn keyring_from_base64(s: &str) -> Option<Keyring> {
 }
 
 pub fn public_key_bundle_to_base64(bundle: &PublicKeyBundle) -> String {
-    STANDARD.encode(bundle.as_bytes())
+    bundle
+        .try_as_bytes()
+        .map(|bytes| STANDARD.encode(bytes))
+        .unwrap_or_default()
 }
 
 pub fn public_key_bundle_from_base64(s: &str) -> Option<PublicKeyBundle> {

iota-util/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod atomic_file;
 pub mod crypto_helper;
 pub mod crypto_util;
 pub mod file_util;