diff --git a/src/windows/disk.rs b/src/windows/disk.rs index e7e5b18a5..5a53ed00f 100644 --- a/src/windows/disk.rs +++ b/src/windows/disk.rs @@ -4,14 +4,16 @@ use crate::{Disk, DiskKind}; use std::ffi::{c_void, OsStr, OsString}; use std::mem::size_of; +use std::os::windows::ffi::OsStringExt; use std::path::Path; use windows::core::PCWSTR; -use windows::Win32::Foundation::{CloseHandle, HANDLE, MAX_PATH}; +use windows::Win32::Foundation::{CloseHandle, GetLastError, HANDLE, + ERROR_MORE_DATA, ERROR_NO_MORE_FILES, MAX_PATH}; use windows::Win32::Storage::FileSystem::{ - CreateFileW, GetDiskFreeSpaceExW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW, + CreateFileW, GetDiskFreeSpaceExW, GetDriveTypeW, GetVolumeInformationW, FILE_ACCESS_RIGHTS, FILE_SHARE_READ, FILE_SHARE_WRITE, OPEN_EXISTING, -}; + GetVolumePathNamesForVolumeNameW, FindVolumeClose, FindNextVolumeW, FindFirstVolumeW}; use windows::Win32::System::Ioctl::{ PropertyStandardQuery, StorageDeviceSeekPenaltyProperty, DEVICE_SEEK_PENALTY_DESCRIPTOR, IOCTL_STORAGE_QUERY_PROPERTY, STORAGE_PROPERTY_QUERY, @@ -19,12 +21,107 @@ use windows::Win32::System::Ioctl::{ use windows::Win32::System::WindowsProgramming::{DRIVE_FIXED, DRIVE_REMOVABLE}; use windows::Win32::System::IO::DeviceIoControl; +/// Creates a copy of the first zero-terminated wide string in `buf`. +/// The copy includes the zero terminator. +fn from_zero_terminated(buf: &[u16]) -> Vec { + let end = buf.iter().position(|&x| x == 0).unwrap_or(buf.len()); + buf[..=end].to_vec() +} + +// Realistically, volume names are probably not longer than 44 characters, +// but the example in the Microsoft documentation uses MAX_PATH as well. +// https://learn.microsoft.com/en-us/windows/win32/fileio/displaying-volume-paths +const VOLUME_NAME_SIZE: usize = MAX_PATH as usize + 1; + +/// Returns a list of zero-terminated wide strings containing volume GUID paths. +/// Volume GUID paths have the form `\\?\{xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx}\`. +/// +/// Rather confusingly, the Win32 API _also_ calls these "volume names". +pub(crate) fn get_volume_guid_paths() -> Vec> { + let mut volume_names = Vec::new(); + unsafe { + let mut buf = Box::new([0u16; VOLUME_NAME_SIZE]); + let handle = match FindFirstVolumeW(&mut buf[..]) { + Ok(handle) => handle, + Err(_) => { + sysinfo_debug!("Error: FindFirstVolumeW() = {:?}", GetLastError()); + return Vec::new(); + } + }; + volume_names.push(from_zero_terminated(&buf[..])); + loop { + match FindNextVolumeW(handle, &mut buf[..]) { + Ok(_) => (), + Err(_) => { + let find_next_err = GetLastError() + .expect_err("GetLastError should return an error after FindNextVolumeW returned zero."); + if find_next_err.code() != ERROR_NO_MORE_FILES.to_hresult() { + sysinfo_debug!("Error: FindNextVolumeW = {}", find_next_err); + } + break; + } + } + volume_names.push(from_zero_terminated(&buf[..])); + } + if FindVolumeClose(handle) != Ok(()) { + sysinfo_debug!("Error: FindVolumeClose = {:?}", GetLastError()); + }; + } + volume_names +} + +/// Given a volume GUID path (`\\?\{xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx}\`), returns all +/// volume paths (drive letters and mount paths) associated with it +/// as zero terminated wide strings. +/// +/// # Safety +/// `volume_name` must contain a zero-terminated wide string. +pub(crate) unsafe fn get_volume_path_names_for_volume_name(volume_guid_path: &[u16]) -> Vec> { + let volume_guid_path = PCWSTR::from_raw(volume_guid_path.as_ptr()); + + // Initial buffer size is just a guess. There is no direct connection between MAX_PATH + // the output of GetVolumePathNamesForVolumeNameW. + let mut path_names_buf = vec![0u16; MAX_PATH as usize]; + let mut path_names_output_size = 0u32; + for _ in 0..10 { + match GetVolumePathNamesForVolumeNameW( + volume_guid_path, + Some(path_names_buf.as_mut_slice()), + &mut path_names_output_size) + .map_err(|_| GetLastError() + .expect_err("GetLastError should return an error after GetVolumePathNamesForVolumeNameW returned zero.") + .code()) { + Ok(()) => break, + Err(e) if e == ERROR_MORE_DATA.to_hresult() => { + // We need a bigger buffer. path_names_output_size contains the required buffer size. + path_names_buf = vec![0u16; path_names_output_size as usize]; + continue; + }, + Err(_e) => { + sysinfo_debug!("Error: GetVolumePathNamesForVolumeNameW() = {}", _e); + return Vec::new(); + } + } + } + + // path_names_buf contains multiple zero terminated wide strings. + // An additional zero terminates the list. + let mut path_names = Vec::new(); + let mut buf = &path_names_buf[..]; + while buf.len() > 0 && buf[0] != 0 { + let path = from_zero_terminated(buf); + buf = &buf[path.len()..]; + path_names.push(path); + } + path_names +} + pub(crate) struct DiskInner { type_: DiskKind, name: OsString, file_system: Vec, mount_point: Vec, - s_mount_point: String, + s_mount_point: OsString, total_space: u64, available_space: u64, is_removable: bool, @@ -44,7 +141,7 @@ impl DiskInner { } pub(crate) fn mount_point(&self) -> &Path { - Path::new(&self.s_mount_point) + self.s_mount_point.as_ref() } pub(crate) fn total_space(&self) -> u64 { @@ -144,33 +241,24 @@ unsafe fn get_drive_size(mount_point: &[u16]) -> Option<(u64, u64)> { } pub(crate) unsafe fn get_list() -> Vec { - let drives = GetLogicalDrives(); - if drives == 0 { - return Vec::new(); - } #[cfg(feature = "multithread")] use rayon::iter::ParallelIterator; - crate::utils::into_iter(0..u32::BITS) - .filter_map(|x| { - if (drives >> x) & 1 == 0 { - return None; - } - let mount_point = [b'A' as u16 + x as u16, b':' as u16, b'\\' as u16, 0]; - - let raw_mount_point = PCWSTR::from_raw(mount_point.as_ptr()); - let drive_type = GetDriveTypeW(raw_mount_point); + crate::utils::into_iter(get_volume_guid_paths()) + .flat_map(|volume_name| { + let raw_volume_name = PCWSTR::from_raw(volume_name.as_ptr()); + let drive_type = GetDriveTypeW(raw_volume_name); let is_removable = drive_type == DRIVE_REMOVABLE; if drive_type != DRIVE_FIXED && drive_type != DRIVE_REMOVABLE { - return None; + return vec![]; } let mut name = [0u16; MAX_PATH as usize + 1]; let mut file_system = [0u16; 32]; let volume_info_res = GetVolumeInformationW( - raw_mount_point, + raw_volume_name, Some(&mut name), None, None, @@ -179,40 +267,32 @@ pub(crate) unsafe fn get_list() -> Vec { ) .is_ok(); if !volume_info_res { - return None; + sysinfo_debug!("Error: GetVolumeInformationW = {:?}", GetLastError()); + return vec![]; } - let mut pos = 0; - for x in name.iter() { - if *x == 0 { - break; - } - pos += 1; + + let mount_paths = get_volume_path_names_for_volume_name(&volume_name[..]); + if mount_paths.len() == 0 { + return vec![]; } - let name = String::from_utf16_lossy(&name[..pos]); - let name = OsStr::new(&name); - pos = 0; - for x in file_system.iter() { - if *x == 0 { - break; + // The device path is the volume name without the trailing backslash. + let device_path = volume_name[..(volume_name.len()-2)].iter().copied().chain([0]).collect::>(); + let handle = match HandleWrapper::new(&device_path[..], Default::default()) { + Some(h) => h, + None => { + return vec![]; } - pos += 1; - } - let file_system: Vec = file_system[..pos].iter().map(|x| *x as u8).collect(); - - let drive_name = [ - b'\\' as u16, - b'\\' as u16, - b'.' as u16, - b'\\' as u16, - b'A' as u16 + x as u16, - b':' as u16, - 0, - ]; - let handle = HandleWrapper::new(&drive_name, Default::default())?; - let (total_space, available_space) = get_drive_size(&mount_point)?; + }; + let (total_space, available_space) = match get_drive_size(&mount_paths[0][..]) { + Some(space) => space, + None => { + return vec![]; + } + }; if total_space == 0 { - return None; + sysinfo_debug!("total_space == 0"); + return vec![]; } let spq_trim = STORAGE_PROPERTY_QUERY { PropertyId: StorageDeviceSeekPenaltyProperty, @@ -245,18 +325,25 @@ pub(crate) unsafe fn get_list() -> Vec { DiskKind::SSD } }; - Some(Disk { + + let name_len = name.iter().position(|&x| x == 0).unwrap_or(name.len()); + let name = OsString::from_wide(&name[..name_len]); + let file_system = file_system.iter() + .take_while(|c| **c != 0) + .map(|c| *c as u8) + .collect::>(); + mount_paths.into_iter().map(move |mount_path| Disk { inner: DiskInner { type_, - name: name.to_owned(), - file_system: file_system.to_vec(), - mount_point: mount_point.to_vec(), - s_mount_point: String::from_utf16_lossy(&mount_point[..mount_point.len() - 1]), + name: name.clone(), + file_system: file_system.clone(), + s_mount_point: OsString::from_wide(&mount_path[..mount_path.len()-1]), + mount_point: mount_path, total_space, available_space, is_removable, }, - }) + }).collect::>() }) .collect::>() }