omni_cli/transport/

relay.rs

//! iOS/watch outbound relay transport (bd-a5452f).
//!
//! watchOS cannot host a TCP server (Apple TN3135), so omni endpoints (iOS
//! extension while open, watch app, Wear service) dial OUTBOUND to a relay and
//! omni-cli posts commands to it. Harry chose relay-in-omni-cli, so the relay is
//! the `omni serve` subcommand. Protocol (apple-utils omni/docs/injection.md,
//! "Relay protocol (omni listen)"): TCP, NDJSON (one JSON object per line).
//!
//! Roles: DEVICE (registers + receives), SENDER (`omni send --relay` / cast),
//! RELAY (`omni serve`). The DEVICE only ever receives **bare**
//! [`InjectionCommand`] lines, so the device drain/apply stays trivial.
//!
//! Frames:
//! - device -> relay: `{"v":1,"op":"register","device":"<id>","platform":"...","token":"<secret>"}`
//! - sender -> relay: `{"v":1,"op":"cast","device":"<id|all>","cmd":{InjectionCommand},"token":"<secret>"}`
//! - relay -> device: a bare InjectionCommand line
//! - local app -> daemon: `{"v":1,"op":"subscribe","token":"..."}`
//! - daemon -> local app: mirrored bare InjectionCommand lines
//! - keepalive: `{"v":1,"op":"ping"}` / `{"v":1,"op":"pong"}`
//!
//! A shared `token` (from SOPS/env, never hard-coded) gates register + cast.
//! Device-ids are operator-assigned (e.g. `harry-watch`).

