Skip to content

Commit

Permalink
Improve object releasers and cleanup usage sites
Browse files Browse the repository at this point in the history
  • Loading branch information
complexspaces committed Oct 30, 2022
1 parent a4e6255 commit 13c32c9
Show file tree
Hide file tree
Showing 4 changed files with 82 additions and 61 deletions.
47 changes: 30 additions & 17 deletions src/apple/disk.rs
Expand Up @@ -62,16 +62,20 @@ impl DiskExt for Disk {

fn refresh(&mut self) -> bool {
unsafe {
let requested_properties = build_requested_properties(&[
if let Some(requested_properties) = build_requested_properties(&[
ffi::kCFURLVolumeAvailableCapacityKey,
ffi::kCFURLVolumeAvailableCapacityForImportantUsageKey,
]);
match get_disk_properties(&self.volume_url, &requested_properties) {
Some(disk_props) => {
self.available_space = get_available_volume_space(&disk_props);
true
]) {
match get_disk_properties(&self.volume_url, &requested_properties) {
Some(disk_props) => {
self.available_space = get_available_volume_space(&disk_props);
true
}
None => false,
}
None => false,
} else {
sysinfo_debug!("failed to create volume key list, skipping refresh");
false
}
}
}
Expand All @@ -97,7 +101,7 @@ pub(super) unsafe fn get_disks() -> Vec<Disk> {
};

// Create a list of properties about the disk that we want to fetch.
let requested_properties = build_requested_properties(&[
let requested_properties = match build_requested_properties(&[
ffi::kCFURLVolumeIsEjectableKey,
ffi::kCFURLVolumeIsRemovableKey,
ffi::kCFURLVolumeIsInternalKey,
Expand All @@ -107,19 +111,30 @@ pub(super) unsafe fn get_disks() -> Vec<Disk> {
ffi::kCFURLVolumeNameKey,
ffi::kCFURLVolumeIsBrowsableKey,
ffi::kCFURLVolumeIsLocalKey,
]);
]) {
Some(properties) => properties,
None => {
sysinfo_debug!("failed to create volume key list");
return Vec::new();
}
};

let mut disks = Vec::with_capacity(raw_disks.len());
for c_disk in raw_disks {
let volume_url = CFReleaser::new(
let volume_url = match CFReleaser::new(
core_foundation_sys::url::CFURLCreateFromFileSystemRepresentation(
kCFAllocatorDefault,
c_disk.f_mntonname.as_ptr() as *const _,
c_disk.f_mntonname.len() as _,
false as _,
),
)
.unwrap();
) {
Some(url) => url,
None => {
sysinfo_debug!("getfsstat returned incompatible paths");
continue;
}
};

let prop_dict = match get_disk_properties(&volume_url, &requested_properties) {
Some(props) => props,
Expand Down Expand Up @@ -173,14 +188,13 @@ type RetainedCFArray = CFReleaser<core_foundation_sys::array::__CFArray>;
type RetainedCFDictionary = CFReleaser<core_foundation_sys::dictionary::__CFDictionary>;
type RetainedCFURL = CFReleaser<core_foundation_sys::url::__CFURL>;

unsafe fn build_requested_properties(properties: &[CFStringRef]) -> RetainedCFArray {
unsafe fn build_requested_properties(properties: &[CFStringRef]) -> Option<RetainedCFArray> {
CFReleaser::new(CFArrayCreate(
ptr::null_mut(),
properties.as_ptr() as *const *const c_void,
properties.len() as _,
&core_foundation_sys::array::kCFTypeArrayCallBacks,
))
.unwrap()
}

fn get_disk_properties(
Expand Down Expand Up @@ -241,8 +255,7 @@ unsafe fn get_dict_value<T, F: FnOnce(*const c_void) -> Option<T>>(
cfs::kCFStringEncodingUTF8,
false as _,
core_foundation_sys::base::kCFAllocatorNull,
))
.unwrap();
))?;

_defined.inner()
}
Expand Down Expand Up @@ -315,7 +328,7 @@ unsafe fn new_disk(
// so we just assume the disk type is an SSD until Rust has a way to conditionally link to
// IOKit in more recent deployment versions.
#[cfg(target_os = "macos")]
let type_ = crate::sys::inner::disk::get_disk_type(&c_disk);
let type_ = crate::sys::inner::disk::get_disk_type(&c_disk).unwrap_or(DiskType::Unknown(-1));
#[cfg(not(target_os = "macos"))]
let type_ = DiskType::SSD;

Expand Down
50 changes: 26 additions & 24 deletions src/apple/macos/disk.rs
Expand Up @@ -13,34 +13,36 @@ use core_foundation_sys::string as cfs;

use std::ffi::CStr;

const UNKNOWN_DISK_TYPE: DiskType = DiskType::Unknown(-1);

pub(crate) fn get_disk_type(disk: &libc::statfs) -> DiskType {
let characteristics_string = CFReleaser::new(unsafe {
cfs::CFStringCreateWithBytesNoCopy(
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,
)
})
.unwrap();
))?
};

// Removes `/dev/` from the value.
let bsd_name = unsafe {
CStr::from_ptr(disk.f_mntfromname.as_ptr())
.to_bytes()
.strip_prefix(b"/dev/")
.expect("device mount point in unknown format")
.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()) };

if matching.is_null() {
return UNKNOWN_DISK_TYPE;
return None;
}

let mut service_iterator: ffi::io_iterator_t = 0;
Expand All @@ -53,10 +55,11 @@ pub(crate) fn get_disk_type(disk: &libc::statfs) -> DiskType {
)
} != libc::KERN_SUCCESS
{
return UNKNOWN_DISK_TYPE;
return None;
}

let service_iterator = IOReleaser::new(service_iterator).unwrap();
// 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) };

let mut parent_entry: ffi::io_registry_entry_t = 0;

Expand All @@ -68,23 +71,22 @@ pub(crate) fn get_disk_type(disk: &libc::statfs) -> DiskType {
// 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 {
let result = unsafe {
if unsafe {
ffi::IORegistryEntryGetParentEntry(
current_service_entry.inner(),
ffi::kIOServicePlane.as_ptr().cast(),
&mut parent_entry,
)
};
if result != libc::KERN_SUCCESS {
} != libc::KERN_SUCCESS
{
break;
}

current_service_entry = IOReleaser::new(parent_entry).unwrap();

// There were no more parents left.
if parent_entry == 0 {
break;
}
current_service_entry = match IOReleaser::new(parent_entry) {
Some(service) => service,
// There were no more parents left
None => break,
};

let properties_result = unsafe {
CFReleaser::new(ffi::IORegistryEntryCreateCFProperty(
Expand All @@ -108,17 +110,17 @@ pub(crate) fn get_disk_type(disk: &libc::statfs) -> DiskType {
_ if medium == ffi::kIOPropertyMediumTypeRotationalKey => Some(DiskType::HDD),
_ => None,
}) {
return disk_type;
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 DiskType::HDD;
return Some(DiskType::HDD);
}
}
}
}

UNKNOWN_DISK_TYPE
None
}
25 changes: 15 additions & 10 deletions src/apple/macos/utils.rs
@@ -1,25 +1,30 @@
// Take a look at the license at the top of the repository in the LICENSE file.

pub(crate) struct IOReleaser(super::ffi::io_object_t);
use std::num::NonZeroU32;

type IoObject = NonZeroU32;

pub(crate) struct IOReleaser(IoObject);

impl IOReleaser {
pub(crate) fn new(obj: u32) -> Option<Self> {
if obj == 0 {
None
} else {
Some(Self(obj))
}
IoObject::new(obj).map(Self)
}

pub(crate) unsafe fn new_unchecked(obj: u32) -> Self {
// Chance at catching in-development mistakes
debug_assert_ne!(obj, 0);
Self(IoObject::new_unchecked(obj))
}

#[inline]
pub(crate) fn inner(&self) -> u32 {
self.0
self.0.get()
}
}

impl Drop for IOReleaser {
fn drop(&mut self) {
if self.0 != 0 {
unsafe { super::ffi::IOObjectRelease(self.0 as _) };
}
unsafe { super::ffi::IOObjectRelease(self.0.get() as _) };
}
}
21 changes: 11 additions & 10 deletions src/apple/utils.rs
Expand Up @@ -2,35 +2,36 @@

use core_foundation_sys::base::CFRelease;
use libc::c_char;
use std::ptr::NonNull;

// A helper using to auto release the resource got from CoreFoundation.
// More information about the ownership policy for CoreFoundation pelease refer the link below:
// https://developer.apple.com/library/archive/documentation/CoreFoundation/Conceptual/CFMemoryMgmt/Concepts/Ownership.html#//apple_ref/doc/uid/20001148-CJBEJBHH
#[repr(transparent)]
pub(crate) struct CFReleaser<T>(*const T);
pub(crate) struct CFReleaser<T>(NonNull<T>);

impl<T> CFReleaser<T> {
pub(crate) fn new(ptr: *const T) -> Option<Self> {
if ptr.is_null() {
None
} else {
Some(Self(ptr))
}
// This cast is OK because `NonNull` is a transparent wrapper
// over a `*const T`. Additionally, mutability doesn't matter with
// pointers here.
NonNull::new(ptr as *mut T).map(Self)
}

pub(crate) fn inner(&self) -> *const T {
self.0
self.0.as_ptr().cast()
}
}

impl<T> Drop for CFReleaser<T> {
fn drop(&mut self) {
if !self.0.is_null() {
unsafe { CFRelease(self.0 as _) }
}
unsafe { CFRelease(self.0.as_ptr().cast()) }
}
}

// Safety: These are safe to implement because we only wrap non-mutable
// CoreFoundation types, which are generally threadsafe unless noted
// otherwise.
unsafe impl<T> Send for CFReleaser<T> {}
unsafe impl<T> Sync for CFReleaser<T> {}

Expand Down

0 comments on commit 13c32c9

Please sign in to comment.