feat: ensure disk has enough space to download file into (#268)

* feat: check disk size before downloading

* fix: use `fs2` to simplify code

* ci: fix broken test due to `check_disk_size`

* fix: ensure directory exists before disk check

* test: run disk size check in test

uplink/Cargo.toml

@@ -41,9 +41,10 @@ tunshell-client = { git = "https://github.com/bytebeamio/tunshell.git", branch =
 fake = { version = "2.5.0", features = ["derive"] }
 rand = "0.8"
 # downloader
+fs2 = "0.4"
 futures-util = "0.3"
-reqwest = { version = "0.11", default-features = false, features = ["stream", "rustls-tls"] }
 human_bytes = "0.4"
+reqwest = { version = "0.11", default-features = false, features = ["stream", "rustls-tls"] }
 # systemstats
 sysinfo = "0.26"
 # logcat

uplink/src/collector/downloader.rs

@@ -85,6 +85,8 @@ pub enum Error {
     FilePathMissing,
     #[error("Download failed, content length zero")]
     EmptyFile,
+    #[error("Disk space is insufficient: {0}")]
+    InsufficientDisk(String),
 }
 
 /// This struct contains the necessary components to download and store file as notified by a download file
@@ -208,6 +210,19 @@ impl FileDownloader {
         let status = status.set_sequence(self.sequence());
         self.bridge_tx.send_action_response(status).await;
 
+        // Ensure that directory for downloading file into, exists
+        let mut download_path = self.config.path.clone();
+        download_path.push(&action.name);
+
+        #[cfg(unix)]
+        self.create_dirs_with_perms(
+            download_path.as_path(),
+            std::os::unix::fs::PermissionsExt::from_mode(0o777),
+        )?;
+
+        #[cfg(not(unix))]
+        std::fs::create_dir_all(&download_path)?;
+
         // Extract url information from action payload
         let mut update = match serde_json::from_str::<DownloadFile>(&action.payload)? {
             DownloadFile { file_name, .. } if file_name.is_empty() => {
@@ -217,10 +232,12 @@ impl FileDownloader {
             u => u,
         };
 
+        self.check_disk_size(&update)?;
+
         let url = update.url.clone();
 
         // Create file to actually download into
-        let (file, file_path) = self.create_file(&action.name, &update.file_name)?;
+        let (file, file_path) = self.create_file(&download_path, &update.file_name)?;
 
         // Create handler to perform download from URL
         // TODO: Error out for 1XX/3XX responses
@@ -239,7 +256,22 @@ impl FileDownloader {
         Ok(())
     }
 
+    fn check_disk_size(&mut self, download: &DownloadFile) -> Result<(), Error> {
+        let disk_free_space = fs2::free_space(&self.config.path)? as usize;
+
+        let req_size = human_bytes(download.content_length as f64);
+        let free_size = human_bytes(disk_free_space as f64);
+        debug!("Download requires {req_size}; Disk free space is {free_size}");
+
+        if download.content_length > disk_free_space {
+            return Err(Error::InsufficientDisk(free_size));
+        }
+
+        Ok(())
+    }
+
     #[cfg(unix)]
+    /// Custom create_dir_all which sets permissions on each created directory, only works on unix
     fn create_dirs_with_perms(&self, path: &Path, perms: Permissions) -> std::io::Result<()> {
         let mut current_path = PathBuf::new();
 
@@ -256,21 +288,11 @@ impl FileDownloader {
     }
 
     /// Creates file to download into
-    fn create_file(&self, name: &str, file_name: &str) -> Result<(File, PathBuf), Error> {
-        // Ensure that directory for downloading file into, exists
-        let mut download_path = self.config.path.clone();
-        download_path.push(name);
-        // do manual create_dir_all while setting permissions on each created directory
-
-        #[cfg(unix)]
-        self.create_dirs_with_perms(
-            download_path.as_path(),
-            std::os::unix::fs::PermissionsExt::from_mode(0o777),
-        )?;
-
-        #[cfg(not(unix))]
-        std::fs::create_dir_all(&download_path)?;
-
+    fn create_file(
+        &self,
+        download_path: &PathBuf,
+        file_name: &str,
+    ) -> Result<(File, PathBuf), Error> {
         let mut file_path = download_path.to_owned();
         file_path.push(file_name);
         // NOTE: if file_path is occupied by a directory due to previous working of uplink, remove it
@@ -309,6 +331,7 @@ impl FileDownloader {
 
             // Calculate percentage on the basis of content_length
             let percentage = 99 * downloaded / content_length;
+
             trace!(
                 "Downloading: size = {size}, percentage = {percentage}, elapsed = {}s",
                 start.elapsed().as_secs()
@@ -411,10 +434,10 @@ mod test {
             download_path: None,
         };
         let mut expected_forward = download_update.clone();
-        let mut download_path = downloader_cfg.path;
-        download_path.push("firmware_update");
-        download_path.push("test.txt");
-        expected_forward.download_path = Some(download_path);
+        let mut path = downloader_cfg.path;
+        path.push("firmware_update");
+        path.push("test.txt");
+        expected_forward.download_path = Some(path);
         let download_action = Action {
             action_id: "1".to_string(),
             kind: "firmware_update".to_string(),
@@ -487,10 +510,10 @@ mod test {
             download_path: None,
         };
         let mut expected_forward = download_update.clone();
-        let mut download_path = downloader_cfg.path;
-        download_path.push("firmware_update");
-        download_path.push("test.txt");
-        expected_forward.download_path = Some(download_path);
+        let mut path = downloader_cfg.path;
+        path.push("firmware_update");
+        path.push("test.txt");
+        expected_forward.download_path = Some(path);
         let download_action = Action {
             action_id: "1".to_string(),
             kind: "firmware_update".to_string(),