win: enumerate drives via FindNextVolumeW + GetVolumePathNamesForVolumeNameW

src/windows/disk.rs

@@ -4,13 +4,17 @@ 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, ERROR_MORE_DATA, ERROR_NO_MORE_FILES, HANDLE, MAX_PATH,
+};
 use windows::Win32::Storage::FileSystem::{
-    CreateFileW, GetDiskFreeSpaceExW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW,
-    FILE_ACCESS_RIGHTS, FILE_SHARE_READ, FILE_SHARE_WRITE, OPEN_EXISTING,
+    CreateFileW, FindFirstVolumeW, FindNextVolumeW, FindVolumeClose, GetDiskFreeSpaceExW,
+    GetDriveTypeW, GetVolumeInformationW, GetVolumePathNamesForVolumeNameW, FILE_ACCESS_RIGHTS,
+    FILE_SHARE_READ, FILE_SHARE_WRITE, OPEN_EXISTING,
 };
 use windows::Win32::System::Ioctl::{
     PropertyStandardQuery, StorageDeviceSeekPenaltyProperty, DEVICE_SEEK_PENALTY_DESCRIPTOR,
@@ -19,12 +23,110 @@ 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<u16> {
+    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<Vec<u16>> {
+    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<Vec<u16>> {
+    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.is_empty() && 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<u8>,
     mount_point: Vec<u16>,
-    s_mount_point: String,
+    s_mount_point: OsString,
     total_space: u64,
     available_space: u64,
     is_removable: bool,
@@ -44,7 +146,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 +246,23 @@ unsafe fn get_drive_size(mount_point: &[u16]) -> Option<(u64, u64)> {
 }
 
 pub(crate) unsafe fn get_list() -> Vec<Disk> {
-    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 +271,36 @@ pub(crate) unsafe fn get_list() -> Vec<Disk> {
             )
             .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.is_empty() {
+                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::<Vec<_>>();
+            let handle = match HandleWrapper::new(&device_path[..], Default::default()) {
+                Some(h) => h,
+                None => {
+                    return vec![];
                 }
-                pos += 1;
-            }
-            let file_system: Vec<u8> = 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 +333,29 @@ pub(crate) unsafe fn get_list() -> Vec<Disk> {
                     DiskKind::SSD
                 }
             };
-            Some(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]),
-                    total_space,
-                    available_space,
-                    is_removable,
-                },
-            })
+
+            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::<Vec<_>>();
+            mount_paths
+                .into_iter()
+                .map(move |mount_path| Disk {
+                    inner: DiskInner {
+                        type_,
+                        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::<Vec<_>>()
         })
         .collect::<Vec<_>>()
 }