notescraft/crates/engine/tests/control_client.rs

//! The control channel under load and hostile input: a 1000-request flood
//! pumped in budgeted slices, and a fuzz pass that must survive arbitrary
//! bytes, truncated JSON and bad params without panicking the server.

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

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

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

/// Idle timeout for the servers under test.
const ALIVE: std::time::Duration = std::time::Duration::from_secs(30);

/// 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 — the spun-up shape the server tests use.
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)).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.
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");
}

/// 1000 pipelined pings are pumped in slices of at most 16, and the client
/// reads exactly its 1000 replies: budget slices spread the load, nothing
/// starves.
#[test]
fn a_flood_of_1000_is_pumped_in_budgeted_slices() {
    let (keep, mut session, registry, mut handle) = spun_up("control-flood", ALIVE);
    let token = handle.token().to_owned();
    let root = keep.path().to_owned();
    let worker = std::thread::spawn(move || -> usize {
        let mut stream = TcpStream::connect(("127.0.0.1", descriptor_port(&root))).expect("connect");
        for request_id in 1..1001u64 {
            let line = format!(r#"{{"id":{request_id},"token":"{token}","method":"ping"}}"#);
            stream.write_all(line.as_bytes()).expect("write a flood request");
            stream.write_all(b"\n").expect("write the flood newline");
        }
        stream.set_read_timeout(Some(std::time::Duration::from_secs(60))).expect("set the flood read timeout");
        let mut buffer = Vec::new();
        stream.read_to_end(&mut buffer).expect("read the flood replies");
        let replies = String::from_utf8_lossy(&buffer);
        replies.matches(r#""ok":true"#).count()
    });
    let mut total = 0usize;
    let mut slices = 0usize;
    let started = std::time::Instant::now();
    while total < 1000 && started.elapsed() < std::time::Duration::from_secs(60) {
        let handled = handle.pump(&mut session, &registry);
        if handled > 0 {
            slices += 1;
            assert!(handled <= 16, "a pump slice ran {handled} requests, over the budget");
        }
        total += handled;
        std::thread::sleep(std::time::Duration::from_millis(1));
    }
    assert_eq!(total, 1000, "every flooded request was pumped");
    let ok_replies = drive(&mut handle, &mut session, &registry, worker);
    assert_eq!(ok_replies, 1000, "the client read one ok reply per flooded request");
    assert!(slices >= 62, "1000 requests at 16 per slice need at least 62 slices, ran {slices}");
}

/// Junk is answered, never fatal, and the failure budget behaves: three
/// unauthorised lines close one connection, two tokened bad methods do not,
/// and a tokened ping on that same connection still succeeds.
#[test]
fn fuzz_input_never_panics_and_the_server_survives() {
    let (keep, mut session, registry, mut handle) = spun_up("control-fuzz", ALIVE);
    let token = handle.token().to_owned();
    let root = keep.path().to_owned();
    let reconnect_root = keep.path().to_owned();
    let worker = std::thread::spawn(move || -> String {
        // Connection one: tokenless junk until the failure budget closes it.
        let mut stream = TcpStream::connect(("127.0.0.1", descriptor_port(&root))).expect("connect");
        for line in [b"not json at all".as_slice(), b"[".repeat(64).as_slice(), br#"{"id":1,"#] {
            stream.write_all(line).expect("write junk");
            stream.write_all(b"\n").expect("write the junk newline");
        }
        stream.set_read_timeout(Some(std::time::Duration::from_secs(10))).expect("set the read timeout");
        let mut buffer = Vec::new();
        let _ = stream.read_to_end(&mut buffer);
        let refused = String::from_utf8_lossy(&buffer).into_owned();
        std::thread::sleep(std::time::Duration::from_secs(2));
        // Connection two: tokened junk methods twice (under the budget), then a
        // good ping on the same, still-open connection.
        let mut stream = TcpStream::connect(("127.0.0.1", descriptor_port(&root))).expect("connect");
        for _ in 0..2 {
            let line = format!(r#"{{"id":2,"token":"{token}","method":"dodgey","params":{{}}}}"#);
            stream.write_all(line.as_bytes()).expect("write a dodgey method");
            stream.write_all(b"\n").expect("write the dodgey newline");
        }
        let line = format!(r#"{{"id":3,"token":"{token}","method":"ping"}}"#);
        stream.write_all(line.as_bytes()).expect("write a good ping");
        stream.write_all(b"\n").expect("write the ping newline");
        let mut answered = Vec::new();
        let _ = stream.read_to_end(&mut answered);
        format!("{refused}{}", String::from_utf8_lossy(&answered))
    });
    let replies = drive(&mut handle, &mut session, &registry, worker);
    assert_eq!(replies.matches("unauthorized").count(), 3, "all three tokenless junk lines were refused: {replies}");
    assert_eq!(replies.matches("unknown_method").count(), 2, "the tokened junk methods were refused: {replies}");
    assert_eq!(replies.matches(r#""ok":true"#).count(), 1, "a tokened ping succeeds on a twice-refused connection: {replies}");
    let worker = std::thread::spawn(move || {
        let mut client = ControlClient::connect(&reconnect_root).expect("reconnect");
        let ping = client.ping().expect("ping after the junk");
        let run = client.run("test.echo", &Value::Null).expect("run after the junk");
        (ping, run)
    });
    let (ping, run) = drive(&mut handle, &mut session, &registry, worker);
    assert_eq!(ping.get("protocol"), Some(&json!(1u64)));
    assert_eq!(run, Value::Null);
}

/// The port from the control descriptor, for raw-socket drivers.
fn descriptor_port(root: &std::path::Path) -> u16 {
    let file = root.join(".notescraft").join("control.json");
    let text = std::fs::read_to_string(file).expect("read the descriptor");
    let port = serde_json::from_str::<Value>(&text).ok().and_then(|file| file.get("port").and_then(Value::as_u64));
    port.map_or(0, |port| u16::try_from(port).unwrap_or(0))
}