vamp/src/ruby.rs

use crate::editor::{Context, Effect, Feature};
use serde::{Deserialize, Serialize};
use std::env;
use std::io::{self, BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::sync::mpsc::{self, Receiver, RecvTimeoutError, SyncSender};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

const RUNNER: &str = include_str!("ruby_host.rb");
const MAX_FRAME: usize = 4 * 1024 * 1024;
const STARTUP_TIMEOUT: Duration = Duration::from_secs(2);
const INVOKE_TIMEOUT: Duration = Duration::from_millis(500);

#[derive(Debug)]
pub struct Config {
    pub plugins: Vec<PathBuf>,
    pub bindings: Vec<Binding>,
    pub leader: String,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            plugins: Vec::new(),
            bindings: Vec::new(),
            leader: "space".into(),
        }
    }
}

/// Relative plugin paths are resolved against the configuration file's directory.
pub fn load_config(path: &Path) -> Result<Config, String> {
    let deadline = Instant::now() + STARTUP_TIMEOUT;
    let path = absolute_file(path)?;
    let mut host = Host::spawn("config", &path)?;
    let response = host.receive(deadline)?;
    match response {
        Response::Config {
            plugins,
            bindings,
            leader,
            ..
        } => {
            let mut config = Config::default();
            let directory = path.parent().expect("absolute file has a parent");
            for plugin in plugins {
                if plugin.is_empty() || plugin.contains('\0') {
                    return Err("Ruby config contains an invalid plugin path".into());
                }
                config.plugins.push(directory.join(plugin));
            }
            config.bindings = bindings;
            config.leader = leader;
            Ok(config)
        }
        Response::Error {
            id: None, message, ..
        } => Err(format!("Ruby config: {message}")),
        _ => Err("Ruby config returned an unexpected protocol frame".into()),
    }
}

pub struct RubyPlugin {
    host: Host,
    next_id: u64,
}

impl RubyPlugin {
    pub fn start(path: &Path) -> Result<(Self, Feature), String> {
        let deadline = Instant::now() + STARTUP_TIMEOUT;
        let path = absolute_file(path)?;
        let mut host = Host::spawn("plugin", &path)?;
        match host.receive(deadline)? {
            Response::Manifest {
                name,
                requires,
                commands,
                ..
            } => {
                if name.is_empty()
                    || requires.iter().any(String::is_empty)
                    || commands.iter().any(String::is_empty)
                {
                    return Err("Ruby plugin manifest contains empty names".into());
                }
                let feature = Feature {
                    name,
                    required_capabilities: requires,
                    commands,
                    core: false,
                };
                Ok((Self { host, next_id: 1 }, feature))
            }
            Response::Error {
                id: None, message, ..
            } => Err(format!("Ruby plugin startup: {message}")),
            _ => Err("Ruby plugin returned an unexpected startup frame".into()),
        }
    }

    /// Context coordinates are one-based; effects are queued until a full valid reply arrives.
    /// Any failure shuts down this host. Starting a fresh plugin is the recovery path.
    pub fn invoke(
        &mut self,
        command: &str,
        args: &str,
        context: &Context,
    ) -> Result<Vec<Effect>, String> {
        let deadline = Instant::now() + INVOKE_TIMEOUT;
        let result = self.invoke_inner(command, args, context, deadline);
        if result.is_err() {
            self.disable();
        }
        result
    }

    fn invoke_inner(
        &mut self,
        command: &str,
        args: &str,
        context: &Context,
        deadline: Instant,
    ) -> Result<Vec<Effect>, String> {
        let id = self.next_id;
        self.next_id = id.checked_add(1).ok_or("Ruby request IDs exhausted")?;
        let request = Request {
            version: 1,
            kind: "invoke",
            id,
            command,
            args,
            context,
        };
        self.host.send(&request, deadline)?;
        match self.host.receive(deadline)? {
            Response::Result {
                id: response_id,
                effects,
                ..
            } if response_id == id => Ok(effects
                .into_iter()
                .map(|effect| match effect {
                    WireEffect::Insert { text } => Effect::Insert(text),
                    WireEffect::Replace { text } => Effect::Replace(text),
                    WireEffect::Message { text } => Effect::Message(text),
                })
                .collect()),
            Response::Error {
                id: Some(response_id),
                message,
                ..
            } if response_id == id => Err(format!("Ruby plugin: {message}")),
            _ => Err("Ruby plugin response does not match the invocation".into()),
        }
    }

