Vectrum
Sign In
218 lines · 8.8 KB · Rust
//! Where a Device panel's devices are: on this machine, through the device
//! service, or on the device a remote thread runs on, through its remote
//! connection. Calls block; make them from a background thread.

use std::{
    collections::HashMap,
    sync::{
        Arc, Mutex,
        atomic::{AtomicU64, Ordering},
    },
    time::Duration,
};

use base64::Engine as _;
use gpui_kit::App;
use serde_json::Value;

use super::DeviceWatch;
use crate::device::{self, Action, AxElement, Detail, DeviceSession, HostStatus, State, stream::HubTarget};
use crate::remote::{ClientMessage, DeviceCall, client::ClientSender, device_proxy};

/// How long a call to another device may take; booting takes the longest.
const CALL_TIMEOUT: Duration = Duration::from_secs(240);

/// Another device's devices, reached through its remote connection.
pub struct RemoteDevices {
    pub env: String,
    /// The device's name, shown where this machine would say "This machine".
    pub name: String,
    sender: ClientSender,
    hub: HubTarget,
    pending: Mutex<HashMap<u64, smol::channel::Sender<Result<Value, String>>>>,
    next: AtomicU64,
}

impl RemoteDevices {
    /// `address` and `key` are the connection's, which the device's stream
    /// proxy answers on too.
    pub fn new(env: String, name: String, sender: ClientSender, address: &str, key: String) -> Self {
        let hub = HubTarget { address: crate::remote::with_port(address), prefix: device_proxy::PREFIX.into(), key: Some(key) };
        Self { env, name, sender, hub, pending: Mutex::default(), next: AtomicU64::new(1) }
    }

    /// Asks the device to send its device state now and on every change.
    pub fn watch(&self) {
        self.sender.send(&ClientMessage::Device { request: 0, call: DeviceCall::Watch });
    }

    /// Hands a reply to the call waiting for it.
    pub fn resolve(&self, request: u64, result: Result<Value, String>) {
        let waiting = self.pending.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).remove(&request);
        if let Some(waiting) = waiting {
            let _ = waiting.try_send(result);
        }
    }

    fn call(&self, call: DeviceCall) -> Result<Value, String> {
        let request = self.next.fetch_add(1, Ordering::Relaxed);
        let (tx, rx) = smol::channel::bounded(1);
        self.pending.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).insert(request, tx);
        self.sender.send(&ClientMessage::Device { request, call });
        let answer = smol::block_on(async {
            smol::future::or(async { rx.recv().await.ok() }, async {
                smol::Timer::after(CALL_TIMEOUT).await;
                None
            })
            .await
        });
        self.pending.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).remove(&request);
        answer.unwrap_or_else(|| {
            Err(format!("{} didn't answer. Its Vectrum Code may be too old for devices, or the connection dropped.", self.name))
        })
    }

    fn call_as<T: serde::de::DeserializeOwned>(&self, call: DeviceCall) -> Result<T, String> {
        serde_json::from_value(self.call(call)?).map_err(|err| format!("Unreadable answer from {}: {err}", self.name))
    }
}

/// A Device panel's devices.
#[derive(Clone)]
pub enum Backend {
    Local,
    Remote(Arc<RemoteDevices>),
}

fn text(err: anyhow::Error) -> String {
    format!("{err:#}")
}

impl Backend {
    /// The devices as last heard of. Another device's read as starting up
    /// until it has said.
    pub fn state(&self, cx: &App) -> State {
        let watch = cx.global::<DeviceWatch>();
        match self {
            Backend::Local => watch.state.clone(),
            Backend::Remote(remote) => watch.remote.get(&remote.env).cloned().unwrap_or_else(|| State {
                enabled: true,
                onboarding_completed: true,
                host_status: HostStatus::Starting,
                ..State::default()
            }),
        }
    }

    /// Where the device hub's streams are, if it runs.
    pub fn hub(&self) -> Option<HubTarget> {
        match self {
            Backend::Local => device::service().hub_origin().map(|origin| HubTarget::local(&origin)),
            Backend::Remote(remote) => Some(remote.hub.clone()),
        }
    }

    pub fn list(&self) {
        match self {
            Backend::Local => {
                device::service().list();
            }
            Backend::Remote(remote) => {
                let _ = remote.call(DeviceCall::List);
            }
        }
    }

