1 files changed, 104 insertions, 177 deletions

diff --git a/vendor/sysinfo/src/apple/macos/disk.rs b/vendor/sysinfo/src/apple/macos/disk.rs
index 7cc5a83ce..3a4372a2f 100644
--- a/vendor/sysinfo/src/apple/macos/disk.rs
+++ b/vendor/sysinfo/src/apple/macos/disk.rs
@@ -1,199 +1,126 @@
 // Take a look at the license at the top of the repository in the LICENSE file.
 
-use crate::sys::{ffi, utils};
-use crate::utils::to_cpath;
-use crate::{Disk, DiskType};
+use crate::sys::ffi;
+use crate::sys::{
+    disk::{get_str_value, DictKey},
+    macos::utils::IOReleaser,
+    utils::CFReleaser,
+};
+use crate::DiskType;
 
-use core_foundation_sys::base::{kCFAllocatorDefault, kCFAllocatorNull, CFRelease};
-use core_foundation_sys::dictionary::{CFDictionaryGetValueIfPresent, CFDictionaryRef};
-use core_foundation_sys::number::{kCFBooleanTrue, CFBooleanRef};
+use core_foundation_sys::base::{kCFAllocatorDefault, kCFAllocatorNull};
 use core_foundation_sys::string as cfs;
 
-use libc::{c_char, c_int, c_void, statfs};
+use std::ffi::CStr;
 
