vamp/src/battery.rs

use std::{
    io,
    sync::{
        Arc,
        atomic::{AtomicU8, Ordering},
        mpsc::{self, Receiver, Sender},
    },
    thread::{self, JoinHandle},
};
#[cfg(not(test))]
use std::{sync::mpsc::RecvTimeoutError, time::Duration};

pub(crate) struct Battery {
    results: Arc<AtomicU8>,
    wake: Option<Sender<()>>,
    thread: Option<JoinHandle<()>>,
}

impl Battery {
    #[cfg(not(test))]
    pub(crate) fn start() -> io::Result<Self> {
        Self::start_with(read_charge, |wake| {
            matches!(
                wake.recv_timeout(Duration::from_secs(30)),
                Ok(()) | Err(RecvTimeoutError::Timeout)
            )
        })
    }

    fn start_with(
        mut read: impl FnMut() -> Option<u8> + Send + 'static,
        mut wait: impl FnMut(&Receiver<()>) -> bool + Send + 'static,
    ) -> io::Result<Self> {
        let results = Arc::new(AtomicU8::new(0));
        let latest = Arc::clone(&results);
        let (wake, requests) = mpsc::channel();
        let thread = thread::Builder::new()
            .name("vamp-battery".into())
            .spawn(move || {
                loop {
                    let charge = read().filter(|charge| *charge <= 100);
                    // A one-slot latest-result mailbox: 0 = no new sample,
                    // 1 = no battery, 2..=102 = charge. Publishing never blocks.
                    latest.store(charge.map_or(1, |charge| charge + 2), Ordering::Relaxed);
                    if !wait(&requests) {
                        break;
                    }
                }
            })?;
        Ok(Self {
            results,
            wake: Some(wake),
            thread: Some(thread),
        })
    }

    pub(crate) fn poll(&mut self) -> Option<Option<u8>> {
        match self.results.swap(0, Ordering::Relaxed) {
            0 => None,
            1 => Some(None),
            charge => Some(Some(charge - 2)),
        }
    }
}

impl Drop for Battery {
    fn drop(&mut self) {
        // Disconnecting interrupts recv_timeout instead of waiting for its deadline.
        self.wake.take();
        if let Some(thread) = self.thread.take() {
            let _ = thread.join();
        }
    }
}

#[cfg(not(test))]
pub(crate) fn read_charge() -> Option<u8> {
    // ponytail: first internal battery on macOS/Linux; aggregate compatible capacities for multi-battery systems.
    #[cfg(target_os = "macos")]
    return macos::read_charge();
    #[cfg(target_os = "linux")]
    return linux::read_charge(std::path::Path::new("/sys/class/power_supply"));
    #[cfg(windows)]
    return windows::read_charge();
    #[cfg(not(any(target_os = "macos", target_os = "linux", windows)))]
    None
}

#[cfg(any(target_os = "macos", target_os = "linux", test))]
fn percentage(current: u64, full: u64) -> Option<u8> {
    if full == 0 || current > full {
        return None;
    }
    Some(((current as u128 * 100) / full as u128) as u8)
}

#[cfg(target_os = "macos")]
mod macos {
    use std::ffi::{CStr, c_char, c_void};

    type Object = *const c_void;

    #[cfg(not(test))]
    #[link(name = "IOKit", kind = "framework")]
    unsafe extern "C" {
        fn IOPSCopyPowerSourcesInfo() -> Object;
        fn IOPSCopyPowerSourcesList(info: Object) -> Object;
        fn IOPSGetPowerSourceDescription(info: Object, source: Object) -> Object;
    }

    #[link(name = "CoreFoundation", kind = "framework")]
    unsafe extern "C" {
        fn CFRelease(object: Object);
        fn CFGetTypeID(object: Object) -> usize;
        #[cfg(not(test))]
        fn CFEqual(left: Object, right: Object) -> u8;
        #[cfg(not(test))]
        fn CFArrayGetTypeID() -> usize;
        #[cfg(not(test))]
        fn CFArrayGetCount(array: Object) -> isize;
        #[cfg(not(test))]
        fn CFArrayGetValueAtIndex(array: Object, index: isize) -> Object;
        #[cfg(not(test))]
        fn CFDictionaryGetTypeID() -> usize;
        #[cfg(not(test))]
        fn CFDictionaryGetValue(dictionary: Object, key: Object) -> Object;
        #[cfg(not(test))]
        fn CFStringGetTypeID() -> usize;
        fn CFStringCreateWithCString(
            allocator: Object,
            string: *const c_char,
            encoding: u32,
        ) -> Object;
        #[cfg(not(test))]
        fn CFBooleanGetTypeID() -> usize;
        #[cfg(not(test))]
        fn CFBooleanGetValue(boolean: Object) -> u8;
        fn CFNumberGetTypeID() -> usize;
        fn CFNumberGetValue(number: Object, kind: isize, value: *mut c_void) -> u8;
    }

    // Only Copy/Create results belong here; power source descriptions are borrowed.
    struct Owned(Object);

    impl Owned {
        fn new(object: Object) -> Option<Self> {
            (!object.is_null()).then(|| Self(object))
        }

        fn string(value: &CStr) -> Option<Self> {
            // kCFStringEncodingUTF8, from CFString.h; a null allocator is the default.
            Self::new(unsafe {
                CFStringCreateWithCString(std::ptr::null(), value.as_ptr(), 0x0800_0100)
            })
        }
    }