use std::collections::HashMap;
use std::io::{BufRead, BufReader, Write};
use std::net::{Shutdown, TcpListener, TcpStream, ToSocketAddrs};
use std::sync::mpsc::{self, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;

use socket2::{SockRef, TcpKeepalive};

use crate::AppError;
use crate::events::SemanticEvent;
use crate::wire::{InjectionCommand, WIRE_VERSION};

pub use omni_sdk::transport::relay::{ControlFrame, DEVICE_ALL, TOKEN_ENV};

/// Idle time on a registered device socket before the OS starts sending TCP
/// keepalive probes (bd-86ab7a).
const KEEPALIVE_IDLE: Duration = Duration::from_secs(20);

/// Interval between TCP keepalive probes once they start (bd-86ab7a).
const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(5);

/// The NDJSON line the relay sends to reject an unauthorized register/cast.
/// Typed (serializes a [`ControlFrame::Error`], op-first like the other frames)
/// so device clients match it by type rather than substring-scanning.
fn unauthorized_line() -> String {
    ControlFrame::Error {
        v: WIRE_VERSION,
        reason: "unauthorized".to_owned(),
    }
    .json_line()
    .unwrap_or_else(|_| r#"{"op":"error","v":1,"reason":"unauthorized"}"#.to_owned())
}

/// The NDJSON line the relay sends back to a SENDER after routing a cast
/// (bd-0b4e45): how many registered devices received it. `delivered == 0` tells
/// `omni send --relay` the target device is not connected.
fn result_line(device: &str, delivered: usize) -> String {
    ControlFrame::Result {
        v: WIRE_VERSION,
        device: device.to_owned(),
        delivered,
    }
    .json_line()
    .unwrap_or_else(|_| format!(r#"{{"op":"result","v":1,"device":"","delivered":{delivered}}}"#))
}

/// A device-bound output frame serialized by the single per-device writer
/// thread. Routing pongs through the same channel as casts guarantees exactly
/// ONE thread ever writes to a device socket, so two NDJSON lines can never be
/// spliced together by concurrent writes (bd-c7bcfc).
#[derive(Debug, Clone, PartialEq, Eq)]
enum DeviceOut {
    /// A bare [`InjectionCommand`] line for the device to apply.
    Cmd(InjectionCommand),
    Event(SemanticEvent),
    /// A keepalive pong reply to a device-initiated ping.
    Pong,
}

/// The relay's device registry: device-id -> a sender feeding that device's
/// single writer thread (commands + pong replies).
#[derive(Clone, Default)]
pub struct Registry {
    inner: Arc<Mutex<HashMap<String, RegisteredDevice>>>,
}

struct RegisteredDevice {
    sender: Sender<DeviceOut>,
    events: bool,
}

impl Registry {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    fn register(&self, device: &str, tx: Sender<DeviceOut>) {
        self.register_capable(device, tx, false);
    }

    fn register_capable(&self, device: &str, tx: Sender<DeviceOut>, events: bool) {
        self.inner
            .lock()
            .unwrap()
            .insert(device.to_owned(), RegisteredDevice { sender: tx, events });
    }

    /// Only explicitly capable registrations receive non-injecting events.
    pub fn route_event(&self, device: &str, event: &SemanticEvent) -> usize {
        if event.validate().is_err() {
            return 0;
        }
        self.inner
            .lock()
            .unwrap()
            .iter()
            .filter(|(id, peer)| peer.events && (device == DEVICE_ALL || id.as_str() == device))
            .filter(|(_, peer)| peer.sender.send(DeviceOut::Event(event.clone())).is_ok())
            .count()
    }

    fn unregister(&self, device: &str) {
        self.inner.lock().unwrap().remove(device);
    }

    /// Whether `device` currently has a live registration.
    #[must_use]
    pub fn is_registered(&self, device: &str) -> bool {
        self.inner.lock().unwrap().contains_key(device)
    }

    /// The device ids currently registered, sorted (bd-bc8b5c). Used to answer a
    /// sender's `list` query so operators can see which cross-device endpoints
    /// are connected to the hub.
    #[must_use]
    pub fn device_ids(&self) -> Vec<String> {
        let mut ids: Vec<String> = self.inner.lock().unwrap().keys().cloned().collect();
        ids.sort();
        ids
    }

    /// Number of currently-registered devices.
    #[must_use]
    pub fn device_count(&self) -> usize {
        self.inner.lock().unwrap().len()
    }

    /// Route a command to one device id, or to all when `device == "all"`.
    /// Returns the number of devices the command was delivered to (queued).
    pub fn route(&self, device: &str, cmd: &InjectionCommand) -> usize {
        let map = self.inner.lock().unwrap();
        if device == DEVICE_ALL {
            let mut n = 0;
            for tx in map.values() {
                if tx.sender.send(DeviceOut::Cmd(cmd.clone())).is_ok() {
                    n += 1;
                }
            }
            n
        } else if let Some(tx) = map.get(device) {
            usize::from(tx.sender.send(DeviceOut::Cmd(cmd.clone())).is_ok())
        } else {
            0
        }
    }
}

/// Options for [`serve`].
#[derive(Debug, Clone)]
pub struct ServeOptions {
    /// Shared auth token required on register + cast frames.
    pub token: String,
    /// Stop accepting after this many connections (for tests); `None` = forever.
    pub max_conns: Option<usize>,
}

/// Bind a relay listener on `addr`.
///
/// # Errors
/// Returns [`AppError::internal`] if binding fails.
pub fn bind(addr: &str) -> Result<TcpListener, AppError> {
    TcpListener::bind(addr).map_err(|e| AppError::internal("relay_bind", format!("{addr}: {e}")))
}

/// Run the relay accept loop on `listener`. One thread per connection; the
/// shared [`Registry`] routes sender casts to registered devices.
///
/// # Errors
/// Returns [`AppError::internal`] on a fatal accept error.
pub fn serve(listener: &TcpListener, opts: &ServeOptions) -> Result<(), AppError> {
    serve_with_registry(listener, opts, &Registry::new())
}

/// Like [`serve`] but routes through a caller-provided [`Registry`], so the
/// device-registration set is observable (used by tests to assert that a dead
/// device is reaped; bd-86ab7a).
///
/// # Errors
/// Returns [`AppError::internal`] on a fatal accept error.
pub fn serve_with_registry(
    listener: &TcpListener,
    opts: &ServeOptions,
    registry: &Registry,
) -> Result<(), AppError> {
    let mut accepted = 0usize;
    for incoming in listener.incoming() {
        let stream = incoming.map_err(|e| AppError::internal("relay_accept", e.to_string()))?;
        let registry = registry.clone();
        let token = opts.token.clone();
        thread::spawn(move || {
            if let Err(error) = handle_conn(stream, &registry, &token) {
                eprintln!("omni serve: connection error: {error}");
            }
        });
        accepted += 1;
        if let Some(max) = opts.max_conns {
            if accepted >= max {
                break;
            }
        }
    }
    Ok(())
}

fn write_line<W: Write>(stream: &mut W, line: &str) -> std::io::Result<()> {
    stream.write_all(line.as_bytes())?;
    stream.write_all(b"\n")?;
    stream.flush()
}

/// Best-effort: enable TCP keepalive on a registered device socket so the OS
/// surfaces a half-open connection (a network drop with no FIN — common for a
/// watch/phone over a flaky tailnet) as a socket error. That lets the relay's
/// reader observe the death and reap the dead registration proactively, rather
/// than waiting for the next cast to be queued into the void (bd-86ab7a).
/// Failures are ignored: keepalive is hardening, never required for delivery.
fn enable_device_keepalive(stream: &TcpStream) {
    let keepalive = TcpKeepalive::new()
        .with_time(KEEPALIVE_IDLE)
        .with_interval(KEEPALIVE_INTERVAL);
    let _ = SockRef::from(stream).set_tcp_keepalive(&keepalive);
}

/// Drain queued [`DeviceOut`] frames to a device `writer` — the SOLE writer to
/// the device socket, so command and pong lines can never interleave on the
/// wire (bd-c7bcfc). On the FIRST write failure the device socket is dead, so
/// reap the registration immediately — otherwise later casts are silently
/// `tx.send`-queued into a dead connection and lost (bd-86ab7a) — and run
/// `on_dead` (used to shut the socket down so the paired reader loop unblocks).
/// Returns when the channel closes (normal disconnect) or after evicting a dead
/// device.
fn run_device_writer<W: Write>(
    rx: &Receiver<DeviceOut>,
    writer: &mut W,
    registry: &Registry,
    device: &str,
    on_dead: impl FnOnce(),
) {
    while let Ok(out) = rx.recv() {
        let line = match out {
            DeviceOut::Cmd(cmd) => match cmd.json_line() {
                Ok(line) => line,
                Err(_) => continue,
            },
            DeviceOut::Event(event) => match event.json_line() {
                Ok(line) => line,
                Err(_) => continue,
            },
            DeviceOut::Pong => r#"{"v":1,"op":"pong"}"#.to_owned(),
        };
        if write_line(writer, &line).is_err() {
            registry.unregister(device);
            on_dead();
            return;
        }
    }
}

fn handle_conn(stream: TcpStream, registry: &Registry, token: &str) -> Result<(), AppError> {
    let mut reader = BufReader::new(
        stream
            .try_clone()
            .map_err(|e| AppError::internal("relay_clone", e.to_string()))?,
    );
    let mut first = String::new();
    let n = reader
        .read_line(&mut first)
        .map_err(|e| AppError::internal("relay_read", e.to_string()))?;
    if n == 0 {
        return Ok(()); // connected and closed without a frame.
    }
    let frame: ControlFrame = serde_json::from_str(first.trim())
        .map_err(|e| AppError::validation("relay_frame", e.to_string()))?;

    match frame {
        ControlFrame::Register {
            device, token: t, ..
        } => {
            if t.trim() != token.trim() {
                let mut s = stream;
                let _ = write_line(&mut s, &unauthorized_line());
                return Ok(());
            }
            // An additive registration capability keeps events off legacy devices.
            let events = serde_json::from_str::<serde_json::Value>(first.trim())
                .ok()
                .is_some_and(|v| v["events"] == true);
            serve_device(stream, reader, registry, &device, events)
        }
        ControlFrame::Cast {
            device,
            cmd,
            token: t,
            ..
        } => {
            if t.trim() != token.trim() {
                let mut s = stream;
                let _ = write_line(&mut s, &unauthorized_line());
                return Ok(());
            }
            let delivered = registry.route(&device, &cmd);
            // Ack the FIRST cast back to the sender (bd-0b4e45) so `omni send
            // --relay` learns whether the target device is connected. Only the
            // first cast is acked: a streaming sender that never reads results
            // accumulates at most one unread line, so its receive buffer cannot
            // fill and stall the relay.
            if let Ok(mut ack) = stream.try_clone() {
                let _ = write_line(&mut ack, &result_line(&device, delivered));
            }
            // A sender may stream more cast/ping frames until it closes.
            serve_sender(stream, reader, registry, token)
        }
        frame @ ControlFrame::EventCast { .. } => {
            serve_event_sender(stream, reader, registry, token, frame)
        }
        // Subscriptions are local `omni daemon` operations, not relay-hub roles.
        ControlFrame::Subscribe { .. } | ControlFrame::SubscribeEvents { .. } => Ok(()),
        ControlFrame::Ping { .. } => {
            let mut s = stream;
            let _ = write_line(&mut s, r#"{"v":1,"op":"pong"}"#);
            Ok(())
        }
        ControlFrame::Pong { .. } => Ok(()),
        ControlFrame::List { token: t, .. } => {
            let mut s = stream;
            if t.trim() != token.trim() {
                let _ = write_line(&mut s, &unauthorized_line());
                return Ok(());
            }
            let frame = ControlFrame::Devices {
                v: WIRE_VERSION,
                devices: registry.device_ids(),
            };
            let _ = write_line(&mut s, &frame.json_line().unwrap_or_default());
            Ok(())
        }
        // Error/Result/Devices are relay->client only; a client never
        // legitimately sends them (bd-0b4e45 Result, bd-bc8b5c Devices).
        ControlFrame::Error { .. }
        | ControlFrame::Result { .. }
        | ControlFrame::EventResult { .. }
        | ControlFrame::EventsSubscribed { .. }
        | ControlFrame::Devices { .. } => Ok(()),
    }
}

/// A registered DEVICE connection: forward queued [`InjectionCommand`]s to the
/// socket (writer thread), and read keepalives until the device disconnects.
///
/// Liveness hardening (bd-86ab7a): TCP keepalive is enabled so the OS surfaces
/// a half-open socket as a read error (reaped by the reader loop below), and
/// the writer thread reaps the registration + shuts the socket down the moment
/// a write fails, so a dead device is not left registered with later casts
/// silently queued into the void.
fn serve_device(
    stream: TcpStream,
    mut reader: BufReader<TcpStream>,
    registry: &Registry,
    device: &str,
    events: bool,
) -> Result<(), AppError> {
    enable_device_keepalive(&stream);

    let (tx, rx) = mpsc::channel::<DeviceOut>();
    if events {
        registry.register_capable(device, tx.clone(), true);
    } else {
        registry.register(device, tx.clone());
    }
    // The reader enqueues pongs to this same writer thread instead of writing
    // to the socket itself, so the device socket has exactly ONE writer and
    // NDJSON lines can't interleave (bd-c7bcfc).
    let pong_tx = tx;

    let mut writer_stream = stream
        .try_clone()
        .map_err(|e| AppError::internal("relay_clone", e.to_string()))?;
    let shutdown_stream = stream
        .try_clone()
        .map_err(|e| AppError::internal("relay_clone", e.to_string()))?;
    let writer_registry = registry.clone();
    let writer_device = device.to_owned();
    let writer = thread::spawn(move || {
        run_device_writer(
            &rx,
            &mut writer_stream,
            &writer_registry,
            &writer_device,
            || {
                // Unblock the (possibly half-open-blocked) reader loop so the
                // connection tears down promptly instead of leaking a thread.
                let _ = shutdown_stream.shutdown(Shutdown::Both);
            },
        );
    });

    let mut line = String::new();
    loop {
        line.clear();
        match reader.read_line(&mut line) {
            Ok(0) | Err(_) => break, // device disconnected (or keepalive surfaced a dead socket)
            Ok(_) => {
                if let Ok(ControlFrame::Ping { .. }) =
                    serde_json::from_str::<ControlFrame>(line.trim())
                {
                    // Enqueue the pong to the single writer thread (never write
                    // to the socket here, bd-c7bcfc). A send error just means
                    // the device/writer is already gone.
                    let _ = pong_tx.send(DeviceOut::Pong);
                }
            }
        }
    }

    drop(pong_tx); // release our writer-channel sender so the writer can finish
    registry.unregister(device); // drops the registry's sender -> writer loop ends
    let _ = writer.join();
    Ok(())
}

/// Event senders require a result for every event; legacy cast acknowledgement
/// behaviour stays unchanged. Nothing on this path becomes an InjectionCommand.
fn serve_event_sender(
    mut writer: TcpStream,
    mut reader: BufReader<TcpStream>,
    registry: &Registry,
    token: &str,
    first: ControlFrame,
) -> Result<(), AppError> {
    let mut frame = first;
    loop {
        let (reply, close) = match frame {
            ControlFrame::EventCast {
                v,
                device,
                event,
                token: got,
            } => {
                if got.trim() != token.trim() {
                    (unauthorized_line(), true)
                } else if v != WIRE_VERSION || event.validate().is_err() {
                    (
                        ControlFrame::Error {
                            v: WIRE_VERSION,
                            reason: "invalid_event".into(),
                        }
                        .json_line()?,
                        true,
                    )
                } else {
                    let delivered = registry.route_event(&device, &event);
                    (
                        ControlFrame::EventResult {
                            v: WIRE_VERSION,
                            device,
                            delivered,
                        }
                        .json_line()?,
                        false,
                    )
                }
            }
            ControlFrame::Ping { .. } => {
                (ControlFrame::Pong { v: WIRE_VERSION }.json_line()?, false)
            }
            _ => (
                ControlFrame::Error {
                    v: WIRE_VERSION,
                    reason: "expected_event_frame".into(),
                }
                .json_line()?,
                true,
            ),
        };
        write_line(&mut writer, &reply)
            .map_err(|e| AppError::internal("event_relay_write", e.to_string()))?;
        if close {
            return Ok(());
        }
        let mut line = String::new();
        let n = std::io::Read::take(&mut reader, (crate::events::MAX_EVENT_BYTES + 4097) as u64)
            .read_line(&mut line)
            .map_err(|e| AppError::internal("event_relay_read", e.to_string()))?;
        if n == 0 {
            return Ok(());
        }
        if n > crate::events::MAX_EVENT_BYTES + 4096 {
            return Err(AppError::validation("event_size", "event frame too large"));
        }
        frame = serde_json::from_str(line.trim())
            .map_err(|e| AppError::validation("event_frame", e.to_string()))?;
    }
}

/// A SENDER connection that has already delivered its first cast: keep reading
/// additional cast/ping frames until it closes.
fn serve_sender(
    stream: TcpStream,
    mut reader: BufReader<TcpStream>,
    registry: &Registry,
    token: &str,
) -> Result<(), AppError> {
    let mut pong_stream = stream;
    let mut line = String::new();
    loop {
        line.clear();
        match reader.read_line(&mut line) {
            Ok(0) | Err(_) => break,
            Ok(_) => match serde_json::from_str::<ControlFrame>(line.trim()) {
                Ok(ControlFrame::Cast {
                    device,
                    cmd,
                    token: t,
                    ..
                }) if t.trim() == token.trim() => {
                    registry.route(&device, &cmd);
                }
                Ok(ControlFrame::Ping { .. }) => {
                    let _ = write_line(&mut pong_stream, r#"{"v":1,"op":"pong"}"#);
                }
                _ => {}
            },
        }
    }
    Ok(())
}

/// How long a sender waits for the relay's delivery result before giving up (a
/// legacy relay sends none). Short so `omni send --relay` never hangs (bd-0b4e45).
const RESULT_READ_TIMEOUT: Duration = Duration::from_millis(800);

/// Outcome of a relay cast (bd-0b4e45): whether the relay acknowledged it and
/// how many registered devices it was routed to.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RelayDelivery {
    /// True when the relay returned a delivery result. A legacy relay (or a read
    /// timeout) leaves this false, so a caller must not treat `delivered` as
    /// authoritative unless `acked` is true.
    pub acked: bool,
    /// Number of registered devices the cast was routed to (0 = the target
    /// device id is not connected to the relay).
    pub delivered: usize,
}

/// Read the relay's delivery result for the first cast on `stream` (bd-0b4e45).
/// Bounded + best-effort: a legacy relay that sends nothing, a timeout, EOF, or
/// a malformed line all yield `acked == false`, so the sender never blocks.
fn read_delivery(stream: &mut TcpStream) -> RelayDelivery {
    let _ = stream.set_read_timeout(Some(RESULT_READ_TIMEOUT));
    let mut reader = BufReader::new(stream);
    let mut line = String::new();
    match reader.read_line(&mut line) {
        Ok(n) if n > 0 => match serde_json::from_str::<ControlFrame>(line.trim()) {
            Ok(ControlFrame::Result { delivered, .. }) => RelayDelivery {
                acked: true,
                delivered,
            },
            _ => RelayDelivery {
                acked: false,
                delivered: 0,
            },
        },
        _ => RelayDelivery {
            acked: false,
            delivered: 0,
        },
    }
}

/// SENDER client: connect to a relay and post one or more commands for `device`,
/// then read the relay's delivery result for the first cast (bd-0b4e45) so the
/// caller can tell whether the target device was actually connected. Casting to
/// a plain `omni daemon` (which sends no result) yields `acked == false`.
///
/// # Errors
/// Returns [`AppError::internal`] on connect / I/O failure.
pub fn cast_via_relay<A: ToSocketAddrs>(
    addr: A,
    token: &str,
    device: &str,
    cmds: &[InjectionCommand],
) -> Result<RelayDelivery, AppError> {
    let mut stream = super::connect_bounded(addr, super::CONNECT_TIMEOUT)
        .map_err(|e| AppError::internal("relay_connect", e.to_string()))?;
    for cmd in cmds {
        let frame = ControlFrame::Cast {
            v: WIRE_VERSION,
            device: device.to_owned(),
            cmd: cmd.clone(),
            token: token.to_owned(),
        };
        let line = frame.json_line()?;
        write_line(&mut stream, &line)
            .map_err(|e| AppError::internal("relay_send", e.to_string()))?;
    }
    Ok(read_delivery(&mut stream))
}

/// Query a relay for the device ids currently registered (bd-bc8b5c), so
/// `omni devices --relay` can show operators which cross-device endpoints are
/// connected. Bounded: a legacy relay that does not understand `list` sends no
/// `devices` reply, so this returns an empty list (never hangs) within the read
/// timeout. An auth rejection surfaces as a validation error.
///
/// # Errors
/// [`AppError::internal`] on connect / I/O failure; [`AppError::validation`]
/// when the relay rejects the token.
pub fn list_relay_devices<A: ToSocketAddrs>(addr: A, token: &str) -> Result<Vec<String>, AppError> {
    let mut stream = super::connect_bounded(addr, super::CONNECT_TIMEOUT)
        .map_err(|e| AppError::internal("relay_connect", e.to_string()))?;
    let frame = ControlFrame::List {
        v: WIRE_VERSION,
        token: token.to_owned(),
    };
    write_line(&mut stream, &frame.json_line()?)
        .map_err(|e| AppError::internal("relay_send", e.to_string()))?;
    let _ = stream.set_read_timeout(Some(RESULT_READ_TIMEOUT));
    let mut reader = BufReader::new(stream);
    let mut line = String::new();
    match reader.read_line(&mut line) {
        Ok(n) if n > 0 => match serde_json::from_str::<ControlFrame>(line.trim()) {
            Ok(ControlFrame::Devices { devices, .. }) => Ok(devices),
            Ok(ControlFrame::Error { reason, .. }) => {
                Err(AppError::validation("relay_list", reason))
            }
            _ => Ok(Vec::new()),
        },
        _ => Ok(Vec::new()),
    }
}

/// A hold-open relay SENDER for STREAMING dispatch (bd-772cb2): connects once
/// and writes a `Cast` frame per command instead of reconnecting per command
/// (the sender side needs no register handshake — the relay routes casts to the
/// registered devices). Lazy-connects on the first cast; reconnects on the next
/// cast after a write error so a transient drop does not abort the stream.
pub struct RelaySender {
    addr: String,
    device: String,
    token: String,
    stream: Option<TcpStream>,
}

impl RelaySender {
    /// Build a sender targeting `addr` (host:port) for `device` with `token`
    /// (no connection is made until the first cast).
    #[must_use]
    pub fn new(
        addr: impl Into<String>,
        device: impl Into<String>,
        token: impl Into<String>,
    ) -> Self {
        Self {
            addr: addr.into(),
            device: device.into(),
            token: token.into(),
            stream: None,
        }
    }

    /// Write one `Cast` frame over the held-open connection, connecting first if
    /// needed. On any I/O error the connection is dropped so the next call
    /// reconnects.
    ///
    /// # Errors
    /// [`AppError::internal`] on connect / write failure (the stream is reset
    /// for a fresh reconnect next time).
    pub fn cast_one(&mut self, cmd: &InjectionCommand) -> Result<(), AppError> {
        if self.stream.is_none() {
            let s = super::connect_bounded(&self.addr, super::CONNECT_TIMEOUT)
                .map_err(|e| AppError::internal("relay_connect", format!("{}: {e}", self.addr)))?;
            self.stream = Some(s);
        }
        let frame = ControlFrame::Cast {
            v: WIRE_VERSION,
            device: self.device.clone(),
            cmd: cmd.clone(),
            token: self.token.clone(),
        };
        let line = frame.json_line()?;
        let stream = self.stream.as_mut().expect("stream just set");
        if let Err(e) = write_line(stream, &line) {
            self.stream = None; // force reconnect next time
            return Err(AppError::internal(
                "relay_send",
                format!("{}: {e}", self.addr),
            ));
        }
        Ok(())
    }
}

/// Resolve the shared token from an explicit flag or the [`TOKEN_ENV`] env var.
///
/// # Errors
/// Returns [`AppError::validation`] when no non-empty token is available.
pub fn resolve_token(flag: Option<&str>) -> Result<String, AppError> {
    resolve_token_with(flag, std::env::var(TOKEN_ENV).ok())
}

/// Pure core of [`resolve_token`] (env value injected for testability).
///
/// # Errors
/// Returns [`AppError::validation`] when neither flag nor env supplies a
/// non-empty token.
pub fn resolve_token_with(
    flag: Option<&str>,
    env_value: Option<String>,
) -> Result<String, AppError> {
    // Trim surrounding whitespace so a token sourced from a FILE (e.g. a
    // sops-decrypted secret with a trailing newline) matches a clean token.
    // The relay auth gate also trims on both sides (bd-790088), so this keeps
    // omni-native senders/hubs consistent with that normalization.
    if let Some(t) = flag {
        let t = t.trim();
        if !t.is_empty() {
            return Ok(t.to_owned());
        }
    }
    match env_value {
        Some(t) if !t.trim().is_empty() => Ok(t.trim().to_owned()),
        _ => Err(AppError::validation(
            "relay_token",
            format!("a shared token is required; pass --token or set {TOKEN_ENV}"),
        )),
    }
}

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

    #[test]
    fn frames_round_trip() {
        let reg = ControlFrame::Register {
            v: 1,
            device: "harry-watch".into(),
            platform: "watchos".into(),
            token: "secret".into(),
        };
        let line = reg.json_line().unwrap();
        assert_eq!(
            line,
            r#"{"op":"register","v":1,"device":"harry-watch","platform":"watchos","token":"secret"}"#
        );
        assert_eq!(serde_json::from_str::<ControlFrame>(&line).unwrap(), reg);

        let cast = ControlFrame::Cast {
            v: 1,
            device: "all".into(),
            cmd: InjectionCommand::text("hi"),
            token: "secret".into(),
        };
        let line = cast.json_line().unwrap();
        assert_eq!(serde_json::from_str::<ControlFrame>(&line).unwrap(), cast);

        let err = ControlFrame::Error {
            v: 1,
            reason: "unauthorized".into(),
        };
        let line = err.json_line().unwrap();
        assert_eq!(line, r#"{"op":"error","v":1,"reason":"unauthorized"}"#);
        assert_eq!(serde_json::from_str::<ControlFrame>(&line).unwrap(), err);
    }

    #[test]
    fn registry_routes_to_device_and_all() {
        let reg = Registry::new();
        let (tx, rx) = mpsc::channel();
        reg.register("dev1", tx);
        assert_eq!(reg.route("dev1", &InjectionCommand::key("enter")), 1);
        assert_eq!(
            rx.recv_timeout(Duration::from_secs(1)).unwrap(),
            DeviceOut::Cmd(InjectionCommand::key("enter"))
        );
        // unknown device -> 0 delivered, no panic
        assert_eq!(reg.route("nope", &InjectionCommand::text("x")), 0);
        // all
        assert_eq!(reg.route(DEVICE_ALL, &InjectionCommand::scroll(2)), 1);
        assert_eq!(
            rx.recv_timeout(Duration::from_secs(1)).unwrap(),
            DeviceOut::Cmd(InjectionCommand::scroll(2))
        );
    }

    #[test]
    fn registry_is_registered_and_count_track_membership() {
        let reg = Registry::new();
        assert_eq!(reg.device_count(), 0);
        assert!(!reg.is_registered("dev1"));
        let (tx, _rx) = mpsc::channel();
        reg.register("dev1", tx);
        assert!(reg.is_registered("dev1"));
        assert_eq!(reg.device_count(), 1);
        reg.unregister("dev1");
        assert!(!reg.is_registered("dev1"));
        assert_eq!(reg.device_count(), 0);
    }

    /// A [`Write`] that always fails, standing in for a dead/half-open device
    /// socket so the eviction path is deterministic without real networking.
    struct DeadWriter;
    impl Write for DeadWriter {
        fn write(&mut self, _buf: &[u8]) -> std::io::Result<usize> {
            Err(std::io::Error::new(
                std::io::ErrorKind::BrokenPipe,
                "dead socket",
            ))
        }
        fn flush(&mut self) -> std::io::Result<()> {
            Ok(())
        }
    }

    #[test]
    fn dead_device_writer_evicts_registration_on_failed_cast() {
        // Regression for bd-86ab7a: a cast routed to a half-open device must NOT
        // leave the registration in place (which would silently queue every
        // later cast into the void). The first failed write reaps it.
        let registry = Registry::new();
        let (tx, rx) = mpsc::channel::<DeviceOut>();
        registry.register("harry-watch", tx.clone());
        assert!(registry.is_registered("harry-watch"));

        // Queue one cast, then drop our handle so recv() ends after that cast.
        tx.send(DeviceOut::Cmd(InjectionCommand::text("lost")))
            .unwrap();
        drop(tx);

        let dead_signalled = std::cell::Cell::new(false);
        run_device_writer(&rx, &mut DeadWriter, &registry, "harry-watch", || {
            dead_signalled.set(true);
        });

        assert!(
            !registry.is_registered("harry-watch"),
            "a dead device must be evicted on the failed write"
        );
        assert!(
            dead_signalled.get(),
            "on_dead must fire so the reader socket is shut down"
        );
        // A subsequent route now honestly reports 0 delivered, rather than
        // silently queueing into a dead connection.
        assert_eq!(
            registry.route("harry-watch", &InjectionCommand::text("after")),
            0
        );
    }

    #[test]
    fn healthy_device_writer_delivers_without_eviction() {
        // The eviction path must not regress the happy path: a healthy writer
        // delivers every queued command and the device stays registered when
        // the channel closes on a normal disconnect.
        let registry = Registry::new();
        let (tx, rx) = mpsc::channel::<DeviceOut>();
        tx.send(DeviceOut::Cmd(InjectionCommand::text("a")))
            .unwrap();
        tx.send(DeviceOut::Cmd(InjectionCommand::key("enter")))
            .unwrap();
        drop(tx); // normal disconnect: feeding channel closes, writer returns cleanly

        // A live registration whose sender stays alive, so we can assert the
        // healthy writer path left the registry untouched (no eviction).
        let (keep_tx, _keep_rx) = mpsc::channel();
        registry.register("dev", keep_tx);

        let mut buf: Vec<u8> = Vec::new();
        let mut dead_signalled = false;
        run_device_writer(&rx, &mut buf, &registry, "dev", || dead_signalled = true);

        assert!(
            !dead_signalled,
            "a healthy channel close must not signal death"
        );
        assert!(
            registry.is_registered("dev"),
            "channel-close (normal disconnect) does not evict via the writer path"
        );
        let text = String::from_utf8(buf).unwrap();
        assert_eq!(text.lines().count(), 2, "both commands delivered as NDJSON");
    }

    #[test]
    fn device_writer_serializes_pongs_and_commands_on_one_writer() {
        // bd-c7bcfc: pongs are routed through the SAME writer as casts, so a
        // device ping reply and a cast can never interleave on the socket.
        // Feed both kinds and assert three whole, un-spliced NDJSON lines.
        let registry = Registry::new();
        let (tx, rx) = mpsc::channel::<DeviceOut>();
        tx.send(DeviceOut::Cmd(InjectionCommand::text("hi")))
            .unwrap();
        tx.send(DeviceOut::Pong).unwrap();
        tx.send(DeviceOut::Cmd(InjectionCommand::key("enter")))
            .unwrap();
        drop(tx); // channel closes after the three frames -> writer returns

        let mut buf: Vec<u8> = Vec::new();
        let mut dead_signalled = false;
        run_device_writer(&rx, &mut buf, &registry, "dev", || dead_signalled = true);
        assert!(!dead_signalled, "no write failed, so no eviction");

        let lines: Vec<&str> = std::str::from_utf8(&buf).unwrap().lines().collect();
        assert_eq!(lines.len(), 3, "three whole NDJSON lines, none spliced");
        assert_eq!(
            InjectionCommand::decode_line(lines[0]).unwrap(),
            InjectionCommand::text("hi")
        );
        assert_eq!(lines[1], r#"{"v":1,"op":"pong"}"#);
        assert_eq!(
            InjectionCommand::decode_line(lines[2]).unwrap(),
            InjectionCommand::key("enter")
        );
    }

    #[test]
    fn end_to_end_dead_device_is_reaped_from_registry() {
        // With an observable registry, a device that drops its connection is
        // reaped so it no longer shows as registered (bd-86ab7a).
        let listener = bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let opts = ServeOptions {
            token: "tok".into(),
            max_conns: Some(1),
        };
        let registry = Registry::new();
        let serve_registry = registry.clone();
        thread::spawn(move || {
            let _ = serve_with_registry(&listener, &opts, &serve_registry);
        });

        let dev = TcpStream::connect(addr).unwrap();
        let mut dev_w = dev.try_clone().unwrap();
        let reg = ControlFrame::Register {
            v: 1,
            device: "harry-watch".into(),
            platform: "watchos".into(),
            token: "tok".into(),
        };
        write_line(&mut dev_w, &reg.json_line().unwrap()).unwrap();

        // Wait for the relay to register the device.
        assert!(
            wait_until(|| registry.is_registered("harry-watch"), 2000),
            "device should register"
        );

        // Device drops its connection; the relay reader must reap it.
        drop(dev_w);
        drop(dev);
        assert!(
            wait_until(|| !registry.is_registered("harry-watch"), 2000),
            "a disconnected device should be reaped from the registry"
        );
    }

    /// Poll `cond` every 10ms until it is true or `timeout_ms` elapses.
    fn wait_until(mut cond: impl FnMut() -> bool, timeout_ms: u64) -> bool {
        let start = std::time::Instant::now();
        while start.elapsed() < Duration::from_millis(timeout_ms) {
            if cond() {
                return true;
            }
            thread::sleep(Duration::from_millis(10));
        }
        cond()
    }

    #[test]
    fn resolve_token_prefers_flag_then_env() {
        assert_eq!(
            resolve_token_with(Some("flagtok"), None).unwrap(),
            "flagtok"
        );
        assert_eq!(
            resolve_token_with(None, Some("envtok".to_owned())).unwrap(),
            "envtok"
        );
        // empty flag falls through to env; empty/absent env -> error
        assert_eq!(
            resolve_token_with(Some(""), None).unwrap_err().code(),
            "relay_token"
        );
        assert_eq!(
            resolve_token_with(None, Some(String::new()))
                .unwrap_err()
                .code(),
            "relay_token"
        );
    }

    #[test]
    fn resolve_token_trims_surrounding_whitespace() {
        // A trailing newline (e.g. from a sops-decrypted token file) is trimmed
        // so it matches a clean token at the relay auth gate (bd-790088).
        assert_eq!(
            resolve_token_with(Some("flagtok\n"), None).unwrap(),
            "flagtok"
        );
        assert_eq!(
            resolve_token_with(None, Some("  envtok \n".to_owned())).unwrap(),
            "envtok"
        );
        // Whitespace-only is treated as empty -> error (no silent empty token).
        assert_eq!(
            resolve_token_with(Some("   \n"), None).unwrap_err().code(),
            "relay_token"
        );
    }

    #[test]
    fn end_to_end_trailing_newline_token_authenticates() {
        // Server has a clean token; a device + sender whose token carries a
        // trailing newline (the sops-file footgun) must still authenticate and
        // route, not be silently rejected (bd-790088).
        let listener = bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let opts = ServeOptions {
            token: "tok".into(),
            max_conns: Some(2),
        };
        thread::spawn(move || {
            let _ = serve(&listener, &opts);
        });

        let dev = TcpStream::connect(addr).unwrap();
        let mut dev_w = dev.try_clone().unwrap();
        let mut dev_r = BufReader::new(dev);
        let reg = ControlFrame::Register {
            v: 1,
            device: "harry-watch".into(),
            platform: "watchos".into(),
            token: "tok\n".into(),
        };
        write_line(&mut dev_w, &reg.json_line().unwrap()).unwrap();
        thread::sleep(Duration::from_millis(100));

        cast_via_relay(
            addr,
            "tok\n",
            "harry-watch",
            &[InjectionCommand::text("hi")],
        )
        .unwrap();

        let mut got = String::new();
        dev_r.read_line(&mut got).unwrap();
        let cmd = InjectionCommand::decode_line(got.trim()).unwrap();
        assert_eq!(cmd, InjectionCommand::text("hi"));
    }

    #[test]
    fn end_to_end_device_receives_cast() {
        let listener = bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let opts = ServeOptions {
            token: "tok".into(),
            max_conns: Some(2),
        };
        thread::spawn(move || {
            let _ = serve(&listener, &opts);
        });

        // DEVICE connects + registers, then reads one delivered bare command.
        let dev = TcpStream::connect(addr).unwrap();
        let mut dev_w = dev.try_clone().unwrap();
        let mut dev_r = BufReader::new(dev);
        let reg = ControlFrame::Register {
            v: 1,
            device: "harry-watch".into(),
            platform: "watchos".into(),
            token: "tok".into(),
        };
        write_line(&mut dev_w, &reg.json_line().unwrap()).unwrap();

        // Give the relay a moment to register the device before casting.
        thread::sleep(Duration::from_millis(100));

        // SENDER casts to the device.
        cast_via_relay(
            addr,
            "tok",
            "harry-watch",
            &[InjectionCommand::text("hello watch")],
        )
        .unwrap();

        // DEVICE should receive the bare InjectionCommand line.
        let mut got = String::new();
        dev_r.read_line(&mut got).unwrap();
        let cmd = InjectionCommand::decode_line(got.trim()).unwrap();
        assert_eq!(cmd, InjectionCommand::text("hello watch"));
    }

    #[test]
    fn end_to_end_unauthorized_register_is_rejected() {
        let listener = bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let opts = ServeOptions {
            token: "right".into(),
            max_conns: Some(1),
        };
        thread::spawn(move || {
            let _ = serve(&listener, &opts);
        });

        let dev = TcpStream::connect(addr).unwrap();
        let mut dev_w = dev.try_clone().unwrap();
        let mut dev_r = BufReader::new(dev);
        let reg = ControlFrame::Register {
            v: 1,
            device: "d".into(),
            platform: "ios".into(),
            token: "wrong".into(),
        };
        write_line(&mut dev_w, &reg.json_line().unwrap()).unwrap();
        let mut got = String::new();
        dev_r.read_line(&mut got).unwrap();
        assert!(
            got.contains("unauthorized"),
            "expected unauthorized, got: {got}"
        );
    }

    #[test]
    fn cast_via_relay_reports_delivered_to_registered_device() {
        // bd-0b4e45: a cast to a CONNECTED device reports acked + delivered==1.
        let listener = bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let opts = ServeOptions {
            token: "tok".into(),
            max_conns: Some(2),
        };
        let registry = Registry::new();
        let serve_registry = registry.clone();
        thread::spawn(move || {
            let _ = serve_with_registry(&listener, &opts, &serve_registry);
        });

        let dev = TcpStream::connect(addr).unwrap();
        let mut dev_w = dev.try_clone().unwrap();
        let reg = ControlFrame::Register {
            v: 1,
            device: "harry-watch".into(),
            platform: "watchos".into(),
            token: "tok".into(),
        };
        write_line(&mut dev_w, &reg.json_line().unwrap()).unwrap();
        assert!(
            wait_until(|| registry.is_registered("harry-watch"), 2000),
            "device should register"
        );

        let delivery =
            cast_via_relay(addr, "tok", "harry-watch", &[InjectionCommand::text("hi")]).unwrap();
        assert!(delivery.acked, "relay should ack the cast");
        assert_eq!(
            delivery.delivered, 1,
            "delivered to the one connected device"
        );
    }

    #[test]
    fn cast_via_relay_reports_zero_for_unconnected_device() {
        // bd-0b4e45: a cast to a device id that is NOT registered reports
        // acked == true with delivered == 0 (no more silent fire-and-forget).
        let listener = bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let opts = ServeOptions {
            token: "tok".into(),
            max_conns: Some(1),
        };
        thread::spawn(move || {
            let _ = serve(&listener, &opts);
        });

        let delivery =
            cast_via_relay(addr, "tok", "ghost-device", &[InjectionCommand::text("hi")]).unwrap();
        assert!(delivery.acked, "relay should still ack (with delivered=0)");
        assert_eq!(delivery.delivered, 0, "no device connected -> delivered 0");
    }

    #[test]
    fn cast_via_relay_does_not_hang_without_a_result() {
        // bd-0b4e45 backward-compat: a legacy relay / plain daemon that sends no
        // result must not hang the sender; cast_via_relay returns acked=false
        // after the bounded read timeout rather than blocking.
        let listener = bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        thread::spawn(move || {
            if let Ok((sock, _)) = listener.accept() {
                let mut r = BufReader::new(sock.try_clone().unwrap());
                let mut line = String::new();
                let _ = r.read_line(&mut line);
                // Hold the socket open WITHOUT replying, simulating a relay that
                // never sends a delivery result.
                thread::sleep(Duration::from_secs(2));
                drop(sock);
            }
        });

        let start = std::time::Instant::now();
        let delivery =
            cast_via_relay(addr, "tok", "local", &[InjectionCommand::text("hi")]).unwrap();
        assert!(!delivery.acked, "no result -> acked must be false");
        assert!(
            start.elapsed() < Duration::from_secs(2),
            "must return on the read timeout, not block for the full hold"
        );
    }

    #[test]
    fn device_ids_returns_sorted_registered() {
        // bd-bc8b5c: device_ids reflects the registered set, sorted.
        let reg = Registry::new();
        let (tx_b, _rx_b) = mpsc::channel();
        let (tx_a, _rx_a) = mpsc::channel();
        reg.register("phone", tx_b);
        reg.register("harry-watch", tx_a);
        assert_eq!(
            reg.device_ids(),
            vec!["harry-watch".to_owned(), "phone".to_owned()]
        );
        reg.unregister("phone");
        assert_eq!(reg.device_ids(), vec!["harry-watch".to_owned()]);
    }

    #[test]
    fn end_to_end_list_devices_returns_registered() {
        // bd-bc8b5c: a sender can query the hub for the connected device ids.
        let listener = bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let opts = ServeOptions {
            token: "tok".into(),
            max_conns: Some(3),
        };
        let registry = Registry::new();
        let serve_registry = registry.clone();
        thread::spawn(move || {
            let _ = serve_with_registry(&listener, &opts, &serve_registry);
        });

        // Hold the device connections open for the duration of the test.
        let mut held = Vec::new();
        for id in ["harry-watch", "phone"] {
            let dev = TcpStream::connect(addr).unwrap();
            let mut dev_w = dev.try_clone().unwrap();
            let reg = ControlFrame::Register {
                v: 1,
                device: id.into(),
                platform: "watchos".into(),
                token: "tok".into(),
            };
            write_line(&mut dev_w, &reg.json_line().unwrap()).unwrap();
            held.push(dev);
            held.push(dev_w);
        }
        assert!(
            wait_until(|| registry.device_ids().len() == 2, 2000),
            "both devices register"
        );

        let mut got = list_relay_devices(addr, "tok").unwrap();
        got.sort();
        assert_eq!(got, vec!["harry-watch".to_owned(), "phone".to_owned()]);
        drop(held);
    }

    #[test]
    fn list_devices_unauthorized_rejected() {
        // bd-bc8b5c: a wrong token on a list query is rejected.
        let listener = bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let opts = ServeOptions {
            token: "right".into(),
            max_conns: Some(1),
        };
        thread::spawn(move || {
            let _ = serve(&listener, &opts);
        });
        let err = list_relay_devices(addr, "wrong").unwrap_err();
        assert_eq!(err.code(), "relay_list");
    }

    #[test]
    fn list_relay_devices_no_hang_without_response() {
        // bd-bc8b5c backward-compat: a server that never replies (legacy relay)
        // yields an empty list within the bounded timeout, not a hang.
        let listener = bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        thread::spawn(move || {
            if let Ok((sock, _)) = listener.accept() {
                let mut r = BufReader::new(sock.try_clone().unwrap());
                let mut line = String::new();
                let _ = r.read_line(&mut line);
                thread::sleep(Duration::from_secs(2));
                drop(sock);
            }
        });
        let start = std::time::Instant::now();
        let devices = list_relay_devices(addr, "tok").unwrap();
        assert!(devices.is_empty(), "no reply -> empty list");
        assert!(
            start.elapsed() < Duration::from_secs(2),
            "must return on the read timeout, not hang"
        );
    }
}