    pub fn disable(&mut self) {
        self.host.shutdown();
    }
}

fn absolute_file(path: &Path) -> Result<PathBuf, String> {
    if path.is_absolute() {
        Ok(path.to_path_buf())
    } else {
        env::current_dir()
            .map(|directory| directory.join(path))
            .map_err(|error| format!("Cannot resolve Ruby file path: {error}"))
    }
}

#[derive(Serialize)]
struct Request<'a> {
    version: u8,
    #[serde(rename = "type")]
    kind: &'static str,
    id: u64,
    command: &'a str,
    args: &'a str,
    context: &'a Context,
}

#[derive(Deserialize)]
#[serde(tag = "type", deny_unknown_fields)]
enum Response {
    #[serde(rename = "manifest")]
    Manifest {
        version: u8,
        name: String,
        requires: Vec<String>,
        commands: Vec<String>,
    },
    #[serde(rename = "config")]
    Config {
        version: u8,
        plugins: Vec<String>,
        bindings: Vec<Binding>,
        leader: String,
    },
    #[serde(rename = "result")]
    Result {
        version: u8,
        id: u64,
        effects: Vec<WireEffect>,
    },
    #[serde(rename = "error")]
    Error {
        version: u8,
        #[serde(default, deserialize_with = "optional_request_id")]
        id: Option<u64>,
        message: String,
    },
}

fn optional_request_id<'de, D: serde::Deserializer<'de>>(
    deserializer: D,
) -> Result<Option<u64>, D::Error> {
    u64::deserialize(deserializer).map(Some)
}

impl Response {
    fn version(&self) -> u8 {
        match self {
            Self::Manifest { version, .. }
            | Self::Config { version, .. }
            | Self::Result { version, .. }
            | Self::Error { version, .. } => *version,
        }
    }
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Binding {
    pub mode: String,
    pub key: String,
    pub command: Option<String>,
}

#[derive(Deserialize)]
#[serde(tag = "type", deny_unknown_fields)]
enum WireEffect {
    #[serde(rename = "insert")]
    Insert { text: String },
    #[serde(rename = "replace")]
    Replace { text: String },
    #[serde(rename = "message")]
    Message { text: String },
}

struct Host {
    child: Option<Child>,
    requests: Option<SyncSender<Vec<u8>>>,
    writes: Receiver<Result<(), String>>,
    frames: Receiver<Result<Response, String>>,
    reader: Option<JoinHandle<()>>,
    writer: Option<JoinHandle<()>>,
}

impl Host {
    fn spawn(mode: &str, path: &Path) -> Result<Self, String> {
        let executable = env::var_os("VAMP_RUBY").unwrap_or_else(|| "ruby".into());
        let mut command = Command::new(&executable);
        command
            .args(["-e", RUNNER, "--", mode])
            .arg(path)
            .current_dir(path.parent().expect("absolute file has a parent"))
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null());
        #[cfg(unix)]
        {
            use std::os::unix::process::CommandExt;
            command.process_group(0);
        }
        let mut child = command.spawn().map_err(|error| {
            format!(
                "Ruby runtime {:?} unavailable or failed to start: {error}",
                executable
            )
        })?;
        let mut stdin = child.stdin.take().expect("piped child stdin");
        let stdout = child.stdout.take().expect("piped child stdout");
        let (request_tx, request_rx) = mpsc::sync_channel::<Vec<u8>>(1);
        let (write_tx, write_rx) = mpsc::sync_channel(1);
        let (frame_tx, frame_rx) = mpsc::sync_channel(1);
        let mut host = Self {
            child: Some(child),
            requests: Some(request_tx),
            writes: write_rx,
            frames: frame_rx,
            reader: None,
            writer: None,
        };
        host.reader = Some(
            thread::Builder::new()
                .name("vamp-ruby-reader".into())
                .spawn(move || {
                    let mut reader = BufReader::new(stdout);
                    loop {
                        let frame = read_frame(&mut reader).and_then(|bytes| {
                            serde_json::from_slice::<Response>(&bytes)
                                .map_err(|error| format!("Malformed Ruby protocol frame: {error}"))
                        });
                        let failed = frame.is_err();
                        // Never block on a flooded host; the bounded queue disconnects after draining.
                        if frame_tx.try_send(frame).is_err() || failed {
                            break;
                        }
                    }
                })
                .map_err(|error| format!("Cannot start Ruby reader: {error}"))?,
        );
        host.writer = Some(
            thread::Builder::new()
                .name("vamp-ruby-writer".into())
                .spawn(move || {
                    while let Ok(bytes) = request_rx.recv() {
                        let result = stdin
                            .write_all(&bytes)
                            .and_then(|_| stdin.flush())
                            .map_err(|error| format!("Ruby host write failed: {error}"));
                        let failed = result.is_err();
                        if write_tx.try_send(result).is_err() || failed {
                            break;
                        }
                    }
                })
                .map_err(|error| format!("Cannot start Ruby writer: {error}"))?,
        );
        Ok(host)
    }