-use std::ffi::{OsStr, OsString};
-use std::mem;
-use std::os::unix::ffi::OsStrExt;
-use std::path::PathBuf;
-use std::ptr;
+pub(crate) fn get_disk_type(disk: &libc::statfs) -> Option<DiskType> {
+    let characteristics_string = unsafe {
+        CFReleaser::new(cfs::CFStringCreateWithBytesNoCopy(
+            kCFAllocatorDefault,
+            ffi::kIOPropertyDeviceCharacteristicsKey.as_ptr(),
+            ffi::kIOPropertyDeviceCharacteristicsKey.len() as _,
+            cfs::kCFStringEncodingUTF8,
+            false as _,
+            kCFAllocatorNull,
+        ))?
+    };
 
-fn to_path(mount_path: &[c_char]) -> Option<PathBuf> {
-    let mut tmp = Vec::with_capacity(mount_path.len());
-    for &c in mount_path {
-        if c == 0 {
-            break;
-        }
-        tmp.push(c as u8);
-    }
-    if tmp.is_empty() {
-        None
-    } else {
-        let path = OsStr::from_bytes(&tmp);
-        Some(PathBuf::from(path))
-    }
-}
+    // Removes `/dev/` from the value.
+    let bsd_name = unsafe {
+        CStr::from_ptr(disk.f_mntfromname.as_ptr())
+            .to_bytes()
+            .strip_prefix(b"/dev/")
+            .or_else(|| {
+                sysinfo_debug!("unknown disk mount path format");
+                None
+            })?
+    };
+
+    // We don't need to wrap this in an auto-releaser because the following call to `IOServiceGetMatchingServices`
+    // will take ownership of one retain reference.
+    let matching =
+        unsafe { ffi::IOBSDNameMatching(ffi::kIOMasterPortDefault, 0, bsd_name.as_ptr().cast()) };
 
-pub(crate) fn get_disks(session: ffi::DASessionRef) -> Vec<Disk> {
-    if session.is_null() {
-        return Vec::new();
+    if matching.is_null() {
+        return None;
     }
-    unsafe {
-        let count = libc::getfsstat(ptr::null_mut(), 0, libc::MNT_NOWAIT);
-        if count < 1 {
-            return Vec::new();
-        }
-        let bufsize = count * mem::size_of::<libc::statfs>() as c_int;
-        let mut disks = Vec::with_capacity(count as _);
-        let count = libc::getfsstat(disks.as_mut_ptr(), bufsize, libc::MNT_NOWAIT);
-        if count < 1 {
-            return Vec::new();
-        }
-        disks.set_len(count as _);
-        disks
-            .into_iter()
-            .filter_map(|c_disk| {
-                let mount_point = to_path(&c_disk.f_mntonname)?;
-                let disk = ffi::DADiskCreateFromBSDName(
-                    kCFAllocatorDefault as _,
-                    session,
-                    c_disk.f_mntfromname.as_ptr(),
-                );
-                let dict = ffi::DADiskCopyDescription(disk);
-                if dict.is_null() {
-                    return None;
-                }
-                // Keeping this around in case one might want the list of the available
-                // keys in "dict".
-                // core_foundation_sys::base::CFShow(dict as _);
-                let name = match get_str_value(dict, b"DAMediaName\0").map(OsString::from) {
-                    Some(n) => n,
-                    None => return None,
-                };
-                let removable = get_bool_value(dict, b"DAMediaRemovable\0").unwrap_or(false);
-                let ejectable = get_bool_value(dict, b"DAMediaEjectable\0").unwrap_or(false);
-                // This is very hackish but still better than nothing...
-                let type_ = if let Some(model) = get_str_value(dict, b"DADeviceModel\0") {
-                    if model.contains("SSD") {
-                        DiskType::SSD
-                    } else {
-                        // We just assume by default that this is a HDD
-                        DiskType::HDD
-                    }
-                } else {
-                    DiskType::Unknown(-1)
-                };
 
-                CFRelease(dict as _);
-                new_disk(name, mount_point, type_, removable || ejectable)
-            })
-            .collect::<Vec<_>>()
+    let mut service_iterator: ffi::io_iterator_t = 0;
+
+    if unsafe {
+        ffi::IOServiceGetMatchingServices(
+            ffi::kIOMasterPortDefault,
+            matching.cast(),
+            &mut service_iterator,
+        )
+    } != libc::KERN_SUCCESS
+    {
+        return None;
     }
-}
 
-unsafe fn get_dict_value<T, F: FnOnce(*const c_void) -> Option<T>>(
-    dict: CFDictionaryRef,
-    key: &[u8],
-    callback: F,
-) -> Option<T> {
-    let key = ffi::CFStringCreateWithCStringNoCopy(
-        ptr::null_mut(),
-        key.as_ptr() as *const c_char,
-        cfs::kCFStringEncodingUTF8,
-        kCFAllocatorNull as _,
-    );
-    let mut value = std::ptr::null();
-    let ret = if CFDictionaryGetValueIfPresent(dict, key as _, &mut value) != 0 {
-        callback(value)
-    } else {
-        None
-    };
-    CFRelease(key as _);
-    ret
-}
+    // Safety: We checked for success, so there is always a valid iterator, even if its empty.
+    let service_iterator = unsafe { IOReleaser::new_unchecked(service_iterator) };
 
-unsafe fn get_str_value(dict: CFDictionaryRef, key: &[u8]) -> Option<String> {
-    get_dict_value(dict, key, |v| {
-        let v = v as cfs::CFStringRef;
-
-        let len_utf16 = cfs::CFStringGetLength(v);
-        let len_bytes = len_utf16 as usize * 2; // Two bytes per UTF-16 codepoint.
-
-        let v_ptr = cfs::CFStringGetCStringPtr(v, cfs::kCFStringEncodingUTF8);
-        if v_ptr.is_null() {
-            // Fallback on CFStringGetString to read the underlying bytes from the CFString.
-            let mut buf = vec![0; len_bytes];
-            let success = cfs::CFStringGetCString(
-                v,
-                buf.as_mut_ptr(),
-                len_bytes as _,
-                cfs::kCFStringEncodingUTF8,
-            );
-
-            if success != 0 {
-                utils::vec_to_rust(buf)
-            } else {
-                None
+    let mut parent_entry: ffi::io_registry_entry_t = 0;
+
+    while let Some(mut current_service_entry) =
+        IOReleaser::new(unsafe { ffi::IOIteratorNext(service_iterator.inner()) })
+    {
+        // Note: This loop is required in a non-obvious way. Due to device properties existing as a tree
+        // in IOKit, we may need an arbitrary number of calls to `IORegistryEntryCreateCFProperty` in order to find
+        // the values we are looking for. The function may return nothing if we aren't deep enough into the registry
+        // tree, so we need to continue going from child->parent node until its found.
+        loop {
+            if unsafe {
+                ffi::IORegistryEntryGetParentEntry(
+                    current_service_entry.inner(),
+                    ffi::kIOServicePlane.as_ptr().cast(),
+                    &mut parent_entry,
+                )
+            } != libc::KERN_SUCCESS
+            {
+                break;
             }
-        } else {
-            utils::cstr_to_rust_with_size(v_ptr, Some(len_bytes))
-        }
-    })
-}
 
-unsafe fn get_bool_value(dict: CFDictionaryRef, key: &[u8]) -> Option<bool> {
-    get_dict_value(dict, key, |v| Some(v as CFBooleanRef == kCFBooleanTrue))
-}
+            current_service_entry = match IOReleaser::new(parent_entry) {
+                Some(service) => service,
+                // There were no more parents left
+                None => break,
+            };
 
-fn new_disk(
-    name: OsString,
-    mount_point: PathBuf,
-    type_: DiskType,
-    is_removable: bool,
-) -> Option<Disk> {
-    let mount_point_cpath = to_cpath(&mount_point);
-    let mut total_space = 0;
-    let mut available_space = 0;
-    let mut file_system = None;
-    unsafe {
-        let mut stat: statfs = mem::zeroed();
-        if statfs(mount_point_cpath.as_ptr() as *const i8, &mut stat) == 0 {
-            // APFS is "special" because its a snapshot-based filesystem, and modern
-            // macOS devices take full advantage of this.
-            //
-            // By default, listing volumes with `statfs` can return both the root-level
-            // "data" partition and any snapshots that exist. However, other than some flags and
-            // reserved(undocumented) bytes, there is no difference between the OS boot snapshot
-            // and the "data" partition.
-            //
-            // To avoid duplicating the number of disks (and therefore available space, etc), only return
-            // a disk (which is really a partition with APFS) if it is the root of the filesystem.
-            let is_root = stat.f_flags & libc::MNT_ROOTFS as u32 == 0;
-            if !is_root {
-                return None;
-            }
+            let properties_result = unsafe {
+                CFReleaser::new(ffi::IORegistryEntryCreateCFProperty(
+                    current_service_entry.inner(),
+                    characteristics_string.inner(),
+                    kCFAllocatorDefault,
+                    0,
+                ))
+            };
+
+            if let Some(device_properties) = properties_result {
+                let disk_type = unsafe {
+                    super::disk::get_str_value(
+                        device_properties.inner(),
+                        DictKey::Defined(ffi::kIOPropertyMediumTypeKey),
+                    )
+                };
 
-            total_space = u64::from(stat.f_bsize).saturating_mul(stat.f_blocks);
-            available_space = u64::from(stat.f_bsize).saturating_mul(stat.f_bavail);
-            let mut vec = Vec::with_capacity(stat.f_fstypename.len());
-            for x in &stat.f_fstypename {
-                if *x == 0 {
-                    break;
+                if let Some(disk_type) = disk_type.and_then(|medium| match medium.as_str() {
+                    _ if medium == ffi::kIOPropertyMediumTypeSolidStateKey => Some(DiskType::SSD),
+                    _ if medium == ffi::kIOPropertyMediumTypeRotationalKey => Some(DiskType::HDD),
+                    _ => None,
+                }) {
+                    return Some(disk_type);
+                } else {
+                    // Many external drive vendors do not advertise their device's storage medium.
+                    //
+                    // In these cases, assuming that there were _any_ properties about them registered, we fallback
+                    // to `HDD` when no storage medium is provided by the device instead of `Unknown`.
+                    return Some(DiskType::HDD);
                 }
-                vec.push(*x as u8);
             }
-            file_system = Some(vec);
-        }
-        if total_space == 0 {
-            return None;
         }
-        Some(Disk {
-            type_,
-            name,
-            file_system: file_system.unwrap_or_else(|| b"<Unknown>".to_vec()),
-            mount_point,
-            total_space,
-            available_space,
-            is_removable,
-        })
     }
+
+    None
 }