notescraft/crates/engine/tests/control.rs

//! The control channel: a live server on loopback whose requests are carried to
//! a pumped [`Session`] and dispatched through the registry (SPEC 10, D9).
//!
//! Every round trip uses a real `127.0.0.1` listener; the pump runs on the test
//! thread exactly where the UI thread would run it.

use std::io::{Read, Write};
use std::net::TcpStream;
use std::time::Instant;

use notescraft_engine::{
    CommandResult, CommandSpec, ControlClient, ControlHandle, ControlOptions, ControlServer, Error, ParamKind, ParamSpec, Registry, Session, register_core,
};
use serde_json::{Value, json};

mod common;
use common::{TempVault, temp_vault};

/// Idle timeouts for tests where connections must survive; the short one is for
/// the timeout test itself.
const ALIVE: std::time::Duration = std::time::Duration::from_secs(30);
const IDLE: std::time::Duration = std::time::Duration::from_millis(200);

/// The `test.echo` body: params back out unchanged.
#[expect(clippy::unnecessary_wraps, reason = "the registry signature demands CommandResult")]
fn echo(_session: &mut Session, params: &Value) -> CommandResult {
    Ok(params.clone())
}

/// A vault, its session, its registry (core plus one stub command) and the live
/// control server over it.
fn spun_up(name: &str, idle: std::time::Duration) -> (TempVault, Session, Registry<Session>, ControlHandle) {
    let temp = temp_vault(name);
    let session = Session::open(temp.path()).expect("open the vault");
    let mut registry = Registry::new();
    register_core(&mut registry).expect("register the core commands");
    registry
        .register(CommandSpec::new("test.echo", "Echo", "echoes its params", |_: &Session| true, echo).with_param(ParamSpec::optional(
            "question",
            ParamKind::Integer,
            "number echoed back",
        )))
        .expect("register test.echo");
    let options = ControlOptions { idle_timeout: idle, ..ControlOptions::default() };
    let handle = ControlServer::start_with(session.vault(), options).expect("start the control server");
    (temp, session, registry, handle)
}

/// Pumps until `worker` finishes, then joins it: the test thread acts as the UI
/// thread. The per-frame budget makes this the same code path the app takes.
fn drive<T>(handle: &mut ControlHandle, session: &mut Session, registry: &Registry<Session>, worker: std::thread::JoinHandle<T>) -> T {
    let started = std::time::Instant::now();
    while started.elapsed() < std::time::Duration::from_secs(60) {
        handle.pump(session, registry);
        if worker.is_finished() {
            return worker.join().expect("worker thread did not panic");
        }
        std::thread::sleep(std::time::Duration::from_millis(1));
    }
    panic!("the worker never finished within 60s");
}

/// One reply line (up to the newline) from a raw socket, `None` at EOF.
fn read_reply(stream: &mut TcpStream) -> Option<Vec<u8>> {
    let mut gathered = Vec::new();
    let mut chunk = [0u8; 512];
    loop {
        if let Some(pos) = gathered.iter().position(|&b| b == b'\n') {
            gathered.truncate(pos);
            return Some(gathered);
        }
        let read = stream.read(&mut chunk).ok()?;
        if read == 0 {
            return gathered.is_empty().then_some(gathered);
        }
        gathered.extend_from_slice(&chunk[..read]);
    }
}

/// `ping`, `info`, `commands` and `run` over a real client.
#[test]
fn a_client_roundtrips_ping_info_commands_and_run() {
    let (keep, mut session, registry, mut handle) = spun_up("control-roundtrip", ALIVE);
    let root = keep.path().to_owned();
    let worker = std::thread::spawn(move || {
        let mut client = ControlClient::connect(&root).expect("connect");
        let ping = client.ping().expect("ping");
        let info = client.info().expect("info");
        let commands = client.commands().expect("commands");
        let run = client.run("test.echo", &json!({ "question": 42 })).expect("run echo");
        (ping, info, commands, run)
    });
    let (ping, info, commands, run) = drive(&mut handle, &mut session, &registry, worker);
    assert_eq!(ping.get("protocol"), Some(&json!(1u64)));
    assert!(info.is_object(), "info is the vault.info result: {info}");
    assert!(commands.as_array().is_some_and(|cmds| cmds.iter().any(|cmd| cmd.get("id") == Some(&json!("test.echo")))), "{commands}");
    assert_eq!(run, json!({ "question": 42 }));
}

/// A command failure comes back as `{ ok: false, error: { code, message } }`
/// and the client turns it into an engine error again.
#[test]
fn a_failed_command_is_an_error_with_a_code() {
    let (keep, mut session, registry, mut handle) = spun_up("control-error", ALIVE);
    let root = keep.path().to_owned();
    let worker = std::thread::spawn(move || {
        let mut client = ControlClient::connect(&root).expect("connect");
        client.run("no.such.command", &Value::Null)
    });
    match drive(&mut handle, &mut session, &registry, worker) {
        Ok(value) => panic!("an unknown command cannot succeed: {value}"),
        Err(Error::Command { id, source }) => {
            assert_eq!(id, "control");
            assert_eq!(source.code, "unknown_command");
        }
        Err(other) => panic!("expected a command error, got {other}"),
    }
}

