Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 14 commits into from
Sep 7, 2023
Merged
11 changes: 11 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 2 additions & 1 deletion uplink/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
71 changes: 47 additions & 24 deletions uplink/src/collector/downloader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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() => {
Expand All @@ -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
Expand All @@ -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();

Expand All @@ -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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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(),
Expand Down
Loading