macos: Wrapper for IOKit iterator

src/platform/macos_iokit/enumeration.rs

@@ -1,4 +1,4 @@
-use std::{io::ErrorKind, iter};
+use std::io::ErrorKind;
 
 use core_foundation::{
     base::{CFType, TCFType},
@@ -8,18 +8,16 @@ use core_foundation::{
 };
 use io_kit_sys::{
     kIOMasterPortDefault, kIORegistryIterateParents, kIORegistryIterateRecursively,
-    keys::kIOServicePlane,
-    ret::kIOReturnSuccess,
-    types::{io_iterator_t, io_object_t},
-    usb::lib::kIOUSBDeviceClassName,
-    IOIteratorNext, IOObjectRelease, IORegistryEntrySearchCFProperty, IOServiceGetMatchingServices,
-    IOServiceMatching,
+    keys::kIOServicePlane, ret::kIOReturnSuccess, usb::lib::kIOUSBDeviceClassName,
+    IORegistryEntrySearchCFProperty, IOServiceGetMatchingServices, IOServiceMatching,
 };
 use log::{error, info};
 
 use crate::{DeviceInfo, Error, Speed};
 
-pub fn list_devices() -> Result<impl Iterator<Item = DeviceInfo>, Error> {
+use super::iokit::{IoService, IoServiceIterator};
+
+fn usb_service_iter() -> Result<IoServiceIterator, Error> {
     unsafe {
         let dictionary = IOServiceMatching(kIOUSBDeviceClassName);
         if dictionary.is_null() {
@@ -32,52 +30,42 @@ pub fn list_devices() -> Result<impl Iterator<Item = DeviceInfo>, Error> {
             return Err(Error::from_raw_os_error(r));
         }
 
-        let devices: Vec<DeviceInfo> = iter::from_fn(|| {
-            let device = IOIteratorNext(iterator);
-            if device != 0 {
-                let dev = probe_device(device);
-                IOObjectRelease(device);
-                Some(dev)
-            } else {
-                None
-            }
-        })
-        .filter_map(|d| d)
-        .collect();
-        IOObjectRelease(iterator);
-
-        Ok(devices.into_iter())
+        Ok(IoServiceIterator::new(iterator))
     }
 }
 
-fn probe_device(device: io_object_t) -> Option<DeviceInfo> {
+pub fn list_devices() -> Result<impl Iterator<Item = DeviceInfo>, Error> {
+    Ok(usb_service_iter()?.filter_map(probe_device))
+}
+
+fn probe_device(device: IoService) -> Option<DeviceInfo> {
     // Can run `ioreg -p IOUSB -l` to see all properties
-    let location_id: u32 = get_integer_property(device, "locationID")?;
+    let location_id: u32 = get_integer_property(&device, "locationID")?;
     log::info!("Probing device {location_id}");
 
     Some(DeviceInfo {
         location_id,
         bus_number: 0, // TODO: does this exist on macOS?
-        device_address: get_integer_property(device, "USB Address")?,
-        vendor_id: get_integer_property(device, "idVendor")?,
-        product_id: get_integer_property(device, "idProduct")?,
-        device_version: get_integer_property(device, "bcdDevice")?,
-        class: get_integer_property(device, "bDeviceClass")?,
-        subclass: get_integer_property(device, "bDeviceSubClass")?,
-        protocol: get_integer_property(device, "bDeviceProtocol")?,
-        speed: get_integer_property(device, "Device Speed").and_then(map_speed),
-        manufacturer_string: get_string_property(device, "USB Vendor Name"),
-        product_string: get_string_property(device, "USB Product Name"),
-        serial_number: get_string_property(device, "USB Serial Number"),
+        device_address: get_integer_property(&device, "USB Address")?,
+        vendor_id: get_integer_property(&device, "idVendor")?,
+        product_id: get_integer_property(&device, "idProduct")?,
+        device_version: get_integer_property(&device, "bcdDevice")?,
+        class: get_integer_property(&device, "bDeviceClass")?,
+        subclass: get_integer_property(&device, "bDeviceSubClass")?,
+        protocol: get_integer_property(&device, "bDeviceProtocol")?,
+        speed: get_integer_property(&device, "Device Speed").and_then(map_speed),
+        manufacturer_string: get_string_property(&device, "USB Vendor Name"),
+        product_string: get_string_property(&device, "USB Product Name"),
+        serial_number: get_string_property(&device, "USB Serial Number"),
     })
 }
 
-fn get_property<T: ConcreteCFType>(device: io_iterator_t, property: &'static str) -> Option<T> {
+fn get_property<T: ConcreteCFType>(device: &IoService, property: &'static str) -> Option<T> {
     unsafe {
         let cf_property = CFString::from_static_string(property);
 
         let raw = IORegistryEntrySearchCFProperty(
-            device,
+            device.get(),
             kIOServicePlane as *mut i8,
             cf_property.as_CFTypeRef() as *const _,
             std::ptr::null(),
@@ -99,14 +87,11 @@ fn get_property<T: ConcreteCFType>(device: io_iterator_t, property: &'static str
     }
 }
 
-fn get_string_property(device: io_iterator_t, property: &'static str) -> Option<String> {
+fn get_string_property(device: &IoService, property: &'static str) -> Option<String> {
     get_property::<CFString>(device, property).map(|s| s.to_string())
 }
 
-fn get_integer_property<T: TryFrom<i64>>(
-    device: io_iterator_t,
-    property: &'static str,
-) -> Option<T> {
+fn get_integer_property<T: TryFrom<i64>>(device: &IoService, property: &'static str) -> Option<T> {
     get_property::<CFNumber>(device, property)
         .and_then(|n| n.to_i64())
         .and_then(|n| n.try_into().ok())

src/platform/macos_iokit/iokit.rs

@@ -0,0 +1,57 @@
+use io_kit_sys::{IOIteratorNext, IOObjectRelease};
+
+pub(crate) struct IoObject(u32);
+
+impl IoObject {
+    // Safety: `handle` must be an IOObject handle. Ownership is transferred.
+    pub unsafe fn new(handle: u32) -> IoObject {
+        IoObject(handle)
+    }
+    pub fn get(&self) -> u32 {
+        self.0
+    }
+}
+
+impl Drop for IoObject {
+    fn drop(&mut self) {
+        unsafe {
+            IOObjectRelease(self.0);
+        }
+    }
+}
+
+pub(crate) struct IoService(IoObject);
+
+impl IoService {
+    // Safety: `handle` must be an IOService handle. Ownership is transferred.
+    pub unsafe fn new(handle: u32) -> IoService {
+        IoService(IoObject(handle))
+    }
+    pub fn get(&self) -> u32 {
+        self.0 .0
+    }
+}
+
+pub(crate) struct IoServiceIterator(IoObject);
+
+impl IoServiceIterator {
+    // Safety: `handle` must be an IoIterator of IoService. Ownership is transferred.
+    pub unsafe fn new(handle: u32) -> IoServiceIterator {
+        IoServiceIterator(IoObject::new(handle))
+    }
+}
+
+impl Iterator for IoServiceIterator {
+    type Item = IoService;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        unsafe {
+            let handle = IOIteratorNext(self.0.get());
+            if handle != 0 {
+                Some(IoService::new(handle))
+            } else {
+                None
+            }
+        }
+    }
+}

src/platform/macos_iokit/mod.rs

@@ -8,3 +8,5 @@ pub use enumeration::list_devices;
 mod device;
 pub(crate) use device::MacDevice as Device;
 pub(crate) use device::MacInterface as Interface;
+
+mod iokit;