irbt/

lib.rs

//! Host-side client for the ir-bt-bridge USB configuration/control protocol.
//!
//! ```no_run
//! # fn main() -> Result<(), irbt::Error> {
//! let device = irbt::Device::open(None)?; // auto-detect the bridge
//! device.set_mapping(0x00FF_02FD, irbt::HidKey::Home)?;
//! for mapping in device.keymap()? {
//!     println!("0x{:08X} -> {}", mapping.ir, mapping.key);
//! }
//! # Ok(())
//! # }
//! ```

use std::collections::HashMap;
use std::io::{self, Read, Write};
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::{mpsc, Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};

use irbt_proto::{Command, Decoded, Decoder, Frame, Reply, Request};

pub use irbt_proto as proto;
pub use irbt_proto::{
    Battery, BuildInfo, Component, DeviceInfo, Event, HidKey, IrCode, Log, LogLevel, Mapping,
    Status, PROTOCOL_VERSION,
};

/// Espressif's USB VID and the PID of the ESP32-S3's built-in USB-Serial-JTAG.
pub const USB_VID: u16 = 0x303A;
pub const USB_PID: u16 = 0x1001;

/// How long to wait for a response to a single request.
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(2);

/// Serial read timeout; doubles as the idle period after which buffered
/// non-frame bytes are delivered as [`Notification::Raw`].
const READ_TIMEOUT: Duration = Duration::from_millis(50);

/// Unread notifications beyond this are dropped rather than buffered forever.
const NOTIFICATION_BACKLOG: usize = 1024;

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("I/O error: {0}")]
    Io(#[from] io::Error),
    #[error("serial port error: {0}")]
    Serial(#[from] serialport::Error),
    #[error("failed to encode request: {0}")]
    Encode(irbt_proto::EncodeError),
    #[error("no ir-bt-bridge found (looking for USB {USB_VID:04x}:{USB_PID:04x}); pass a port explicitly")]
    NoDevice,
    #[error("several candidate devices found ({}); pass a port explicitly", .0.join(", "))]
    MultipleDevices(Vec<String>),
    #[error("device speaks protocol v{device}, this tool speaks v{host}; update the firmware or the tool")]
    ProtocolMismatch { device: u16, host: u16 },
    #[error("timed out waiting for the device")]
    Timeout,
    #[error("device disconnected")]
    Disconnected,
    #[error("device error: {0}")]
    Device(String),
    #[error("unexpected reply from device: {0:?}")]
    UnexpectedReply(Reply),
}

pub type Result<T, E = Error> = std::result::Result<T, E>;

/// Something the device sent without being asked.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Notification {
    Event(Event),
    Log(Log),
    /// Output that isn't framed: ROM bootloader text, panic dumps, or a
    /// corrupted frame. Decoded lossily as UTF-8.
    Raw(String),
}

type Pending = Arc<Mutex<HashMap<u32, mpsc::Sender<Result<Reply, String>>>>>;

/// A connection to one bridge.
pub struct Device {
    writer: Mutex<Box<dyn Write + Send>>,
    pending: Pending,
    notifications: Mutex<mpsc::Receiver<Notification>>,
    alive: Arc<AtomicBool>,
    next_id: AtomicU32,
    timeout: Duration,
}

/// Serial port paths of every connected device that looks like a bridge.
///
/// # Errors
///
/// [`Error::Serial`] if the serial ports can't be enumerated.
pub fn find_ports() -> Result<Vec<String>> {
    Ok(serialport::available_ports()?
        .into_iter()
        .filter(|port| {
            matches!(&port.port_type, serialport::SerialPortType::UsbPort(usb)
                if usb.vid == USB_VID && usb.pid == USB_PID)
        })
        .map(|port| port.port_name)
        .collect())
}

fn find_single_port() -> Result<String, Error> {
    let mut ports = find_ports()?;
    if ports.len() > 1 {
        return Err(Error::MultipleDevices(ports));
    }
    ports.pop().ok_or(Error::NoDevice)
}