    pub fn retry(&self) {
        match self {
            Backend::Local => {
                device::service().retry();
            }
            Backend::Remote(remote) => {
                let _ = remote.call(DeviceCall::Retry);
            }
        }
    }

    pub fn configure(&self, enabled: Option<bool>, agent_access: Option<bool>, onboarding: Option<bool>) -> Result<(), String> {
        match self {
            Backend::Local => device::service().configure(enabled, agent_access, onboarding).map(|_| ()).map_err(text),
            Backend::Remote(remote) => remote.call(DeviceCall::Configure { enabled, agent_access, onboarding }).map(|_| ()),
        }
    }

    pub fn open(&self, thread: &str, device: &str, boot: bool) -> Result<DeviceSession, String> {
        match self {
            Backend::Local => device::service().open(thread, device, boot).map_err(text),
            Backend::Remote(remote) => {
                remote.call_as(DeviceCall::Open { thread: thread.into(), device: device.into(), boot })
            }
        }
    }

    pub fn close(&self, thread: &str, device: Option<&str>, shutdown: bool) -> Result<(), String> {
        match self {
            Backend::Local => device::service().close(thread, device, shutdown).map_err(text),
            Backend::Remote(remote) => remote
                .call(DeviceCall::Close { thread: thread.into(), device: device.map(str::to_string), shutdown })
                .map(|_| ()),
        }
    }

    pub fn detail(&self, device: &str) -> Result<Detail, String> {
        match self {
            Backend::Local => device::service().detail(device).map_err(text),
            Backend::Remote(remote) => remote.call_as(DeviceCall::Detail { device: device.into() }),
        }
    }

    pub fn action(&self, device: &str, action: &Action) -> Result<Detail, String> {
        match self {
            Backend::Local => device::service().action(device, action).map_err(text),
            Backend::Remote(remote) => remote.call_as(DeviceCall::Action { device: device.into(), action: action.clone() }),
        }
    }

    pub fn screenshot(&self, device: &str) -> Result<Vec<u8>, String> {
        match self {
            Backend::Local => device::service().screenshot(device).map(|(_, png)| png).map_err(text),
            Backend::Remote(remote) => {
                let data: String = remote.call_as(DeviceCall::Screenshot { device: device.into() })?;
                base64::engine::general_purpose::STANDARD.decode(data).map_err(|err| err.to_string())
            }
        }
    }

    pub fn accessibility(&self, device: &str) -> Result<Vec<AxElement>, String> {
        match self {
            Backend::Local => device::service().accessibility(device).map_err(text),
            Backend::Remote(remote) => remote.call_as(DeviceCall::Accessibility { device: device.into() }),
        }
    }
}

/// Answers a client's device call on the host. Blocks.
pub fn answer(call: DeviceCall) -> Result<Value, String> {
    let service = device::service();
    let json = |value: Result<Value, serde_json::Error>| value.map_err(|err| err.to_string());
    match call {
        DeviceCall::Watch => json(serde_json::to_value(service.state())),
        DeviceCall::List => json(serde_json::to_value(service.list())),
        DeviceCall::Retry => json(serde_json::to_value(service.retry())),
        DeviceCall::Configure { enabled, agent_access, onboarding } => {
            json(serde_json::to_value(service.configure(enabled, agent_access, onboarding).map_err(text)?))
        }
        DeviceCall::Open { thread, device, boot } => json(serde_json::to_value(service.open(&thread, &device, boot).map_err(text)?)),
        DeviceCall::Close { thread, device, shutdown } => {
            service.close(&thread, device.as_deref(), shutdown).map_err(text)?;
            Ok(Value::Null)
        }
        DeviceCall::Detail { device } => json(serde_json::to_value(service.detail(&device).map_err(text)?)),
        DeviceCall::Action { device, action } => json(serde_json::to_value(service.action(&device, &action).map_err(text)?)),
        DeviceCall::Screenshot { device } => {
            let (_, png) = service.screenshot(&device).map_err(text)?;
            Ok(Value::String(base64::engine::general_purpose::STANDARD.encode(png)))
        }
        DeviceCall::Accessibility { device } => json(serde_json::to_value(service.accessibility(&device).map_err(text)?)),
    }
}