    fn send(&mut self, request: &Request<'_>, deadline: Instant) -> Result<(), String> {
        let sender = self.requests.as_ref().ok_or("Ruby plugin is disabled")?;
        match self.frames.try_recv() {
            Ok(_) => return Err("Ruby host emitted an unsolicited protocol frame".into()),
            Err(mpsc::TryRecvError::Disconnected) => {
                return Err("Ruby host crashed or closed its output".into());
            }
            Err(mpsc::TryRecvError::Empty) => {}
        }
        let mut frame = FrameWriter(Vec::new());
        serde_json::to_writer(&mut frame, request)
            .map_err(|error| format!("Cannot encode bounded Ruby request: {error}"))?;
        frame.0.push(b'\n');
        sender
            .try_send(frame.0)
            .map_err(|_| "Ruby host request channel failed")?;
        receive_until(&self.writes, deadline, "write")?
    }

    fn receive(&mut self, deadline: Instant) -> Result<Response, String> {
        let response = receive_until(&self.frames, deadline, "response")??;
        if response.version() != 1 {
            return Err("Unsupported Ruby protocol version".into());
        }
        Ok(response)
    }

    fn shutdown(&mut self) {
        self.requests.take();
        if let Some(mut child) = self.child.take() {
            // Descendants can inherit pipes; kill the host group before joining IO threads.
            #[cfg(unix)]
            unsafe {
                kill(-(child.id() as i32), 9);
            }
            let _ = child.kill();
            let _ = child.wait();
        }
        if let Some(writer) = self.writer.take() {
            let _ = writer.join();
        }
        if let Some(reader) = self.reader.take() {
            let _ = reader.join();
        }
    }
}

impl Drop for Host {
    fn drop(&mut self) {
        self.shutdown();
    }
}

#[cfg(unix)]
unsafe extern "C" {
    fn kill(pid: i32, signal: i32) -> i32;
}

fn receive_until<T>(receiver: &Receiver<T>, deadline: Instant, phase: &str) -> Result<T, String> {
    let remaining = deadline.saturating_duration_since(Instant::now());
    if remaining.is_zero() {
        return Err(format!("Ruby host {phase} timeout"));
    }
    receiver
        .recv_timeout(remaining)
        .map_err(|error| match error {
            RecvTimeoutError::Timeout => format!("Ruby host {phase} timeout"),
            RecvTimeoutError::Disconnected => {
                format!("Ruby host crashed or closed its {phase} channel")
            }
        })
}

