standard_plugin/testing/

machine.rs

//! The daemon world of the in-process host: children a plugin spawns
//! (recorded with their argv, environment and directory, answered from a
//! script) and file watches (checked with the host's own path rules).

use alloc::collections::VecDeque;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use std::path::Path;

use standard_plugin_manifest::paths::{Excludes, PathGrants, real};

use super::MockHost;
use crate::daemon::process::{Command, Stdio};
use crate::error::{Error, Result};
use crate::ui_runtime::Event;

/// What a scripted program does when the plugin spawns it: the output
/// it writes (delivered as `Event::ProcessOutput` for a piped stream) and
/// its exit status (`Event::ProcessExited`, and what `wait` answers).
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ProcessScript {
    pub stdout: Vec<u8>,
    pub stderr: Vec<u8>,
    pub status: i32,
}

impl ProcessScript {
    /// Exits with `status` and writes nothing.
    pub fn exit(status: i32) -> Self {
        Self {
            status,
            ..Self::default()
        }
    }

    pub fn stdout(mut self, bytes: impl AsRef<[u8]>) -> Self {
        self.stdout = bytes.as_ref().to_vec();
        self
    }

    pub fn stderr(mut self, bytes: impl AsRef<[u8]>) -> Self {
        self.stderr = bytes.as_ref().to_vec();
        self
    }
}

/// One child the plugin spawned, as the host saw it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MockProcess {
    /// Distinct per child, from 1000.
    pub pid: u32,
    pub program: String,
    pub args: Vec<String>,
    /// The environment the child got: the plugin's
    /// ([`MockHost::env`]), less `env_remove`, plus `env` (or only `env`
    /// with `clear_env`), sorted.
    pub env: Vec<(String, String)>,
    /// The directory it ran in: the command's, else the plugin's `HOME`.
    pub cwd: Option<String>,
    pub stdin: Vec<u8>,
    pub stdin_closed: bool,
    /// Its exit status once it exited (a script, [`DaemonHarness::exit`],
    /// or a kill: `-9`).
    ///
    /// [`DaemonHarness::exit`]: super::DaemonHarness::exit
    pub status: Option<i32>,
    pub killed: bool,
    stdout: Stdio,
    stderr: Stdio,
}

impl MockProcess {
    /// One variable of the child's environment.
    pub fn var(&self, name: &str) -> Option<&str> {
        self.env
            .iter()
            .find(|(key, _)| key == name)
            .map(|(_, value)| value.as_str())
    }
}

/// One watch the plugin holds.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MockWatch {
    pub handle: u32,
    pub path: String,
    pub recursive: bool,
    pub exclude: Vec<String>,
}

/// One pane the plugin opened (`daemon::panes::create`), as the host saw
/// it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MockPaneLaunch {
    /// The pane's id, `pane-1` onwards; its generation is 1.
    pub id: String,
    /// The project whose root holds `cwd`.
    pub project: String,
    pub cwd: String,
    pub command: Option<String>,
    pub env: Vec<(String, String)>,
    /// The title it asked for (`NewPane::title`); the plugin's id shows when
    /// none.
    pub title: Option<String>,
}

/// The daemon-world state of a [`MockHost`].
#[derive(Clone, Debug, Default)]
pub(crate) struct Machine {
    pub processes: Vec<MockProcess>,
    pub pane_launches: Vec<MockPaneLaunch>,
    pub scripts: Vec<(String, Option<Vec<String>>, ProcessScript)>,
    pub watches: Vec<MockWatch>,
    /// WebSockets the plugin opened (`daemon-net`), by handle.
    pub websockets: Vec<MockWebSocket>,
    next_pid: u32,
    next_watch: u32,
    /// Machine events for the plugin, delivered before the next drive.
    pub pending: VecDeque<Event>,
}

/// A WebSocket the plugin opened on the test host.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MockWebSocket {
    pub socket: u32,
    pub url: String,
    pub headers: Vec<(String, String)>,
    /// The text frames it sent.
    pub sent: Vec<String>,
    pub closed: bool,
}

/// Watches one plugin may hold, as the host allows.
const MAX_WATCHES: usize = 256;