    impl Drop for Owned {
        fn drop(&mut self) {
            unsafe { CFRelease(self.0) };
        }
    }

    unsafe fn is_type(object: Object, kind: usize) -> bool {
        !object.is_null() && unsafe { CFGetTypeID(object) == kind }
    }

    unsafe fn number(object: Object) -> Option<u64> {
        if !unsafe { is_type(object, CFNumberGetTypeID()) } {
            return None;
        }
        let mut value: i64 = 0;
        // kCFNumberSInt64Type = 4. A lossy/out-of-range conversion returns false.
        let converted = unsafe { CFNumberGetValue(object, 4, (&mut value as *mut i64).cast()) };
        (converted != 0).then_some(value)?.try_into().ok()
    }

    #[cfg(not(test))]
    pub(super) fn read_charge() -> Option<u8> {
        // The opaque snapshot is only passed to IOPS accessors. All owned values
        // are released on every return path, including failed type checks.
        unsafe {
            let info = Owned::new(IOPSCopyPowerSourcesInfo())?;
            let sources = Owned::new(IOPSCopyPowerSourcesList(info.0))?;
            if !is_type(sources.0, CFArrayGetTypeID()) {
                return None;
            }
            let type_key = Owned::string(c"Type")?;
            let internal = Owned::string(c"InternalBattery")?;
            let present_key = Owned::string(c"Is Present")?;
            let current_key = Owned::string(c"Current Capacity")?;
            let full_key = Owned::string(c"Max Capacity")?;
            for index in 0..CFArrayGetCount(sources.0) {
                let source = CFArrayGetValueAtIndex(sources.0, index);
                if source.is_null() {
                    continue;
                }
                let description = IOPSGetPowerSourceDescription(info.0, source);
                if !is_type(description, CFDictionaryGetTypeID()) {
                    continue;
                }
                let kind = CFDictionaryGetValue(description, type_key.0);
                if !is_type(kind, CFStringGetTypeID()) || CFEqual(kind, internal.0) == 0 {
                    continue;
                }
                let present = CFDictionaryGetValue(description, present_key.0);
                if !is_type(present, CFBooleanGetTypeID()) || CFBooleanGetValue(present) == 0 {
                    continue;
                }
                let current = number(CFDictionaryGetValue(description, current_key.0))?;
                let full = number(CFDictionaryGetValue(description, full_key.0))?;
                return super::percentage(current, full);
            }
            None
        }
    }

    #[cfg(test)]
    mod tests {
        include!("../tests/unit/battery_macos.rs");
    }
}

#[cfg(any(target_os = "linux", test))]
mod linux {
    use std::{fs, io, path::Path};

    pub(super) fn read_charge(root: &Path) -> Option<u8> {
        for entry in fs::read_dir(root).ok()? {
            let path = entry.ok()?.path();
            if fs::read_to_string(path.join("type"))
                .ok()
                .as_deref()
                .map(str::trim)
                != Some("Battery")
            {
                continue;
            }
            match fs::read_to_string(path.join("scope")) {
                Ok(scope) if matches!(scope.trim(), "System" | "Unknown") => {}
                Err(error) if error.kind() == io::ErrorKind::NotFound => {}
                _ => continue,
            }
            if fs::read_to_string(path.join("present"))
                .ok()
                .as_deref()
                .map(str::trim)
                != Some("1")
            {
                continue;
            }
            return read_capacity(&path);
        }
        None
    }

    fn read_capacity(path: &Path) -> Option<u8> {
        match fs::read_to_string(path.join("capacity")) {
            Ok(capacity) => {
                return capacity
                    .trim()
                    .parse::<u8>()
                    .ok()
                    .filter(|value| *value <= 100);
            }
            Err(error) if error.kind() == io::ErrorKind::NotFound => {}
            Err(_) => return None,
        }
        // Never mix energy and charge units or use design capacity as full capacity.
        let (current, full) = match fs::read_to_string(path.join("energy_now")) {
            Ok(current) => (current, fs::read_to_string(path.join("energy_full")).ok()?),
            Err(error) if error.kind() == io::ErrorKind::NotFound => (
                fs::read_to_string(path.join("charge_now")).ok()?,
                fs::read_to_string(path.join("charge_full")).ok()?,
            ),
            Err(_) => return None,
        };
        super::percentage(current.trim().parse().ok()?, full.trim().parse().ok()?)
    }
}

#[cfg(any(windows, test))]
fn windows_charge(percent: u8, flags: u8) -> Option<u8> {
    (percent <= 100 && flags & 128 == 0).then_some(percent)
}

#[cfg(all(windows, not(test)))]
mod windows {
    #[repr(C)]
    struct SystemPowerStatus {
        ac_line_status: u8,
        battery_flag: u8,
        battery_life_percent: u8,
        system_status_flag: u8,
        battery_life_time: u32,
        battery_full_life_time: u32,
    }

    #[link(name = "kernel32")]
    unsafe extern "system" {
        fn GetSystemPowerStatus(status: *mut SystemPowerStatus) -> i32;
    }

    pub(super) fn read_charge() -> Option<u8> {
        let mut status = std::mem::MaybeUninit::uninit();
        if unsafe { GetSystemPowerStatus(status.as_mut_ptr()) } == 0 {
            return None;
        }
        let status = unsafe { status.assume_init() };
        super::windows_charge(status.battery_life_percent, status.battery_flag)
    }
}

#[cfg(test)]
#[path = "../tests/unit/battery.rs"]
mod tests;