blubat-core 0.4.0

Bluetooth battery model, macOS data sources and polling engine behind the blubat CLI
//! The native IOKit source: Apple HID peripherals.
//!
//! Apple's own peripherals report a battery level here and nowhere else, so
//! this is both the only source for them and, at well under a millisecond, the
//! one cheap enough to sit on a poll tick.

use std::collections::HashMap;
use std::ffi::CString;

use objc2_core_foundation::{CFBoolean, CFDictionary, CFNumber, CFRetained, CFString, CFType};
use objc2_io_kit::{
    IOIteratorNext, IOObjectRelease, IORegistryEntryCreateCFProperty, IOServiceGetMatchingServices,
    IOServiceMatching, io_iterator_t, io_object_t, kIOMainPortDefault,
};

use crate::address::Address;
use crate::config::{Advertised, AdvertisedThresholds};
use crate::device::{ChargeState, Device, Levels, Source};
use crate::timestamp::Timestamp;

/// The class Apple's HID peripherals register under.
///
/// It bounds the iteration only. `BatteryPercent` is what decides whether an
/// entry counts as a reading, so a driver that carries the key but not the
/// class is a matter of widening this constant, not of reworking the parse.
/// [`crate::presence`] watches the same class, so what is read and what is
/// noticed arriving cannot come apart.
pub(crate) const SERVICE_CLASS: &str = "AppleDeviceManagementHIDEventService";

const KEYS: [&str; 8] = [
    "Product",
    "DeviceAddress",
    "BatteryPercent",
    "BatteryStatusFlags",
    "HasBattery",
    "Transport",
    "LowBatteryNotificationPercentage",
    "CriticallyLowBatteryNotificationPercentage",
];

/// Reads every Apple HID peripheral that reports a battery.
pub(crate) fn read(read_at: Timestamp, warnings: &mut Vec<String>) -> Vec<Device> {
    matching_entries()
        .into_iter()
        .filter_map(|properties| device(&properties, read_at, warnings))
        .collect()
}

/// Reads the thresholds Apple's own nodes publish for themselves.
///
/// A separate pass from a reading, and deliberately not on the poll tick: these
/// numbers describe the device rather than its charge and never move, so a
/// caller resolves them once per run. Devices that publish neither are absent
/// rather than empty, which keeps the map to what it is for.
pub fn advertised() -> AdvertisedThresholds {
    matching_entries()
        .iter()
        .filter_map(advertised_by)
        .collect()
}

/// The thresholds one registry entry publishes, absent where it publishes none.
fn advertised_by(properties: &Properties) -> Option<(Address, Advertised)> {
    let percentage = |key| {
        properties
            .get(key)
            .and_then(Property::number)
            .and_then(|percent| u8::try_from(percent).ok())
            .filter(|&percent| percent <= 100)
    };

    let address = properties
        .get("DeviceAddress")
        .and_then(Property::text)
        .and_then(Address::parse)?;
    let advertised = Advertised {
        low: percentage("LowBatteryNotificationPercentage"),
        critical: percentage("CriticallyLowBatteryNotificationPercentage"),
    };

    (advertised != Advertised::NONE).then_some((address, advertised))
}

/// One registry property, narrowed to the three types these keys use.
enum Property {
    Text(String),
    Number(i64),
    Flag(bool),
}

impl Property {
    fn text(&self) -> Option<&str> {
        match self {
            Property::Text(text) => Some(text),
            _ => None,
        }
    }

    fn number(&self) -> Option<i64> {
        match self {
            Property::Number(number) => Some(*number),
            _ => None,
        }
    }

    /// Whether the property denies its key, as either a boolean or a zero.
    ///
    /// IOKit is free to hand back a flag as a CFBoolean or as a CFNumber, and
    /// a denial has to be honoured whichever shape it arrives in.
    fn is_false(&self) -> bool {
        matches!(self, Property::Flag(false) | Property::Number(0))
    }
}

type Properties = HashMap<&'static str, Property>;