impl MockHost {
    /// The WebSockets the plugin opened, in order.
    pub fn websockets(&self) -> &[MockWebSocket] {
        &self.machine.websockets
    }

    /// Delivers a socket's news to the plugin as the host would
    /// (`system.net.websocket`): `kind` `open`, `message` (with `text`) or
    /// `closed` (with `text` as the reason).
    pub fn websocket_event(&mut self, socket: u32, kind: &str, text: &str) {
        let payload = match kind {
            "message" => serde_json::json!({"socket": socket, "kind": kind, "text": text}),
            "closed" => serde_json::json!({"socket": socket, "kind": kind, "reason": text}),
            _ => serde_json::json!({"socket": socket, "kind": kind}),
        };
        self.machine.pending.push_back(Event::Plugin {
            name: String::from(crate::daemon::net::WEBSOCKET_EVENT),
            payload: crate::Json(payload.to_string()),
        });
    }

    pub(crate) fn websocket_open(
        &mut self,
        url: &str,
        headers: &[(String, String)],
    ) -> Result<u32> {
        let target = url
            .strip_prefix("wss://")
            .map(|rest| rest.split('/').next().unwrap_or_default())
            .ok_or_else(|| Error::Invalid(alloc::format!("{url:?} is not a wss:// URL")))?;
        let (host, port) = match target.rsplit_once(':') {
            Some((host, port)) => (host, port.parse::<u16>().unwrap_or(0)),
            None => (target, 443),
        };
        if !self.grants.socket_ports(host).contains(&port) {
            return Err(self.deny(standard_plugin_manifest::GrantDenied {
                grant: alloc::format!("socket.connect:{host}:{port}"),
            }));
        }
        let socket = self.machine.websockets.len() as u32 + 1;
        self.machine.websockets.push(MockWebSocket {
            socket,
            url: url.to_string(),
            headers: headers.to_vec(),
            sent: Vec::new(),
            closed: false,
        });
        Ok(socket)
    }

    pub(crate) fn websocket_send(&mut self, socket: u32, text: &str) -> Result<()> {
        let open = self
            .machine
            .websockets
            .iter_mut()
            .find(|open| open.socket == socket && !open.closed)
            .ok_or_else(|| Error::Invalid(alloc::format!("socket {socket} is not open")))?;
        open.sent.push(text.to_string());
        Ok(())
    }

    pub(crate) fn websocket_close(&mut self, socket: u32) {
        if let Some(open) = self
            .machine
            .websockets
            .iter_mut()
            .find(|open| open.socket == socket)
        {
            open.closed = true;
        }
    }

    /// Answers every spawn of `program` (as spawned, any arguments) with
    /// `script`. A later script for the same program wins.
    pub fn script(mut self, program: &str, script: ProcessScript) -> Self {
        self.machine
            .scripts
            .insert(0, (program.to_string(), None, script));
        self
    }

    /// Answers spawns of `program` whose arguments start with `args`.
    pub fn script_args(mut self, program: &str, args: &[&str], script: ProcessScript) -> Self {
        let args = args.iter().map(|arg| arg.to_string()).collect();
        self.machine
            .scripts
            .insert(0, (program.to_string(), Some(args), script));
        self
    }

    /// Every child the plugin spawned, in order.
    pub fn processes(&self) -> &[MockProcess] {
        &self.machine.processes
    }

    /// The watches the plugin holds.
    /// The panes the plugin opened, in order.
    pub fn pane_launches(&self) -> &[MockPaneLaunch] {
        &self.machine.pane_launches
    }