struct FrameWriter(Vec<u8>);

impl Write for FrameWriter {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        if bytes.len() > MAX_FRAME - self.0.len() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "protocol frame exceeds 4 MiB",
            ));
        }
        self.0.extend_from_slice(bytes);
        Ok(bytes.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

fn read_frame(reader: &mut impl BufRead) -> Result<Vec<u8>, String> {
    let mut frame = Vec::new();
    loop {
        let chunk = reader
            .fill_buf()
            .map_err(|error| format!("Ruby host read failed: {error}"))?;
        if chunk.is_empty() {
            return Err("Ruby host crashed or closed its output before a complete frame".into());
        }
        let newline = chunk.iter().position(|byte| *byte == b'\n');
        let length = newline.unwrap_or(chunk.len());
        if length > MAX_FRAME - frame.len() {
            return Err("Ruby protocol frame exceeds 4 MiB".into());
        }
        frame.extend_from_slice(&chunk[..length]);
        reader.consume(length + usize::from(newline.is_some()));
        if newline.is_some() {
            return Ok(frame);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU64, Ordering};

    static NEXT_FIXTURE: AtomicU64 = AtomicU64::new(1);

    struct Fixture(PathBuf);

    impl Fixture {
        fn new() -> Self {
            let path = env::temp_dir().join(format!(
                "vamp-ruby-{}-{}",
                std::process::id(),
                NEXT_FIXTURE.fetch_add(1, Ordering::Relaxed)
            ));
            std::fs::create_dir(&path).unwrap();
            Self(path)
        }

        fn file(&self, name: &str, text: &str) -> PathBuf {
            let path = self.0.join(name);
            std::fs::write(&path, text).unwrap();
            path
        }

        fn plugin(&self, text: &str) -> (RubyPlugin, Feature) {
            RubyPlugin::start(&self.file("plugin.rb", text)).unwrap_or_else(|error| {
                panic!("Actual Ruby runtime required (VAMP_RUBY or ruby): {error}")
            })
        }
    }

    impl Drop for Fixture {
        fn drop(&mut self) {
            let _ = std::fs::remove_dir_all(&self.0);
        }
    }

    fn context() -> Context {
        Context {
            text: "α\nhello".into(),
            row: 2,
            column: 3,
            mode: "NORMAL".into(),
            path: Some("buffer.txt".into()),
        }
    }

    #[test]
    fn actual_ruby_dsl_context_queued_effects_and_output_isolation() {
        let fixture = Fixture::new();
        let (mut plugin, feature) = fixture.plugin(r#"
puts "startup noise"
STDOUT.write "direct startup noise\n"
warn "warning"
Vamp.plugin "greeting", requires: ["text"]
Vamp.command "greet" do |editor, args|
  puts "invocation noise"
  STDOUT.write "direct invocation noise\n"
  raise "bad context" unless editor.text == "α\nhello" && editor.row == 2 && editor.column == 3 && editor.mode == "NORMAL" && editor.path == "buffer.txt"
  editor.insert "Hello #{args}!"
  editor.replace editor.text
  editor.message "Greeting inserted"
end
"#);
        assert_eq!(feature.name, "greeting");
        assert_eq!(feature.required_capabilities, ["text"]);
        assert_eq!(feature.commands, ["greet"]);
        assert!(!feature.core);
        for _ in 0..2 {
            let effects = plugin.invoke("greet", "friend", &context()).unwrap();
            assert!(
                matches!(&effects[..], [Effect::Insert(text), Effect::Replace(replacement), Effect::Message(message)]
                if text == "Hello friend!" && replacement == "α\nhello" && message == "Greeting inserted")
            );
        }
        plugin.disable();
        plugin.disable();
        assert!(
            plugin
                .invoke("greet", "", &context())
                .unwrap_err()
                .contains("disabled")
        );
        assert!(
            plugin.host.child.is_none()
                && plugin.host.reader.is_none()
                && plugin.host.writer.is_none()
        );
    }

    #[test]
    fn config_resolves_relative_paths_and_rejects_invalid_types_and_core() {
        let fixture = Fixture::new();
        let config = fixture.file(
            "config.rb",
            r#"
puts "configuration noise"
Vamp.configure do |config|
  config.plugin "plugins/greeting.rb"
  config.bind "g", "greet"
end
"#,
        );
        let loaded = load_config(&config).unwrap();
        assert_eq!(loaded.plugins, [fixture.0.join("plugins/greeting.rb")]);
        let invalid = fixture.file("invalid.rb", "Vamp.configure { |config| config.plugin 7 }");
        assert!(
            load_config(&invalid)
                .unwrap_err()
                .contains("must be a String")
        );
        let invalid = fixture.file("core.rb", "Vamp.plugin 'evil', core: true");
        assert!(RubyPlugin::start(&invalid).is_err());
        let invalid = fixture.file("syntax.rb", "this is not valid ruby (");
        assert!(RubyPlugin::start(&invalid).is_err());
    }

    #[test]
    fn config_leader_defaults_custom_keys_validation_and_required_wire_field() {
        let fixture = Fixture::new();
        let defaults = Config::default();
        assert_eq!(defaults.leader, "space");
        assert!(defaults.plugins.is_empty() && defaults.bindings.is_empty());
        let path = fixture.file("config.rb", "Vamp.configure { |_| }");
        assert_eq!(load_config(&path).unwrap().leader, "space");
        for (source, key) in [
            (r#"Vamp.configure { |config| config.leader = "," }"#, ","),
            (r#"Vamp.configure { |config| config.leader = "\\" }"#, "\\"),
        ] {
            let path = fixture.file("config.rb", source);
            assert_eq!(load_config(&path).unwrap().leader, key);
        }
        for (value, error) in [
            ("7", "leader key must be a String"),
            ("nil", "leader key must be a String"),
            (r#""""#, "leader key must not be empty"),
        ] {
            let path = fixture.file(
                "config.rb",
                &format!("Vamp.configure {{ |config| config.leader = {value} }}"),
            );
            assert!(load_config(&path).unwrap_err().contains(error), "{value}");
        }
        for frame in [
            r#"{"type":"config","version":1,"plugins":[],"bindings":[]}"#,
            r#"{"type":"config","version":1,"plugins":[],"bindings":[],"leader":null}"#,
            r#"{"type":"config","version":1,"plugins":[],"bindings":[],"leader":7}"#,
        ] {
            assert!(serde_json::from_str::<Response>(frame).is_err(), "{frame}");
        }
    }

    #[test]
    fn crash_hang_callback_failure_and_fresh_host_recovery() {
        let fixture = Fixture::new();
        for callback in [
            "exit! 19",
            "sleep 10",
            "editor.insert 'must not escape'; raise 'callback failed'",
            "editor.insert 7",
        ] {
            let (mut plugin, _) = fixture.plugin(&format!(
                "Vamp.plugin 'failure'\nVamp.command('fail') {{ |editor, args| {callback} }}"
            ));
            let error = plugin.invoke("fail", "", &context()).unwrap_err();
            if callback == "sleep 10" {
                assert!(error.contains("timeout"), "{error}");
            }
            assert!(
                plugin.host.child.is_none()
                    && plugin.host.reader.is_none()
                    && plugin.host.writer.is_none()
            );
            assert!(
                plugin
                    .invoke("fail", "", &context())
                    .unwrap_err()
                    .contains("disabled")
            );
            let (mut recovered, _) = fixture.plugin(
                "Vamp.plugin 'recovered'\nVamp.command('ok') { |editor, args| editor.message 'alive' }");
            assert!(
                matches!(&recovered.invoke("ok", "", &context()).unwrap()[..],
                [Effect::Message(message)] if message == "alive")
            );
        }
    }

    #[test]
    fn startup_and_config_hangs_are_bounded() {
        let fixture = Fixture::new();
        let path = fixture.file("hang.rb", "sleep 10");
        assert!(RubyPlugin::start(&path).err().unwrap().contains("timeout"));
        assert!(load_config(&path).unwrap_err().contains("timeout"));
    }

    #[test]
    fn invalid_protocol_version_shape_effect_and_request_id_shutdown() {
        let fixture = Fixture::new();
        for mutation in [
            "frame['version'] = 2",
            "frame['id'] += 1",
            "frame['effects'] = [{'type' => 'delete', 'text' => 'x'}]",
            "frame['effects'] = [{'type' => 'insert', 'text' => 7}]",
            "frame['unexpected'] = true",
        ] {
            let source = format!(
                r#"
Vamp.plugin 'protocol'
Vamp.command('ok') {{ |editor, args| editor.insert 'x' }}
module JSON
  class << self
    alias_method :original_generate, :generate
    def generate(frame)
      if frame['type'] == 'result'
        {mutation}
      end
      original_generate(frame)
    end
  end
end
"#
            );
            let (mut plugin, _) = fixture.plugin(&source);
            assert!(plugin.invoke("ok", "", &context()).is_err(), "{mutation}");
            assert!(plugin.host.child.is_none());
        }
    }

    #[test]
    fn bounded_frames_and_stalled_writes_shutdown() {
        let fixture = Fixture::new();
        let (mut plugin, _) = fixture.plugin(
            "Vamp.plugin 'large'\nVamp.command('large') { |editor, args| editor.insert 'x' * (4 * 1024 * 1024 + 1) }");
        assert!(plugin.invoke("large", "", &context()).is_err());
        assert!(plugin.host.child.is_none());

        // Stop Ruby reading the request pipe after it sends the manifest.
        let (mut plugin, _) = fixture.plugin(
            r#"
Vamp.plugin 'blocked'
Vamp.command('ok') { |editor, args| editor.message 'ok' }
class IO
  def gets(*args)
    sleep 10
  end
end
"#,
        );
        let mut big_context = context();
        big_context.text = "x".repeat(2 * 1024 * 1024);
        assert!(
            plugin
                .invoke("ok", "", &big_context)
                .unwrap_err()
                .contains("timeout")
        );
        assert!(plugin.host.child.is_none());

        let (mut plugin, _) = fixture.plugin(
            "Vamp.plugin 'request'\nVamp.command('ok') { |editor, args| editor.message 'ok' }",
        );
        big_context.text = "x".repeat(MAX_FRAME + 1);
        assert!(
            plugin
                .invoke("ok", "", &big_context)
                .unwrap_err()
                .contains("4 MiB")
        );
        assert!(plugin.host.child.is_none());

        for payload in ["'not json'", "'x' * (4 * 1024 * 1024 + 1)"] {
            let source = format!(
                r#"
Vamp.plugin 'bad_frame'
Vamp.command('ok') {{ |editor, args| editor.message 'ok' }}
class IO
  alias_method :original_write, :write
  def write(payload)
    if payload.include?('"type":"result"')
      original_write({payload})
    else
      original_write(payload)
    end
  end
end
"#
            );
            let (mut plugin, _) = fixture.plugin(&source);
            assert!(plugin.invoke("ok", "", &context()).is_err());
            assert!(plugin.host.child.is_none());
        }
    }

    #[cfg(unix)]
    #[test]
    fn inherited_pipes_do_not_prevent_drop_or_timeout_cleanup() {
        let fixture = Fixture::new();
        let (mut plugin, _) = fixture.plugin(
            r#"
Vamp.plugin 'forked'
Vamp.command('hang') do |editor, args|
  fork { sleep 10 }
  sleep 10
end
"#,
        );
        assert!(
            plugin
                .invoke("hang", "", &context())
                .unwrap_err()
                .contains("timeout")
        );
        assert!(plugin.host.reader.is_none() && plugin.host.writer.is_none());
        let (plugin, _) = fixture.plugin("Vamp.plugin 'drop'\nfork { sleep 10 }");
        drop(plugin);
    }
}