impl Device {
    /// Open the bridge at `port`, or auto-detect it when `None`, and check
    /// that it speaks this library's protocol version.
    ///
    /// # Errors
    ///
    /// [`Error::NoDevice`] or [`Error::MultipleDevices`] when auto-detecting
    /// doesn't find exactly one bridge, [`Error::Serial`]/[`Error::Io`] if
    /// the port can't be opened or configured,
    /// [`Error::ProtocolMismatch`] if the firmware speaks another protocol
    /// version, and any error from [`Device::info`].
    pub fn open(port: Option<&str>) -> Result<Self> {
        let path = match port {
            Some(path) => path.to_owned(),
            None => find_single_port()?,
        };

        let mut port = serialport::new(&path, 115_200)
            .timeout(READ_TIMEOUT)
            .open()?;
        // The USB-Serial-JTAG resets the chip when RTS is asserted while DTR
        // is not (the classic auto-reset circuit, emulated). Drop RTS first so
        // that state is never passed through.
        port.write_request_to_send(false)?;
        port.write_data_terminal_ready(false)?;

        let reader = port.try_clone()?;
        let device = Self::from_io(reader, port);
        device.handshake()?;
        Ok(device)
    }

    /// Talk to a device over an arbitrary byte stream (tests, TCP bridges, …).
    /// No handshake is performed; call [`Device::info`] to check the version.
    ///
    /// Reads returning [`io::ErrorKind::TimedOut`] or
    /// [`io::ErrorKind::WouldBlock`] are treated as idle periods; EOF or any
    /// other error ends the connection.
    ///
    /// # Panics
    ///
    /// If the OS can't spawn the reader thread.
    pub fn from_io(
        reader: impl Read + Send + 'static,
        writer: impl Write + Send + 'static,
    ) -> Self {
        let pending = Pending::default();
        let alive = Arc::new(AtomicBool::new(true));
        let (tx, rx) = mpsc::sync_channel(NOTIFICATION_BACKLOG);

        {
            let pending = pending.clone();
            let alive = alive.clone();
            thread::spawn(move || {
                read_loop(reader, &pending, &tx);
                alive.store(false, Ordering::SeqCst);
                // Dropping the senders wakes every waiting caller.
                pending.lock().unwrap().clear();
            });
        }

        Self {
            writer: Mutex::new(Box::new(writer)),
            pending,
            notifications: Mutex::new(rx),
            alive,
            next_id: AtomicU32::new(1),
            timeout: DEFAULT_TIMEOUT,
        }
    }

    pub fn set_timeout(&mut self, timeout: Duration) {
        self.timeout = timeout;
    }

    fn handshake(&self) -> Result<DeviceInfo> {
        let info = self.info()?;
        if info.protocol_version != PROTOCOL_VERSION {
            return Err(Error::ProtocolMismatch {
                device: info.protocol_version,
                host: PROTOCOL_VERSION,
            });
        }
        Ok(info)
    }

    /// Send one command and wait for its reply.
    ///
    /// # Errors
    ///
    /// [`Error::Encode`] if the request can't be encoded, [`Error::Io`] if
    /// writing fails, [`Error::Timeout`] if no reply arrives in time,
    /// [`Error::Disconnected`] if the connection is gone, or
    /// [`Error::Device`] if the device reports an error.
    ///
    /// # Panics
    ///
    /// If an internal lock is poisoned (another thread panicked holding it).
    pub fn call(&self, command: Command) -> Result<Reply> {
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        let wire =
            irbt_proto::encode(&Frame::Request(Request { id, command })).map_err(Error::Encode)?;

        let (tx, rx) = mpsc::channel();
        self.pending.lock().unwrap().insert(id, tx);
        let forget = || {
            self.pending.lock().unwrap().remove(&id);
        };

        if !self.alive.load(Ordering::SeqCst) {
            forget();
            return Err(Error::Disconnected);
        }
        let written = {
            let mut writer = self.writer.lock().unwrap();
            writer.write_all(&wire).and_then(|()| writer.flush())
        };
        if let Err(e) = written {
            forget();
            return Err(e.into());
        }

        match rx.recv_timeout(self.timeout) {
            Ok(Ok(reply)) => Ok(reply),
            Ok(Err(message)) => Err(Error::Device(message)),
            Err(mpsc::RecvTimeoutError::Timeout) => {
                forget();
                Err(Error::Timeout)
            }
            Err(mpsc::RecvTimeoutError::Disconnected) => Err(Error::Disconnected),
        }
    }