    /// Opens a pane as `standardd` does: in a directory inside a project
    /// root of this machine (`account.projects`, by whole components after
    /// symlinks; any machine's while [`MockHost::machine`] is unset), with
    /// the command and variables recorded, listed in `panes` at
    /// generation 1.
    pub(crate) fn create_pane(
        &mut self,
        cwd: &str,
        command: Option<&str>,
        env: &[(String, String)],
        title: Option<&str>,
    ) -> Result<(String, u64)> {
        if title.is_some_and(|title| {
            title.trim().is_empty()
                || title.chars().count() > 128
                || title.chars().any(char::is_control)
        }) {
            return Err(Error::Invalid(
                "a pane title is 1 to 128 characters with no control characters".to_string(),
            ));
        }
        if env.len() > 64
            || env.iter().any(|(name, value)| {
                name.is_empty() || name.contains(['=', '\0']) || value.contains('\0')
            })
        {
            return Err(Error::Invalid(
                "variables a pane cannot be given".to_string(),
            ));
        }
        let path = Path::new(cwd);
        if !path.is_absolute() {
            return Err(Error::Invalid(alloc::format!(
                "{cwd} is not an absolute path"
            )));
        }
        let resolved = real(path);
        let machine = self.machine_id.clone();
        let project = self
            .account
            .projects
            .iter()
            .filter(|project| machine.as_ref().is_none_or(|id| &project.machine == id))
            .filter(|project| resolved.starts_with(real(Path::new(&project.path))))
            .max_by_key(|project| real(Path::new(&project.path)).components().count())
            .ok_or_else(|| {
                Error::Invalid(alloc::format!(
                    "{cwd} is not inside a project root on this machine"
                ))
            })?
            .id
            .clone();
        let id = alloc::format!("pane-{}", self.machine.pane_launches.len() + 1);
        self.panes.push(crate::api::account::Pane {
            id: id.clone(),
            generation: 1,
            machine: machine.unwrap_or_default(),
            project: Some(project.clone()),
            title: title.map_or_else(|| self.plugin_id.clone(), str::to_string),
            cwd: cwd.to_string(),
            cols: 80,
            rows: 24,
            ..Default::default()
        });
        self.machine.pane_launches.push(MockPaneLaunch {
            id: id.clone(),
            project,
            cwd: cwd.to_string(),
            command: command.map(str::to_string),
            env: env.to_vec(),
            title: title.map(str::to_string),
        });
        Ok((id, 1))
    }

    pub fn watches(&self) -> &[MockWatch] {
        &self.machine.watches
    }

    pub(crate) fn spawn(&mut self, command: &Command) -> Result<u32> {
        if command.program.is_empty() {
            return Err(Error::Invalid("no program".into()));
        }
        if !self.grants.allows("machine.full") {
            self.check(&alloc::format!("process.exec:{}", command.program))?;
        }
        let mut env: Vec<(String, String)> = if command.clear_env {
            Vec::new()
        } else {
            self.environment
                .iter()
                .filter(|(name, _)| !command.env_remove.contains(name))
                .cloned()
                .collect()
        };
        for (name, value) in &command.env {
            env.retain(|(key, _)| key != name);
            env.push((name.clone(), value.clone()));
        }
        env.sort();
        let cwd = command.cwd.clone().or_else(|| {
            self.environment
                .iter()
                .find(|(name, _)| name == "HOME")
                .map(|(_, home)| home.clone())
        });
        let pid = 1000 + self.machine.next_pid;
        self.machine.next_pid += 1;
        let script = self
            .machine
            .scripts
            .iter()
            .find(|(program, args, _)| {
                *program == command.program
                    && args
                        .as_ref()
                        .is_none_or(|args| command.args.starts_with(args))
            })
            .map(|(_, _, script)| script.clone());
        self.machine.processes.push(MockProcess {
            pid,
            program: command.program.clone(),
            args: command.args.clone(),
            env,
            cwd,
            stdin: Vec::new(),
            stdin_closed: false,
            status: None,
            killed: false,
            stdout: command.stdout,
            stderr: command.stderr,
        });
        if let Some(script) = script {
            self.output(pid, false, &script.stdout);
            self.output(pid, true, &script.stderr);
            self.exited(pid, script.status);
        }
        Ok(pid)
    }

    fn process(&mut self, pid: u32) -> Result<&mut MockProcess> {
        self.machine
            .processes
            .iter_mut()
            .find(|process| process.pid == pid)
            .ok_or_else(|| Error::Invalid(alloc::format!("no child {pid} of this plugin")))
    }

    /// Queues output of `pid` for the plugin, when that stream is piped.
    pub(crate) fn output(&mut self, pid: u32, stderr: bool, bytes: &[u8]) {
        let Ok(process) = self.process(pid) else {
            return;
        };
        let piped = if stderr {
            process.stderr == Stdio::Piped
        } else {
            process.stdout == Stdio::Piped
        };
        if piped && !bytes.is_empty() {
            self.machine.pending.push_back(Event::ProcessOutput {
                pid,
                stderr,
                bytes: bytes.to_vec(),
            });
        }
    }

