diff --git a/src/main.rs b/src/main.rs index dec8992..4fd38cf 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,7 @@ use std::{ cmp::{max, min}, collections::{HashMap, HashSet}, io::{stdout, Write}, + os::unix::fs::FileTypeExt, path::{Path, PathBuf}, str::FromStr, sync::Arc, @@ -331,11 +332,27 @@ async fn main() -> anyhow::Result<()> { }); } - let file_size = target_file.file.metadata().unwrap().len(); + let is_blockdevice = std::fs::metadata(target_file.file.as_path()) + .unwrap() + .file_type() + .is_block_device(); + + let file_size = if is_blockdevice { + let sectors: u64 = std::fs::read_to_string(format!( + "/sys/block/{}/size", + target_file.file.file_name().unwrap().to_str().unwrap() + )) + .unwrap() + .parse() + .unwrap(); + sectors.checked_mul(512).unwrap() + } else { + target_file.file.metadata().unwrap().len() + }; let file_chunk_size = length.unwrap_or(1 * 1024 * 1024 * 1024); // 16GB - let mut chunk_count = file_size / file_chunk_size; + let mut chunk_count = file_size.checked_div(file_chunk_size).unwrap(); let mut size_of_last_chunk = file_size % file_chunk_size; if size_of_last_chunk == 0 {