storebound/apps/desktop/src-tauri/src/sidecar.rs

//! The storebound sidecar: `storebound serve --stdio`, started from Rust only (no capability grants the shell
//! plugin to the webview), with the fixed arguments in policy::SIDECAR_ARGS. One process per app run, started on
//! first use and again after it stops. Its first request is `session.hello` with a fresh 32-byte secret, so only this
//! process can mint the confirm tokens the sidecar accepts for real writes.
//!
//! Jobs: the webview may tag a call with a short job id. Progress notifications of that request then carry the job id
//! (so the right toast or screen shows them), and `cancel_job` sends `$/cancelRequest` for every request of the job.
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use serde::Serialize;
use serde_json::{json, Value};
use tauri::{AppHandle, Emitter, Manager, Runtime};
use tauri_plugin_shell::process::{CommandChild, CommandEvent};
use tauri_plugin_shell::ShellExt;
use tokio::sync::oneshot;

use crate::confirm;
use crate::policy::{self, Kind};

/// A failed call, as the webview sees it.
#[derive(Debug, Clone, Serialize)]
pub struct Fail {
    pub code: i64,
    pub message: String,
    #[serde(skip_serializing_if = "Value::is_null")]
    pub data: Value,
}

impl Fail {
    pub fn new(code: i64, message: impl Into<String>) -> Fail {
        Fail { code, message: message.into(), data: Value::Null }
    }
    /// The host refused (window, origin, params, state).
    pub fn host(message: impl Into<String>) -> Fail {
        Fail::new(-1, message)
    }
}

impl From<String> for Fail {
    fn from(s: String) -> Fail {
        Fail::host(s)
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct Descriptor {
    pub method: String,
    pub title: String,
    pub description: String,
    pub kind: String,
    pub network: bool,
    pub destructive: bool,
    pub json: bool,
}

type Pending = Arc<Mutex<HashMap<u64, oneshot::Sender<Result<Value, Fail>>>>>;
/// Request id -> the webview's job id.
type Jobs = Arc<Mutex<HashMap<u64, String>>>;

/// Job ids come from the webview and are echoed back in events: short and plain only.
pub fn valid_job(job: &str) -> bool {
    !job.is_empty() && job.len() <= 64 && job.bytes().all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b':' | b'.'))
}

pub struct Running {
    child: Mutex<Option<CommandChild>>,
    secret: [u8; 32],
    pending: Pending,
    jobs: Jobs,
    next: AtomicU64,
    alive: Arc<AtomicBool>,
    pub version: String,
    pub protocol: String,
    pub methods: Vec<Descriptor>,
}

impl Running {
    pub fn kind(&self, method: &str) -> Option<Kind> {
        self.methods.iter().find(|d| d.method == method).and_then(|d| Kind::parse(&d.kind))
    }
    pub fn descriptor(&self, method: &str) -> Option<&Descriptor> {
        self.methods.iter().find(|d| d.method == method)
    }
    pub fn mint(&self, method: &str, params: &Value) -> Result<String, String> {
        let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map_err(|e| e.to_string())?.as_millis() as u64;
        confirm::mint(&self.secret, method, params, now, &confirm::new_nonce()?)
    }

    fn send(&self, id: u64, method: &str, params: &Value) -> Result<oneshot::Receiver<Result<Value, Fail>>, Fail> {
        if !self.alive.load(Ordering::SeqCst) {
            return Err(Fail::host("the storebound sidecar stopped"));
        }
        let (tx, rx) = oneshot::channel();
        self.pending.lock().map_err(|_| Fail::host("sidecar state poisoned"))?.insert(id, tx);
        let mut line = serde_json::to_vec(&json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params })).map_err(|e| Fail::host(e.to_string()))?;
        line.push(b'\n');
        let mut child = self.child.lock().map_err(|_| Fail::host("sidecar state poisoned"))?;
        match child.as_mut() {
            Some(c) => c.write(&line).map_err(|e| Fail::host(format!("sidecar write: {e}")))?,
            None => return Err(Fail::host("the storebound sidecar stopped")),
        }
        Ok(rx)
    }

    pub async fn call(&self, method: &str, params: Value) -> Result<Value, Fail> {
        self.call_job(method, params, None).await
    }

    /// A call whose progress events carry `job`, and which `cancel_job` can stop.
    pub async fn call_job(&self, method: &str, params: Value, job: Option<&str>) -> Result<Value, Fail> {
        let id = self.next.fetch_add(1, Ordering::SeqCst);
        if let Some(j) = job {
            if let Ok(mut m) = self.jobs.lock() {
                m.insert(id, j.to_string());
            }
        }
        let res = match self.send(id, method, &params) {
            Ok(rx) => rx.await.map_err(|_| Fail::host("the storebound sidecar stopped")).and_then(|r| r),
            Err(e) => Err(e),
        };
        if job.is_some() {
            if let Ok(mut m) = self.jobs.lock() {
                m.remove(&id);
            }
        }
        res
    }

    /// Cancels every running request of a job (`$/cancelRequest`; they answer -32800). Returns how many.
    pub fn cancel_job(&self, job: &str) -> usize {
        let ids: Vec<u64> = self.jobs.lock().map(|m| m.iter().filter(|(_, j)| j.as_str() == job).map(|(id, _)| *id).collect()).unwrap_or_default();
        for id in &ids {
            self.notify("$/cancelRequest", json!({ "id": id }));
        }
        ids.len()
    }

    pub fn notify(&self, method: &str, params: Value) {
        let mut line = match serde_json::to_vec(&json!({ "jsonrpc": "2.0", "method": method, "params": params })) {
            Ok(l) => l,
            Err(_) => return,
        };
        line.push(b'\n');
        if let Ok(mut c) = self.child.lock() {
            if let Some(c) = c.as_mut() {
                let _ = c.write(&line);
            }
        }
    }

    pub fn alive(&self) -> bool {
        self.alive.load(Ordering::SeqCst)
    }

    pub fn kill(&self) {
        self.alive.store(false, Ordering::SeqCst);
        if let Ok(mut c) = self.child.lock() {
            if let Some(c) = c.take() {
                let _ = c.kill();
            }
        }
    }
}