    /// `pid` exits with `status`: the plugin hears it after its output.
    pub(crate) fn exited(&mut self, pid: u32, status: i32) {
        let Ok(process) = self.process(pid) else {
            return;
        };
        if process.status.is_some() {
            return;
        }
        process.status = Some(status);
        self.machine
            .pending
            .push_back(Event::ProcessExited { pid, status });
    }

    pub(crate) fn wait(&mut self, pid: u32) -> Result<i32> {
        let process = self.process(pid)?;
        process.status.ok_or_else(|| {
            Error::Unavailable(alloc::format!(
                "child {pid} ({}) has no script and would block forever: script it \
                 (MockHost::script) or end it (DaemonHarness::exit) first",
                process.program
            ))
        })
    }

    pub(crate) fn try_wait(&mut self, pid: u32) -> Result<Option<i32>> {
        Ok(self.process(pid)?.status)
    }

    pub(crate) fn write(&mut self, pid: u32, bytes: &[u8]) -> Result<()> {
        let process = self.process(pid)?;
        if process.stdin_closed {
            return Err(Error::Invalid(alloc::format!(
                "child {pid}'s stdin is closed"
            )));
        }
        process.stdin.extend_from_slice(bytes);
        Ok(())
    }

    pub(crate) fn close_stdin(&mut self, pid: u32) -> Result<()> {
        self.process(pid)?.stdin_closed = true;
        Ok(())
    }

    pub(crate) fn kill(&mut self, pid: u32) -> Result<()> {
        let process = self.process(pid)?;
        if process.status.is_none() {
            process.killed = true;
            self.exited(pid, -9);
        }
        Ok(())
    }

    /// `daemon-watch.watch` with the host's own rules: an absolute path a
    /// read grant covers after resolving symlinks (or `machine.full`), and
    /// exclude globs the host accepts.
    pub(crate) fn watch(&mut self, path: &str, recursive: bool, exclude: &[String]) -> Result<u32> {
        let named = Path::new(path);
        if !named.is_absolute() {
            return Err(Error::Invalid(alloc::format!(
                "{path} is not an absolute path"
            )));
        }
        if !PathGrants::from_grants(&self.grants).allows(&real(named), false) {
            return Err(self.deny(standard_plugin_manifest::GrantDenied {
                grant: alloc::format!("fs.read:{path}"),
            }));
        }
        Excludes::new(exclude).map_err(|error| Error::Invalid(error.to_string()))?;
        if self.machine.watches.len() >= MAX_WATCHES {
            return Err(Error::RateLimited);
        }
        self.machine.next_watch += 1;
        let handle = self.machine.next_watch;
        self.machine.watches.push(MockWatch {
            handle,
            path: path.to_string(),
            recursive,
            exclude: exclude.to_vec(),
        });
        Ok(handle)
    }

    pub(crate) fn unwatch(&mut self, handle: u32) -> Result<()> {
        let before = self.machine.watches.len();
        self.machine.watches.retain(|watch| watch.handle != handle);
        if self.machine.watches.len() == before {
            return Err(Error::Invalid(alloc::format!("no watch {handle}")));
        }
        Ok(())
    }

    /// Queues `Event::FileChanged` for every watch that hears a change at
    /// `path`, as the host decides it (depth and exclude globs). Returns
    /// how many watches hear it.
    pub(crate) fn changed(&mut self, path: &str) -> usize {
        let path = Path::new(path);
        let mut heard = 0;
        for watch in &self.machine.watches {
            let Ok(rest) = path.strip_prefix(&watch.path) else {
                continue;
            };
            if !watch.recursive && rest.components().count() > 1 {
                continue;
            }
            let excluded = Excludes::new(&watch.exclude).is_ok_and(|globs| globs.excludes(rest));
            if excluded {
                continue;
            }
            heard += 1;
            self.machine.pending.push_back(Event::FileChanged {
                watch: watch.handle,
                paths: alloc::vec![path.to_string_lossy().into_owned()],
            });
        }
        heard
    }
}