Skip to main content

workspace/
disk.rs

1use std::path::Path;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub struct DiskSpace {
5    pub available_bytes: u64,
6    pub total_bytes: u64,
7}
8
9#[cfg(unix)]
10#[allow(unsafe_code)] // FFI: statvfs(3) — SAFETY contract inside the body.
11pub fn get_disk_space(path: &Path) -> std::io::Result<DiskSpace> {
12    use std::ffi::CString;
13    use std::os::unix::ffi::OsStrExt;
14
15    // Ensure directory exists or check its parent
16    let target = if path.exists() {
17        path.to_path_buf()
18    } else if let Some(parent) = path.parent() {
19        if parent.exists() {
20            parent.to_path_buf()
21        } else {
22            std::env::current_dir().unwrap_or_else(|_| Path::new(".").to_path_buf())
23        }
24    } else {
25        std::env::current_dir().unwrap_or_else(|_| Path::new(".").to_path_buf())
26    };
27
28    let c_path = CString::new(target.as_os_str().as_bytes())
29        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidInput, e))?;
30
31    // `statvfs` is a plain-old-data C struct of integers, so all-zero is a
32    // valid bit pattern; we still go through MaybeUninit so the struct is
33    // only ever assumed initialized after a successful syscall.
34    let mut stat = std::mem::MaybeUninit::<libc::statvfs>::zeroed();
35    // SAFETY: `c_path` is a live NUL-terminated CString and `stat` is a
36    // writable, properly aligned slot for exactly the struct statvfs(3)
37    // fills in. On success statvfs fully initializes the struct, which is
38    // the only case where we call assume_init below.
39    let ret = unsafe { libc::statvfs(c_path.as_ptr(), stat.as_mut_ptr()) };
40    if ret != 0 {
41        return Err(std::io::Error::last_os_error());
42    }
43    // SAFETY: statvfs returned 0, so the struct is fully initialized.
44    let stat = unsafe { stat.assume_init() };
45
46    // The statvfs block counters are u32 on some targets (macOS) and u64 on
47    // others (Linux). `as_u64` is a real widening conversion on the former
48    // and the identity on the latter, and (unlike a bare `as u64` or
49    // `u64::from`) it is correct on every platform without tripping
50    // clippy's same-type conversion lints on Linux.
51    fn as_u64(v: impl Into<u64>) -> u64 {
52        v.into()
53    }
54
55    let frsize = if stat.f_frsize > 0 {
56        as_u64(stat.f_frsize)
57    } else {
58        as_u64(stat.f_bsize)
59    };
60
61    let available_bytes = as_u64(stat.f_bavail) * frsize;
62    let total_bytes = as_u64(stat.f_blocks) * frsize;
63
64    Ok(DiskSpace {
65        available_bytes,
66        total_bytes,
67    })
68}
69
70#[cfg(windows)]
71#[allow(unsafe_code)] // FFI: GetDiskFreeSpaceExW — SAFETY contract inside the body.
72pub fn get_disk_space(path: &Path) -> std::io::Result<DiskSpace> {
73    use std::os::windows::ffi::OsStrExt;
74
75    let target = if path.exists() {
76        path.to_path_buf()
77    } else if let Some(parent) = path.parent() {
78        if parent.exists() {
79            parent.to_path_buf()
80        } else {
81            std::env::current_dir().unwrap_or_else(|_| Path::new(".").to_path_buf())
82        }
83    } else {
84        std::env::current_dir().unwrap_or_else(|_| Path::new(".").to_path_buf())
85    };
86
87    let mut wide_path: Vec<u16> = target.as_os_str().encode_wide().collect();
88    wide_path.push(0);
89
90    let mut free_bytes_available: u64 = 0;
91    let mut total_number_of_bytes: u64 = 0;
92    let mut total_number_of_free_bytes: u64 = 0;
93
94    extern "system" {
95        fn GetDiskFreeSpaceExW(
96            lpDirectoryName: *const u16,
97            lpFreeBytesAvailableToCaller: *mut u64,
98            lpTotalNumberOfBytes: *mut u64,
99            lpTotalNumberOfFreeBytes: *mut u64,
100        ) -> i32;
101    }
102
103    // SAFETY: `wide_path` is a NUL-terminated UTF-16 buffer (the trailing 0
104    // is pushed above) and outlives the call, so `lpDirectoryName` is a valid
105    // null-terminated wide string. The three output pointers address live,
106    // writable `u64` locals, and GetDiskFreeSpaceExW is documented to fill
107    // each of them (or fail without writing) — it never retains the pointers.
108    let ret = unsafe {
109        GetDiskFreeSpaceExW(
110            wide_path.as_ptr(),
111            &mut free_bytes_available,
112            &mut total_number_of_bytes,
113            &mut total_number_of_free_bytes,
114        )
115    };
116
117    if ret == 0 {
118        return Err(std::io::Error::last_os_error());
119    }
120
121    Ok(DiskSpace {
122        available_bytes: free_bytes_available,
123        total_bytes: total_number_of_bytes,
124    })
125}
126
127#[cfg(not(any(unix, windows)))]
128pub fn get_disk_space(_path: &Path) -> std::io::Result<DiskSpace> {
129    Ok(DiskSpace {
130        available_bytes: u64::MAX,
131        total_bytes: u64::MAX,
132    })
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[test]
140    fn test_disk_space_query() {
141        let cwd = std::env::current_dir().unwrap();
142        let space = get_disk_space(&cwd).expect("get_disk_space should succeed");
143        assert!(space.total_bytes > 0, "total disk space should be > 0");
144        assert!(space.available_bytes <= space.total_bytes);
145    }
146}