#[derive(Default)]
pub struct Sidecar {
    current: tokio::sync::Mutex<Option<Arc<Running>>>,
}

/// Progress of a running request, forwarded to the main window as `sidecar-progress`.
#[derive(Clone, Serialize)]
pub struct Progress {
    pub params: Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub job: Option<String>,
}

impl Sidecar {
    /// The running sidecar, started (and greeted) if needed.
    pub async fn get<R: Runtime>(&self, app: &AppHandle<R>) -> Result<Arc<Running>, Fail> {
        let mut cur = self.current.lock().await;
        if let Some(r) = cur.as_ref() {
            if r.alive() {
                return Ok(r.clone());
            }
        }
        let r = match start(app).await {
            Ok(r) => Arc::new(r),
            Err(e) => {
                crate::logs::host("error", &format!("sidecar did not start: {}", e.message));
                return Err(e);
            }
        };
        crate::logs::host("info", &format!("sidecar started: storebound {} ({})", r.version, r.protocol));
        *cur = Some(r.clone());
        Ok(r)
    }

    /// The running sidecar, if there is one (never starts it).
    pub async fn current(&self) -> Option<Arc<Running>> {
        self.current.lock().await.as_ref().filter(|r| r.alive()).cloned()
    }

    /// Stops the sidecar; the next call starts a fresh one (after a tool install, so it sees the new tool paths).
    pub async fn shutdown(&self) {
        if let Some(r) = self.current.lock().await.take() {
            r.kill();
        }
    }

    /// Non-async variant for the exit path (the event loop is ending).
    pub fn shutdown_now(&self) {
        if let Ok(mut cur) = self.current.try_lock() {
            if let Some(r) = cur.take() {
                r.kill();
            }
        }
    }
}

/// Tools the app manages in its data folder (Dependencies screen): `tools/bin` (ffmpeg, ffprobe) goes in front of
/// PATH for every storebound command, `tools/chrome/<executable>` is the Chrome for Testing it renders with. Only
/// paths that exist are passed.
pub fn managed_tool_env(data_dir: &std::path::Path) -> Vec<(String, String)> {
    let mut env = Vec::new();
    let bin = data_dir.join("tools").join("bin");
    if bin.is_dir() {
        env.push(("STOREBOUND_TOOL_PATH".to_string(), bin.to_string_lossy().to_string()));
    }
    let chrome = data_dir.join("tools").join("chrome").join(policy::chrome_executable());
    if chrome.is_file() {
        env.push(("STOREBOUND_CHROME".to_string(), chrome.to_string_lossy().to_string()));
    }
    env
}

