Add abstraction for atomic file I/O

4 files changed, 81 insertions, 30 deletions

diff --git a/mullvad-api/src/address_cache.rs b/mullvad-api/src/address_cache.rs
index 92b7b6054f..a8b1a91a44 100644
--- a/mullvad-api/src/address_cache.rs
+++ b/mullvad-api/src/address_cache.rs
@@ -83,15 +83,11 @@ impl AddressCache {
             None => return Ok(()),
         };
 
-        let temp_path = write_path.with_file_name("api-cache.temp");
-
-        let mut file = fs::File::create(&temp_path).await?;
+        let mut file = crate::fs::AtomicFile::new(write_path.to_path_buf()).await?;
         let mut contents = address.to_string();
         contents += "\n";
         file.write_all(contents.as_bytes()).await?;
-        file.sync_data().await?;
-
-        fs::rename(&temp_path, write_path).await
+        file.finalize().await
     }
 }
 
diff --git a/mullvad-api/src/fs.rs b/mullvad-api/src/fs.rs
new file mode 100644
index 0000000000..d643563c9a
--- /dev/null
+++ b/mullvad-api/src/fs.rs
@@ -0,0 +1,72 @@
+use std::{
+    ops::{Deref, DerefMut},
+    path::{Path, PathBuf},
+};
+use talpid_types::ErrorExt;
+use tokio::{fs, io};
+
+/// Stores content in a temporary file before moving it to the
+/// final destination, ensuring that consumers of the file never
+/// end up with partial content. Must be moved with `finalize`.
+pub struct AtomicFile {
+    file: Option<fs::File>,
+    temp_path: PathBuf,
+    target_path: PathBuf,
+}
+
+impl AtomicFile {
+    pub async fn new(target_path: PathBuf) -> io::Result<Self> {
+        let temp_path = target_path.with_file_name(uuid::Uuid::new_v4().to_string());
+        Ok(Self {
+            file: Some(fs::File::create(&temp_path).await?),
+            temp_path,
+            target_path,
+        })
+    }
+
+    /// Flushes and moves the file to `self.target_path`, replacing it if it exists.
+    pub async fn finalize(mut self) -> io::Result<()> {
+        let result = async {
+            let file = self.file.take().unwrap();
+            file.sync_all().await?;
+            let std_file = file.into_std().await;
+            let _ = tokio::task::spawn_blocking(move || drop(std_file)).await;
+            fs::rename(&self.temp_path, &self.target_path).await
+        }
+        .await;
+        if result.is_err() {
+            let _ = tokio::task::spawn_blocking(move || try_remove_file(&self.temp_path)).await;
+        }
+        result
+    }
+}
+
+impl Drop for AtomicFile {
+    fn drop(&mut self) {
+        if self.file.is_some() {
+            log::error!("{} was not finalized", self.target_path.display());
+            try_remove_file(&self.temp_path);
+        }
+    }
+}
+
+fn try_remove_file(temp_path: &Path) {
+    if let Err(error) = std::fs::remove_file(temp_path) {
+        let msg = format!("Failed to delete temp file: {}", temp_path.display());
+        log::warn!("{}", error.display_chain_with_msg(&msg));
+    }
+}
+
+impl Deref for AtomicFile {
+    type Target = fs::File;
+
+    fn deref(&self) -> &Self::Target {
+        self.file.as_ref().unwrap()
+    }
+}
+
+impl DerefMut for AtomicFile {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.file.as_mut().unwrap()
+    }
+}
diff --git a/mullvad-api/src/lib.rs b/mullvad-api/src/lib.rs
index 17bbc5b3a2..b1019ef155 100644
--- a/mullvad-api/src/lib.rs
+++ b/mullvad-api/src/lib.rs
@@ -32,6 +32,7 @@ pub use crate::https_client_with_sni::SocketBypassRequest;
 mod access;
 mod address_cache;
 pub mod device;
+mod fs;
 mod relay_list;
 pub use address_cache::AddressCache;
 pub use device::DevicesProxy;
diff --git a/mullvad-api/src/proxy.rs b/mullvad-api/src/proxy.rs
index 009a1960dc..21fa39c9c6 100644
--- a/mullvad-api/src/proxy.rs
+++ b/mullvad-api/src/proxy.rs
@@ -1,7 +1,6 @@
 use crate::tls_stream::TlsStream;
 use futures::Stream;
 use hyper::client::connect::{Connected, Connection};
-use rand::{distributions::Alphanumeric, Rng};
 use serde::{Deserialize, Serialize};
 use shadowsocks::relay::tcprelay::ProxyClientStream;
 use std::{
@@ -90,30 +89,13 @@ impl ApiConnectionMode {
     }
 
     /// Stores this config to `CURRENT_CONFIG_FILENAME`.
-    /// The content is saved to a temporary file first, which ensures that
-    /// consumers of the file never end up with partial content.
     pub async fn save(&self, cache_dir: &Path) -> io::Result<()> {
-        let path = cache_dir.join(CURRENT_CONFIG_FILENAME);
-        let mut temp_ext = String::from("temp");
-        temp_ext.push_str(
-            &rand::thread_rng()
-                .sample_iter(&Alphanumeric)
-                .take(5)
-                .map(char::from)
-                .collect::<String>(),
-        );
-        let temp_path = path.with_extension(temp_ext);
-
-        {
-            let mut file = fs::File::create(&temp_path).await?;
-            let json = serde_json::to_string_pretty(self)
-                .map_err(|_| io::Error::new(io::ErrorKind::Other, "serialization failed"))?;
-            file.write_all(json.as_bytes()).await?;
-            file.write_all(b"\n").await?;
-            file.sync_data().await?;
-        }
-
-        fs::rename(&temp_path, path).await
+        let mut file = crate::fs::AtomicFile::new(cache_dir.join(CURRENT_CONFIG_FILENAME)).await?;
+        let json = serde_json::to_string_pretty(self)
+            .map_err(|_| io::Error::new(io::ErrorKind::Other, "serialization failed"))?;
+        file.write_all(json.as_bytes()).await?;
+        file.write_all(b"\n").await?;
+        file.finalize().await
     }
 
     /// Attempts to remove `CURRENT_CONFIG_FILENAME`, if it exists.