use crate::error::{Aria2Error, FatalError, Result};
use std::fmt;
use std::path::Path;
const DEFAULT_MARGIN_MB: u64 = 100;
#[cfg(unix)]
fn path_to_cstring(path: &Path) -> Option<std::ffi::CString> {
use std::os::unix::ffi::OsStrExt;
std::ffi::CString::new(path.as_os_str().as_bytes()).ok()
}
#[derive(Debug, Clone)]
pub enum DiskError {
InsufficientSpace {
required: u64,
available: Option<u64>,
},
IoError(String),
PermissionDenied(String),
}
impl fmt::Display for DiskError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
DiskError::InsufficientSpace {
required,
available,
} => {
write!(
f,
"Not enough disk space: need {}, have {}",
format_bytes(*required),
available.map_or_else(|| "unknown".into(), format_bytes)
)
}
DiskError::IoError(msg) => write!(f, "Disk I/O error: {}", msg),
DiskError::PermissionDenied(path) => write!(f, "Permission denied: {}", path),
}
}
}
impl std::error::Error for DiskError {}
fn format_bytes(bytes: u64) -> String {
if bytes >= 1024 * 1024 * 1024 {
format!("{:.2} GiB", bytes as f64 / (1024.0 * 1024.0 * 1024.0))
} else if bytes >= 1024 * 1024 {
format!("{:.2} MiB", bytes as f64 / (1024.0 * 1024.0))
} else if bytes >= 1024 {
format!("{:.2} KiB", bytes as f64 / 1024.0)
} else {
format!("{} B", bytes)
}
}
pub fn check_disk_space(path: &Path, required_bytes: u64) -> std::result::Result<(), String> {
#[cfg(unix)]
{
let check_path = if path.as_os_str().is_empty() {
Path::new(".")
} else {
path
};
let Some(c_path) = path_to_cstring(check_path) else {
tracing::warn!(
path = %check_path.display(),
required = required_bytes,
"path contains embedded null byte, skipping disk space check"
);
return Ok(());
};
let mut stat: libc::statvfs = unsafe { std::mem::zeroed() };
let ret = unsafe { libc::statvfs(c_path.as_ptr(), &mut stat) };
if ret != 0 {
tracing::warn!(
path = %check_path.display(),
required = required_bytes,
"statvfs failed, skipping disk space check: {}",
std::io::Error::last_os_error()
);
return Ok(());
}
let available = stat.f_bavail as u64 * stat.f_frsize as u64;
let needed_with_headroom = required_bytes.saturating_add(required_bytes / 10);
if available < needed_with_headroom {
let available_gb = available as f64 / (1024.0 * 1024.0 * 1024.0);
let required_gb = required_bytes as f64 / (1024.0 * 1024.0 * 1024.0);
return Err(format!(
"Insufficient disk space: need {:.2} GiB but only {:.2} GiB available on '{}'",
required_gb,
available_gb,
check_path.display()
));
}
Ok(())
}
#[cfg(windows)]
{
use std::os::windows::ffi::OsStrExt;
let check_path = if path.as_os_str().is_empty() {
Path::new(".")
} else {
path
};
let abs_path = match std::fs::canonicalize(check_path) {
Ok(p) => p,
Err(_) => {
std::env::current_dir().unwrap_or_else(|_| Path::new(".").to_path_buf())
}
};
let mut wide_path: Vec<u16> = abs_path.as_os_str().encode_wide().collect();
wide_path.push(0);
let mut free_bytes_available: u64 = 0;
let mut total_number_of_bytes: u64 = 0;
let mut total_number_of_free_bytes: u64 = 0;
let result = unsafe {
windows_sys::Win32::Storage::FileSystem::GetDiskFreeSpaceExW(
wide_path.as_ptr(),
&mut free_bytes_available as *mut u64 as *mut _,
&mut total_number_of_bytes as *mut u64 as *mut _,
&mut total_number_of_free_bytes as *mut u64 as *mut _,
)
};
if result == 0 {
tracing::warn!(
path = %check_path.display(),
required = required_bytes,
"GetDiskFreeSpaceEx failed, skipping disk space check"
);
return Ok(());
}
let needed_with_headroom = required_bytes.saturating_add(required_bytes / 10);
if free_bytes_available < needed_with_headroom {
let available_gb = free_bytes_available as f64 / (1024.0 * 1024.0 * 1024.0);
let required_gb = required_bytes as f64 / (1024.0 * 1024.0 * 1024.0);
return Err(format!(
"Insufficient disk space: need {:.2} GiB but only {:.2} GiB available on '{}'",
required_gb,
available_gb,
check_path.display()
));
}
Ok(())
}
#[cfg(not(any(unix, windows)))]
{
tracing::warn!(
path = %path.display(),
required = required_bytes,
"Disk space check skipped on unsupported platform"
);
Ok(())
}
}
pub fn check_disk_space_typed(
path: &Path,
required_bytes: u64,
) -> std::result::Result<(), DiskError> {
#[cfg(unix)]
{
let check_path = if path.as_os_str().is_empty() {
Path::new(".")
} else {
path
};
let Some(c_path) = path_to_cstring(check_path) else {
tracing::warn!(
path = %check_path.display(),
required = required_bytes,
"path contains embedded null byte, skipping disk space check"
);
return Ok(());
};
let mut stat: libc::statvfs = unsafe { std::mem::zeroed() };
let ret = unsafe { libc::statvfs(c_path.as_ptr(), &mut stat) };
if ret != 0 {
tracing::warn!(
path = %check_path.display(),
required = required_bytes,
"statvfs failed, skipping disk space check: {}",
std::io::Error::last_os_error()
);
return Ok(());
}
let available = stat.f_bavail as u64 * stat.f_frsize as u64;
let needed_with_headroom = required_bytes.saturating_add(required_bytes / 10);
if available < needed_with_headroom {
return Err(DiskError::InsufficientSpace {
required: needed_with_headroom,
available: Some(available),
});
}
Ok(())
}
#[cfg(windows)]
{
use std::os::windows::ffi::OsStrExt;
let check_path = if path.as_os_str().is_empty() {
Path::new(".")
} else {
path
};
let abs_path = match std::fs::canonicalize(check_path) {
Ok(p) => p,
Err(_) => std::env::current_dir().unwrap_or_else(|_| Path::new(".").to_path_buf()),
};
let mut wide_path: Vec<u16> = abs_path.as_os_str().encode_wide().collect();
wide_path.push(0);
let mut free_bytes_available: u64 = 0;
let mut total_number_of_bytes: u64 = 0;
let mut total_number_of_free_bytes: u64 = 0;
let result = unsafe {
windows_sys::Win32::Storage::FileSystem::GetDiskFreeSpaceExW(
wide_path.as_ptr(),
&mut free_bytes_available as *mut u64 as *mut _,
&mut total_number_of_bytes as *mut u64 as *mut _,
&mut total_number_of_free_bytes as *mut u64 as *mut _,
)
};
if result == 0 {
return Ok(());
}
let needed_with_headroom = required_bytes.saturating_add(required_bytes / 10);
if free_bytes_available < needed_with_headroom {
return Err(DiskError::InsufficientSpace {
required: needed_with_headroom,
available: Some(free_bytes_available),
});
}
Ok(())
}
#[cfg(not(any(unix, windows)))]
{
let _ = path;
let _ = required_bytes;
Ok(())
}
}
pub fn available_space(path: &Path) -> Result<u64> {
let path = if path.as_os_str().is_empty() {
Path::new(".")
} else {
path
};
#[cfg(unix)]
{
let c_path = path_to_cstring(path).ok_or_else(|| {
Aria2Error::Fatal(FatalError::Config(format!(
"path contains embedded null byte: {}",
path.display()
)))
})?;
let mut stat: libc::statvfs = unsafe { std::mem::zeroed() };
let ret = unsafe { libc::statvfs(c_path.as_ptr(), &mut stat) };
if ret != 0 {
return Err(Aria2Error::Fatal(FatalError::Config(format!(
"statvfs failed: {}",
std::io::Error::last_os_error()
))));
}
Ok(stat.f_bavail as u64 * stat.f_frsize as u64)
}
#[cfg(windows)]
{
use std::os::windows::ffi::OsStrExt;
let abs_path = match std::fs::canonicalize(path) {
Ok(p) => p,
Err(_) => std::env::current_dir().unwrap_or_else(|_| Path::new(".").to_path_buf()),
};
let mut wide_path: Vec<u16> = abs_path.as_os_str().encode_wide().collect();
wide_path.push(0);
let mut free_bytes_available: u64 = 0;
let mut total_number_of_bytes: u64 = 0;
let mut total_number_of_free_bytes: u64 = 0;
let result = unsafe {
windows_sys::Win32::Storage::FileSystem::GetDiskFreeSpaceExW(
wide_path.as_ptr(),
&mut free_bytes_available as *mut u64 as *mut _,
&mut total_number_of_bytes as *mut u64 as *mut _,
&mut total_number_of_free_bytes as *mut u64 as *mut _,
)
};
if result == 0 {
return Err(Aria2Error::Fatal(FatalError::Config(
"GetDiskFreeSpaceExW failed".to_string(),
)));
}
Ok(free_bytes_available)
}
#[cfg(not(any(unix, windows)))]
{
let _ = path;
Ok(u64::MAX)
}
}
pub fn has_enough_space(path: &Path, required: u64) -> bool {
available_space(path).is_ok_and(|avail| avail >= required)
}
pub fn check_with_margin(path: &Path, required: u64, margin_mb: Option<u64>) -> Result<()> {
let margin = margin_mb.unwrap_or(DEFAULT_MARGIN_MB) * 1024 * 1024;
let total_needed = required.saturating_add(margin);
let avail = match available_space(path) {
Ok(v) => v,
Err(err) => {
tracing::warn!(
path = %path.display(),
required = total_needed,
"available_space failed, skipping disk space check: {}",
err
);
return Ok(());
}
};
if avail < total_needed {
Err(Aria2Error::Fatal(FatalError::DiskSpaceExhausted))
} else {
Ok(())
}
}
pub fn total_space(path: &Path) -> Result<u64> {
let path = if path.as_os_str().is_empty() {
Path::new(".")
} else {
path
};
#[cfg(unix)]
{
let c_path = path_to_cstring(path).ok_or_else(|| {
Aria2Error::Fatal(FatalError::Config(format!(
"path contains embedded null byte: {}",
path.display()
)))
})?;
let mut stat: libc::statvfs = unsafe { std::mem::zeroed() };
let ret = unsafe { libc::statvfs(c_path.as_ptr(), &mut stat) };
if ret != 0 {
return Err(Aria2Error::Fatal(FatalError::Config(format!(
"statvfs failed: {}",
std::io::Error::last_os_error()
))));
}
Ok(stat.f_blocks as u64 * stat.f_frsize as u64)
}
#[cfg(windows)]
{
use std::os::windows::ffi::OsStrExt;
let abs_path = match std::fs::canonicalize(path) {
Ok(p) => p,
Err(_) => std::env::current_dir().unwrap_or_else(|_| Path::new(".").to_path_buf()),
};
let mut wide_path: Vec<u16> = abs_path.as_os_str().encode_wide().collect();
wide_path.push(0);
let mut free_bytes_available: u64 = 0;
let mut total_number_of_bytes: u64 = 0;
let mut total_number_of_free_bytes: u64 = 0;
let result = unsafe {
windows_sys::Win32::Storage::FileSystem::GetDiskFreeSpaceExW(
wide_path.as_ptr(),
&mut free_bytes_available as *mut u64 as *mut _,
&mut total_number_of_bytes as *mut u64 as *mut _,
&mut total_number_of_free_bytes as *mut u64 as *mut _,
)
};
if result == 0 {
return Err(Aria2Error::Fatal(FatalError::Config(
"GetDiskFreeSpaceExW failed".to_string(),
)));
}
Ok(total_number_of_bytes)
}
#[cfg(not(any(unix, windows)))]
{
let _ = path;
Ok(u64::MAX)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_sufficient_space_ok() {
let required = 100 * 1024 * 1024; let result = check_disk_space(Path::new("."), required);
assert!(
result.is_ok() || result.is_err(),
"Should return Ok or Err without panicking"
);
}
#[test]
fn test_insufficient_space_err() {
let required = u64::MAX / 2; let result = check_disk_space(Path::new("."), required);
if let Err(error_msg) = result {
assert!(
error_msg.to_lowercase().contains("insufficient")
|| error_msg.to_lowercase().contains("space"),
"Error message should mention space: {}",
error_msg
);
}
}
#[test]
fn test_disk_error_display() {
let err = DiskError::InsufficientSpace {
required: 1024 * 1024 * 1024, available: Some(512 * 1024 * 1024), };
let display_str = format!("{}", err);
assert!(
display_str.contains("Not enough disk space"),
"Should mention insufficient space"
);
assert!(
display_str.contains("1.00 GiB") || display_str.contains("GiB"),
"Should show required size in GiB"
);
assert!(
display_str.contains("512.00 MiB") || display_str.contains("MiB"),
"Should show available size in MiB"
);
let io_err = DiskError::IoError("Failed to write block".to_string());
let io_display = format!("{}", io_err);
assert!(
io_display.contains("Disk I/O error"),
"Should mention I/O error"
);
assert!(
io_display.contains("Failed to write block"),
"Should include original message"
);
let perm_err = DiskError::PermissionDenied("/root/secret".to_string());
let perm_display = format!("{}", perm_err);
assert!(
perm_display.contains("Permission denied"),
"Should mention permission denied"
);
assert!(perm_display.contains("/root/secret"), "Should include path");
assert_eq!(format_bytes(0), "0 B", "Zero bytes");
assert_eq!(format_bytes(500), "500 B", "Small bytes");
assert_eq!(format_bytes(1024), "1.00 KiB", "Exactly 1 KiB");
assert_eq!(format_bytes(1024 * 1024), "1.00 MiB", "Exactly 1 MiB");
assert_eq!(
format_bytes(1024 * 1024 * 1024),
"1.00 GiB",
"Exactly 1 GiB"
);
assert_eq!(
format_bytes(1536 * 1024 * 1024),
"1.50 GiB",
"1.5 GiB with decimal"
);
}
#[test]
fn test_check_disk_space_typed_returns_structured_errors() {
let huge_request = u64::MAX / 2;
match check_disk_space_typed(Path::new("."), huge_request) {
Ok(()) => {
}
Err(DiskError::InsufficientSpace {
required,
available,
}) => {
assert!(required > 0, "Required should be positive");
if let Some(avail) = available {
let _ = avail;
}
}
Err(DiskError::IoError(msg)) => {
assert!(!msg.is_empty(), "Error message should not be empty");
}
Err(DiskError::PermissionDenied(_)) => {
}
}
}
#[test]
fn test_check_disk_space_empty_path() {
let result = check_disk_space(Path::new(""), 1024);
assert!(
result.is_ok() || result.is_err(),
"Empty path should be handled gracefully"
);
}
#[test]
fn test_check_disk_space_zero_bytes() {
let result = check_disk_space(Path::new("."), 0);
assert!(
result.is_ok(),
"Requesting zero bytes should always succeed"
);
}
}