async fn start<R: Runtime>(app: &AppHandle<R>) -> Result<Running, Fail> {
    let tool_env = app.path().app_data_dir().map(|d| managed_tool_env(&d)).unwrap_or_default();
    let cmd = app
        .shell()
        .sidecar(policy::SIDECAR)
        .map_err(|e| Fail::host(format!("sidecar missing: {e} (bun tools/build-sidecars.ts --host)")))?
        .args(policy::SIDECAR_ARGS)
        .envs(tool_env)
        .set_raw_out(true)
        .set_process_group(true);
    let (mut rx, child) = cmd.spawn().map_err(|e| Fail::host(format!("sidecar did not start: {e}")))?;
    let pending: Pending = Arc::new(Mutex::new(HashMap::new()));
    let jobs: Jobs = Arc::new(Mutex::new(HashMap::new()));
    let alive = Arc::new(AtomicBool::new(true));

    let (p, a, handle, jb) = (pending.clone(), alive.clone(), app.clone(), jobs.clone());
    tauri::async_runtime::spawn(async move {
        let mut out: Vec<u8> = Vec::new();
        let mut err: Vec<u8> = Vec::new();
        // the last stderr lines, written to host.log when the sidecar stops
        let mut tail: std::collections::VecDeque<String> = std::collections::VecDeque::new();
        let mut end = String::from("closed");
        while let Some(ev) = rx.recv().await {
            match ev {
                CommandEvent::Stdout(bytes) => {
                    out.extend_from_slice(&bytes);
                    while let Some(nl) = out.iter().position(|b| *b == b'\n') {
                        let line: Vec<u8> = out.drain(..=nl).collect();
                        route(&handle, &p, &jb, &line);
                    }
                }
                CommandEvent::Stderr(bytes) => {
                    err.extend_from_slice(&bytes);
                    while let Some(nl) = err.iter().position(|b| *b == b'\n') {
                        let line: Vec<u8> = err.drain(..=nl).collect();
                        let text = String::from_utf8_lossy(&line).to_string();
                        if cfg!(debug_assertions) {
                            eprint!("[sidecar] {text}");
                        }
                        if tail.len() == 20 {
                            tail.pop_front();
                        }
                        tail.push_back(text.trim_end().to_string());
                    }
                }
                CommandEvent::Terminated(t) => {
                    end = format!("exit code {:?}, signal {:?}", t.code, t.signal);
                    break;
                }
                CommandEvent::Error(e) => {
                    end = format!("error: {e}");
                    break;
                }
                _ => {}
            }
        }
        let expected = !a.swap(false, Ordering::SeqCst);
        let tail: Vec<String> = tail.into_iter().collect();
        crate::logs::host(if expected { "info" } else { "warn" }, &format!("sidecar stopped ({end}){}", if tail.is_empty() { String::new() } else { format!("\n{}", tail.join("\n")) }));
        if let Ok(mut map) = p.lock() {
            for (_, tx) in map.drain() {
                let _ = tx.send(Err(Fail::host("the storebound sidecar stopped")));
            }
        }
    });

    let secret = confirm::new_secret()?;
    let mut running = Running {
        child: Mutex::new(Some(child)),
        secret,
        pending,
        jobs,
        next: AtomicU64::new(1),
        alive,
        version: String::new(),
        protocol: String::new(),
        methods: Vec::new(),
    };
    let hello = tokio::time::timeout(
        Duration::from_secs(20),
        running.call("session.hello", json!({ "secret": confirm::hex(&secret), "client": { "name": "storebound-desktop", "version": env!("CARGO_PKG_VERSION") } })),
    )
    .await
    .map_err(|_| Fail::host("the storebound sidecar did not answer session.hello"))??;
    running.protocol = hello["protocol"].as_str().unwrap_or_default().to_string();
    running.version = hello["server"]["version"].as_str().unwrap_or_default().to_string();
    if running.protocol != "storebound-desktop/1" {
        running.kill();
        return Err(Fail::host(format!("unsupported sidecar protocol '{}'", running.protocol)));
    }
    running.methods = hello["methods"]
        .as_array()
        .map(|a| {
            a.iter()
                .map(|m| Descriptor {
                    method: m["method"].as_str().unwrap_or_default().into(),
                    title: m["title"].as_str().unwrap_or_default().into(),
                    description: m["description"].as_str().unwrap_or_default().into(),
                    kind: m["kind"].as_str().unwrap_or_default().into(),
                    network: m["network"].as_bool().unwrap_or(true),
                    destructive: m["destructive"].as_bool().unwrap_or(true),
                    json: m["json"].as_bool().unwrap_or(false),
                })
                .collect()
        })
        .unwrap_or_default();
    Ok(running)
}

fn route<R: Runtime>(app: &AppHandle<R>, pending: &Pending, jobs: &Jobs, line: &[u8]) {
    let Ok(msg) = serde_json::from_slice::<Value>(line) else { return };
    if let Some(id) = msg.get("id").and_then(Value::as_u64) {
        let tx = pending.lock().ok().and_then(|mut m| m.remove(&id));
        if let Some(tx) = tx {
            let res = match msg.get("error") {
                Some(e) => Err(Fail {
                    code: e["code"].as_i64().unwrap_or(-32603),
                    message: e["message"].as_str().unwrap_or("error").to_string(),
                    data: e.get("data").cloned().unwrap_or(Value::Null),
                }),
                None => Ok(msg.get("result").cloned().unwrap_or(Value::Null)),
            };
            let _ = tx.send(res);
        }
    } else if msg.get("method").and_then(Value::as_str) == Some("progress") {
        let job = msg["params"]["requestId"].as_u64().and_then(|id| jobs.lock().ok().and_then(|m| m.get(&id).cloned()));
        let _ = app.emit_to(policy::MAIN, "sidecar-progress", Progress { params: msg["params"].clone(), job });
    }
}