/// A wrong or missing token is refused with a fixed delay, and the connection
/// closes after three failures.
#[test]
fn wrong_tokens_are_refused_and_the_connection_closes_after_three() {
    let (_keep, _session, _registry, handle) = spun_up("control-bad-token", ALIVE);
    let mut stream = TcpStream::connect(("127.0.0.1", handle.port())).expect("connect");
    for _ in 0..3 {
        stream.write_all(br#"{"id":7,"token":"wrong","method":"ping"}"#).ok();
        stream.write_all(b"\n").ok();
        let reply = read_reply(&mut stream).expect("an unauthorized reply");
        let reply: Value = serde_json::from_slice(&reply).expect("parse the reply");
        assert_eq!(reply["id"], json!(7u64));
        assert_eq!(reply["error"]["code"], json!("unauthorized"), "{reply}");
    }
    let mut leftover = [0u8; 16];
    let read = stream.read(&mut leftover).ok().unwrap_or_default();
    assert_eq!(read, 0, "the connection should be closed after three failures");
}

/// More than 1 MiB on one line is refused, then the socket is closed.
#[test]
fn an_oversized_line_is_refused_and_the_connection_closed() {
    let (_keep, _session, _registry, handle) = spun_up("control-oversize", ALIVE);
    let mut stream = TcpStream::connect(("127.0.0.1", handle.port())).expect("connect");
    stream.write_all(&vec![b'"'; 1024 * 1024]).ok();
    stream.write_all(b"x\"\n").ok();
    let reply = read_reply(&mut stream).expect("an oversized reply");
    let reply: Value = serde_json::from_slice(&reply).expect("parse the reply");
    assert_eq!(reply["error"]["code"], json!("line_too_large"), "{reply}");
    let mut leftover = [0u8; 16];
    let read = stream.read(&mut leftover).ok().unwrap_or_default();
    assert_eq!(read, 0, "the connection should be closed");
}

/// Nine concurrent connections: eight are served, the ninth is refused.
#[test]
fn the_ninth_concurrent_connection_is_refused() {
    let (_keep, _session, _registry, handle) = spun_up("control-ninth", ALIVE);
    let mut served = Vec::new();
    for _ in 0..8 {
        served.push(TcpStream::connect(("127.0.0.1", handle.port())).expect("connect"));
    }
    let mut ninth = TcpStream::connect(("127.0.0.1", handle.port())).expect("connect");
    let mut leftover = [0u8; 16];
    let read = ninth.read(&mut leftover).ok().unwrap_or_default();
    assert_eq!(read, 0, "the ninth connection should be closed at once");
    drop(served);
}

/// A connection that sends nothing is closed after the idle timeout, and the
/// server keeps working for the next client.
#[test]
fn an_idle_connection_is_closed_after_the_timeout() {
    let (keep, mut session, registry, mut handle) = spun_up("control-idle", IDLE);
    let mut stream = TcpStream::connect(("127.0.0.1", handle.port())).expect("connect");
    let started = Instant::now();
    let mut leftover = [0u8; 16];
    let read = stream.read(&mut leftover).ok().unwrap_or_default();
    assert!(started.elapsed() >= std::time::Duration::from_millis(150), "closed too early");
    assert!(started.elapsed() < std::time::Duration::from_secs(5), "not closed soon after the timeout");
    assert_eq!(read, 0, "expected EOF after the idle timeout");
    let root = keep.path().to_owned();
    let worker = std::thread::spawn(move || {
        let mut client = ControlClient::connect(&root).expect("connect");
        client.ping().expect("ping")
    });
    let ping = drive(&mut handle, &mut session, &registry, worker);
    assert_eq!(ping.get("protocol"), Some(&json!(1u64)));
}

/// A stale `control.json` (dead pid) is refused by the client and overwritten
/// by `start`, and dropping the server removes the file.
#[test]
fn a_stale_control_file_is_overwritten_and_removed_on_drop() {
    let keep = temp_vault("control-stale");
    let file = keep.path().join(".notescraft");
    std::fs::create_dir_all(&file).expect("create the private folder");
    let file = file.join("control.json");
    std::fs::write(&file, br#"{"port":1,"token":"stale","pid":999999999,"started":0}"#).expect("seed the stale file");
    match ControlClient::connect(keep.path()) {
        Ok(_) => panic!("a control file with a dead pid must refuse connections"),
        Err(Error::Control(_)) => {}
        Err(other) => panic!("expected a control error, got {other}"),
    }
    let session = Session::open(keep.path()).expect("open the vault");
    let handle = ControlServer::start_with(session.vault(), ControlOptions { idle_timeout: IDLE, ..ControlOptions::default() }).expect("start");
    let written: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&file).expect("read the control file")).expect("parse the control file");
    assert_eq!(written.get("port"), Some(&json!(handle.port())), "{written}");
    assert_eq!(written.get("pid"), Some(&json!(std::process::id())), "{written}");
    drop(handle);
    assert!(!file.exists(), "dropping the server removes control.json");
}