fn device(
    properties: &Properties,
    read_at: Timestamp,
    warnings: &mut Vec<String>,
) -> Option<Device> {
    let percent = properties
        .get("BatteryPercent")
        .and_then(Property::number)?;
    if properties.get("HasBattery").is_some_and(Property::is_false) {
        return None;
    }

    let address = properties
        .get("DeviceAddress")
        .and_then(Property::text)
        .and_then(Address::parse);
    let Some(address) = address else {
        warnings.push("skipping an IOKit battery reading with no usable DeviceAddress".to_string());
        return None;
    };

    let name = properties
        .get("Product")
        .and_then(Property::text)
        .map(str::to_string)
        .unwrap_or_else(|| address.to_string());

    Some(Device {
        address,
        name,
        // The device category comes from the other source, which names it better.
        kind: None,
        transport: properties
            .get("Transport")
            .and_then(Property::text)
            .map(str::to_string),
        levels: Levels {
            main: u8::try_from(percent).ok().filter(|&level| level <= 100),
            ..Levels::default()
        },
        charge: properties
            .get("BatteryStatusFlags")
            .and_then(Property::number)
            .map_or(ChargeState::Unknown, ChargeState::from_status_flags),
        source: Source::IoKit,
        // The registry lists a device only while it is present.
        connected: true,
        read_at,
    })
}

/// Runs one pass over the matching services, releasing every handle it takes.
///
/// Each `unsafe` wraps the single call it vouches for, so the loop and the
/// pushes stay plainly safe code.
fn matching_entries() -> Vec<Properties> {
    let class = CString::new(SERVICE_CLASS).expect("class name has no interior nul");
    let mut entries = Vec::new();

    let Some(matching) = (unsafe { IOServiceMatching(class.as_ptr()) }) else {
        return entries;
    };
    // IOServiceMatching hands back the mutable subtype and the getter wants the
    // immutable one. Same object, so the reinterpret is sound.
    let matching: CFRetained<CFDictionary> = unsafe { CFRetained::cast_unchecked(matching) };

    let mut iterator: io_iterator_t = 0;
    let result =
        unsafe { IOServiceGetMatchingServices(kIOMainPortDefault, Some(matching), &mut iterator) };
    if result != 0 || iterator == 0 {
        return entries;
    }

    loop {
        let entry = IOIteratorNext(iterator);
        if entry == 0 {
            break;
        }

        entries.push(read_properties(entry));
        IOObjectRelease(entry);
    }
    IOObjectRelease(iterator);

    entries
}

fn read_properties(entry: io_object_t) -> Properties {
    KEYS.iter()
        .filter_map(|&key| read_property(entry, key).map(|value| (key, value)))
        .collect()
}

/// Copies one property off a registry entry.
///
/// `CFRetained` owns the create-rule reference, so the release is its `Drop`
/// and every type test is a checked downcast.
fn read_property(entry: io_object_t, key: &str) -> Option<Property> {
    let key = CFString::from_str(key);
    let value: CFRetained<CFType> =
        unsafe { IORegistryEntryCreateCFProperty(entry, Some(&key), None, 0) }?;

    value
        .downcast_ref::<CFString>()
        .map(|text| Property::Text(text.to_string()))
        .or_else(|| {
            value
                .downcast_ref::<CFNumber>()
                .and_then(CFNumber::as_i64)
                .map(Property::Number)
        })
        .or_else(|| {
            value
                .downcast_ref::<CFBoolean>()
                .map(|flag| Property::Flag(flag.value()))
        })
}

#[cfg(test)]
mod tests {
    use super::*;