    fn call_ok(&self, command: Command) -> Result<()> {
        match self.call(command)? {
            Reply::Ok => Ok(()),
            other => Err(Error::UnexpectedReply(other)),
        }
    }

    /// Check that the device is responding.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::UnexpectedReply`] if the device answers
    /// with the wrong reply type.
    pub fn ping(&self) -> Result<()> {
        match self.call(Command::Ping)? {
            Reply::Pong => Ok(()),
            other => Err(Error::UnexpectedReply(other)),
        }
    }

    /// Firmware and protocol version.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::UnexpectedReply`] if the device answers
    /// with the wrong reply type.
    pub fn info(&self) -> Result<DeviceInfo> {
        match self.call(Command::GetInfo)? {
            Reply::Info(info) => Ok(info),
            other => Err(Error::UnexpectedReply(other)),
        }
    }

    /// Git commit, toolchain and component versions the firmware was built
    /// with. Firmware older than this command drops the request, so this
    /// times out on it; check [`DeviceInfo::firmware_version`] first.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::UnexpectedReply`] if the device answers
    /// with the wrong reply type.
    pub fn build_info(&self) -> Result<BuildInfo> {
        match self.call(Command::GetBuildInfo)? {
            Reply::BuildInfo(info) => Ok(info),
            other => Err(Error::UnexpectedReply(other)),
        }
    }

    /// Uptime, BLE state, mapping count and current mode.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::UnexpectedReply`] if the device answers
    /// with the wrong reply type.
    pub fn status(&self) -> Result<Status> {
        match self.call(Command::GetStatus)? {
            Reply::Status(status) => Ok(status),
            other => Err(Error::UnexpectedReply(other)),
        }
    }

    /// The full keymap, fetched page by page.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::UnexpectedReply`] if the device answers
    /// with the wrong reply type.
    pub fn keymap(&self) -> Result<Vec<Mapping>> {
        let mut mappings = Vec::new();
        loop {
            let offset = u16::try_from(mappings.len()).unwrap_or(u16::MAX);
            let page = match self.call(Command::ListKeymap { offset })? {
                Reply::Keymap(page) => page,
                other => return Err(Error::UnexpectedReply(other)),
            };
            let done = page.entries.is_empty()
                || usize::from(page.offset) + page.entries.len() >= usize::from(page.total);
            mappings.extend(page.entries);
            if done {
                return Ok(mappings);
            }
        }
    }

    /// Map `ir` to `key`, replacing any existing mapping for `ir`.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::Device`] also covers a full keymap.
    pub fn set_mapping(&self, ir: IrCode, key: HidKey) -> Result<()> {
        self.call_ok(Command::SetMapping(Mapping { ir, key }))
    }

    /// Remove the mapping for `ir`.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::Device`] also covers `ir` not being mapped.
    pub fn remove_mapping(&self, ir: IrCode) -> Result<()> {
        self.call_ok(Command::RemoveMapping { ir })
    }

    /// Remove every mapping.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::UnexpectedReply`] if the device answers
    /// with the wrong reply type.
    pub fn clear_keymap(&self) -> Result<()> {
        self.call_ok(Command::ClearKeymap)
    }

    /// Restore the firmware's built-in default keymap.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::UnexpectedReply`] if the device answers
    /// with the wrong reply type.
    pub fn reset_keymap(&self) -> Result<()> {
        self.call_ok(Command::ResetKeymap)
    }