    fn properties(pairs: Vec<(&'static str, Property)>) -> Properties {
        pairs.into_iter().collect()
    }

    fn trackpad() -> Properties {
        properties(vec![
            (
                "Product",
                Property::Text("Paul\u{2019}s Magic Trackpad".to_string()),
            ),
            (
                "DeviceAddress",
                Property::Text("30-82-16-f2-24-90".to_string()),
            ),
            ("BatteryPercent", Property::Number(85)),
            ("BatteryStatusFlags", Property::Number(0)),
            ("HasBattery", Property::Flag(true)),
            ("Transport", Property::Text("Bluetooth".to_string())),
        ])
    }

    #[test]
    fn builds_a_device_from_every_key_it_reads() {
        let device = device(&trackpad(), Timestamp::from_unix(0), &mut Vec::new())
            .expect("a battery reading");

        assert_eq!(device.name, "Paul\u{2019}s Magic Trackpad");
        assert_eq!(device.address.as_str(), "30-82-16-f2-24-90");
        assert_eq!(device.levels.main, Some(85));
        assert_eq!(device.charge, ChargeState::Discharging);
        assert_eq!(device.transport.as_deref(), Some("Bluetooth"));
        assert_eq!(device.source, Source::IoKit);
        assert!(device.connected);
    }

    #[test]
    fn a_plugged_in_device_reads_as_charging() {
        let mut charging = trackpad();
        charging.insert("BatteryStatusFlags", Property::Number(3));

        let device =
            device(&charging, Timestamp::from_unix(0), &mut Vec::new()).expect("a battery reading");
        assert_eq!(device.charge, ChargeState::Charging);
    }

    #[test]
    fn charge_state_is_unknown_without_the_status_flags() {
        let mut no_flags = trackpad();
        no_flags.remove("BatteryStatusFlags");

        let device =
            device(&no_flags, Timestamp::from_unix(0), &mut Vec::new()).expect("a battery reading");
        assert_eq!(device.charge, ChargeState::Unknown);
    }

    #[test]
    fn an_entry_without_a_battery_percent_is_not_a_reading() {
        let mut no_battery = trackpad();
        no_battery.remove("BatteryPercent");

        assert!(device(&no_battery, Timestamp::from_unix(0), &mut Vec::new()).is_none());
    }

    #[test]
    fn has_battery_false_overrules_a_reported_percentage() {
        let mut denied = trackpad();
        denied.insert("HasBattery", Property::Flag(false));

        assert!(device(&denied, Timestamp::from_unix(0), &mut Vec::new()).is_none());
    }

    #[test]
    fn has_battery_denies_a_reading_as_a_number_as_well_as_a_boolean() {
        let mut denied = trackpad();
        denied.insert("HasBattery", Property::Number(0));
        assert!(device(&denied, Timestamp::from_unix(0), &mut Vec::new()).is_none());

        let mut allowed = trackpad();
        allowed.insert("HasBattery", Property::Number(1));
        assert!(device(&allowed, Timestamp::from_unix(0), &mut Vec::new()).is_some());
    }

    #[test]
    fn an_unusable_address_drops_the_entry() {
        let mut bad_address = trackpad();
        bad_address.insert("DeviceAddress", Property::Text("nonsense".to_string()));
        let mut warnings = Vec::new();

        assert!(device(&bad_address, Timestamp::from_unix(0), &mut warnings).is_none());
        assert_eq!(warnings.len(), 1, "returned rather than printed");
    }

    #[test]
    fn the_address_stands_in_for_a_missing_product_name() {
        let mut anonymous = trackpad();
        anonymous.remove("Product");

        let device = device(&anonymous, Timestamp::from_unix(0), &mut Vec::new())
            .expect("a battery reading");
        assert_eq!(device.name, "30-82-16-f2-24-90");
    }

    #[test]
    fn apples_own_thresholds_are_read_where_the_node_publishes_them() {
        let mut publishing = trackpad();
        publishing.insert("LowBatteryNotificationPercentage", Property::Number(6));
        publishing.insert(
            "CriticallyLowBatteryNotificationPercentage",
            Property::Number(3),
        );

        let (address, advertised) = advertised_by(&publishing).expect("both thresholds");

        assert_eq!(address.as_str(), "30-82-16-f2-24-90");
        assert_eq!(advertised.low, Some(6));
        assert_eq!(advertised.critical, Some(3));
    }

    #[test]
    fn a_node_publishing_one_threshold_contributes_only_that_one() {
        let mut half = trackpad();
        half.insert("LowBatteryNotificationPercentage", Property::Number(6));

        let (_, advertised) = advertised_by(&half).expect("the one it publishes");

        assert_eq!(advertised.low, Some(6));
        assert_eq!(advertised.critical, None);
    }

    #[test]
    fn a_node_publishing_nothing_usable_contributes_nothing() {
        let mut unusable = trackpad();
        unusable.insert("LowBatteryNotificationPercentage", Property::Number(101));

        assert_eq!(advertised_by(&trackpad()), None, "no thresholds at all");
        assert_eq!(advertised_by(&unusable), None, "not a percentage");

        let mut anonymous = trackpad();
        anonymous.insert("LowBatteryNotificationPercentage", Property::Number(6));
        anonymous.remove("DeviceAddress");
        assert_eq!(advertised_by(&anonymous), None, "nothing to key it on");
    }

    #[test]
    fn an_out_of_range_percentage_leaves_the_level_absent() {
        for percent in [-1, 101, 255] {
            let mut odd = trackpad();
            odd.insert("BatteryPercent", Property::Number(percent));

            let device =
                device(&odd, Timestamp::from_unix(0), &mut Vec::new()).expect("a battery reading");
            assert_eq!(device.levels.main, None, "at {percent}");
        }
    }
}