    /// Put the device in learn mode and block until the next IR code is
    /// bound to `key`. Returns `None` if no code arrived within `timeout`.
    ///
    /// Other notifications received while waiting are discarded.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::Device`] if the code couldn't be stored
    /// (e.g. keymap full), [`Error::Timeout`] if the device never reports
    /// the outcome.
    pub fn learn(&self, key: HidKey, timeout: Duration) -> Result<Option<Mapping>> {
        let timeout_ms = u32::try_from(timeout.as_millis()).unwrap_or(u32::MAX);
        self.call_ok(Command::Learn { key, timeout_ms })?;

        // The device enforces `timeout`; the extra margin covers latency.
        let deadline = Instant::now() + timeout + self.timeout;
        loop {
            let left = deadline.saturating_duration_since(Instant::now());
            match self.next_notification(left)? {
                Some(Notification::Event(Event::Learned(mapping))) => return Ok(Some(mapping)),
                Some(Notification::Event(Event::LearnTimedOut)) => return Ok(None),
                Some(Notification::Event(Event::LearnFailed { reason })) => {
                    return Err(Error::Device(reason))
                }
                Some(_) => {}
                None => return Err(Error::Timeout),
            }
        }
    }

    /// Leave learn mode without binding anything.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::UnexpectedReply`] if the device answers
    /// with the wrong reply type.
    pub fn cancel_learn(&self) -> Result<()> {
        self.call_ok(Command::CancelLearn)
    }

    /// Press and release `key` over BLE, as if a mapped IR code arrived.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::UnexpectedReply`] if the device answers
    /// with the wrong reply type.
    pub fn send_hid(&self, key: HidKey) -> Result<()> {
        self.call_ok(Command::SendHid { key })
    }

    /// Transmit a NEC code through the IR blaster.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::UnexpectedReply`] if the device answers
    /// with the wrong reply type.
    pub fn send_ir(&self, code: IrCode) -> Result<()> {
        self.call_ok(Command::SendIr { code })
    }

    /// Enable/disable [`Event::IrReceived`] notifications.
    ///
    /// # Errors
    ///
    /// See [`Device::call`]; [`Error::UnexpectedReply`] if the device answers
    /// with the wrong reply type.
    pub fn set_monitor(&self, enabled: bool) -> Result<()> {
        self.call_ok(Command::SetMonitor { enabled })
    }

    /// Wait up to `timeout` for the next event, log line or raw output.
    /// `Ok(None)` means nothing arrived in time.
    ///
    /// # Errors
    ///
    /// [`Error::Disconnected`] once the connection is gone and every
    /// buffered notification has been read.
    ///
    /// # Panics
    ///
    /// If an internal lock is poisoned (another thread panicked holding it).
    pub fn next_notification(&self, timeout: Duration) -> Result<Option<Notification>> {
        match self.notifications.lock().unwrap().recv_timeout(timeout) {
            Ok(notification) => Ok(Some(notification)),
            Err(mpsc::RecvTimeoutError::Timeout) => Ok(None),
            Err(mpsc::RecvTimeoutError::Disconnected) => Err(Error::Disconnected),
        }
    }
}

fn read_loop(mut reader: impl Read, pending: &Pending, notify: &mpsc::SyncSender<Notification>) {
    let mut decoder = Decoder::new();
    let mut buf = [0u8; 512];

    let mut dispatch = |decoded: Decoded| {
        let notification = match decoded {
            Decoded::Frame(Frame::Response(response)) => {
                // A missing entry means the caller already timed out.
                if let Some(tx) = pending.lock().unwrap().remove(&response.id) {
                    let _ = tx.send(response.result);
                }
                return;
            }
            Decoded::Frame(Frame::Event(event)) => Notification::Event(event),
            Decoded::Frame(Frame::Log(log)) => Notification::Log(log),
            Decoded::Frame(Frame::Request(_)) => return,
            Decoded::Raw(bytes) => Notification::Raw(String::from_utf8_lossy(&bytes).into_owned()),
        };
        // Full means nobody is consuming notifications; dropping is fine.
        let _ = notify.try_send(notification);
    };

    loop {
        match reader.read(&mut buf) {
            Ok(0) => break,
            Ok(n) => buf[..n]
                .iter()
                .filter_map(|&byte| decoder.push(byte))
                .for_each(&mut dispatch),
            Err(e)
                if matches!(
                    e.kind(),
                    io::ErrorKind::TimedOut | io::ErrorKind::WouldBlock
                ) =>
            {
                decoder.flush().into_iter().for_each(&mut dispatch);
            }
            Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
            Err(_) => break,
        }
    }
    decoder.flush().into_iter().for_each(dispatch);
}

#[cfg(test)]
mod tests;