sausage/cli/tests/plugins/probe/src/main.rs

//! Raw JSON-RPC methods used by the host's protocol tests.
use serde_json::{Value, json};
use std::io::{self, BufRead, Read, Write};
type Lines = io::Lines<io::BufReader<io::Stdin>>;

fn send(value: Value) {
    let mut output = io::stdout().lock();
    writeln!(output, "\n{value}").unwrap();
    output.flush().unwrap();
}
fn receive(lines: &mut Lines) -> Value {
    serde_json::from_str(&lines.next().unwrap().unwrap()).unwrap()
}
fn echo(params: &Value) -> Value {
    params.clone()
}
fn chat(params: &Value) -> Value {
    send(json!({"jsonrpc":"2.0", "method":"core.notify", "params":{
        "name":"text", "params":{"text":"interim"}
    }}));
    json!({"text":params["text"]})
}
fn discover(params: &Value) -> Value {
    if let Ok(bytes) = std::fs::read("/state/probe/discover.json") {
        serde_json::from_slice(&bytes).unwrap()
    } else {
        params.clone()
    }
}
fn environment() -> Value {
    json!({
        "pid": std::env::var("SAUSAGE_PID").unwrap(),
        "plugin": std::env::var("SAUSAGE_PLUGIN").unwrap(),
        "parent": std::env::var("SAUSAGE_PARENT").unwrap(),
    })
}
fn multicall(lines: &mut Lines) -> Value {
    send(json!({
        "jsonrpc": "2.0",
        "id": 1,
        "method": "probe.echo",
        "params": {"n": 1}
    }));
    send(json!({
        "jsonrpc": "2.0",
        "id": "1",
        "method": "probe.echo",
        "params": {"n": 2}
    }));
    for _ in 0..2 {
        let response = receive(lines);
        if response["id"] == 1 {
            assert_eq!(response["result"]["n"], 1);
        } else {
            assert_eq!(response["id"], "1");
            assert_eq!(response["result"]["n"], 2);
        }
    }
    json!({"count":2})
}
fn notification(params: &Value) -> Value {
    send(json!({
        "jsonrpc": "2.0",
        "method": "probe.write",
        "params": params
    }));
    let before = std::time::Instant::now();
    loop {
        let contents = std::fs::read_to_string(params["path"].as_str().unwrap()).ok();
        if contents.as_deref() == params["text"].as_str() {
            break;
        }
        assert!(before.elapsed() < std::time::Duration::from_secs(10));
    }
    json!({})
}
fn unknown_id() {
    send(json!({
        "jsonrpc": "2.0",
        "id": "late",
        "result": 7
    }));
}
fn wait_eof(lines: &mut Lines) {
    std::fs::write("/state/waiting", "yes").unwrap();
    send(json!({"jsonrpc":"2.0","method":"core.notify","params":{
        "name":"waiting-eof", "params":{"pid":std::env::var("SAUSAGE_PID").unwrap()}
    }}));
    for _ in lines {}
    std::fs::write("/state/stopped", "yes").unwrap();
}
fn busy() -> Value {
    loop {
        std::hint::spin_loop();
    }
}
fn exit_without_answer() {
    eprintln!("last line");
}
fn junk() -> Value {
    io::stdout().write_all(&[255, 10]).unwrap();
    println!("junk\n{{}}\n[]");
    eprintln!("stderr line");
    json!({})
}
fn invalid_utf8() -> Value {
    io::stdout().write_all(&[255, 10]).unwrap();
    json!({})
}
fn blank() -> Value {
    println!("\n \t\n");
    json!({})
}
fn burst(request: &Value) {
    let count = request["params"]["count"].as_u64().unwrap();
    let mut output = String::new();
    for index in 0..count {
        let notification = json!({
            "jsonrpc": "2.0",
            "method": "core.notify",
            "params": {"name": "burst", "params": {"index": index}}
        });
        output.push_str(&notification.to_string());
        output.push('\n');
    }
    let response = json!({"jsonrpc":"2.0", "id":request["id"], "result":count});
    output.push_str(&response.to_string());
    output.push('\n');
    let mut stdout = io::stdout().lock();
    stdout.write_all(output.as_bytes()).unwrap();
    stdout.flush().unwrap();
}
fn notify() -> Value {
    send(json!({
        "jsonrpc": "2.0",
        "method": "core.notify",
        "params": {
            "name": "progress",
            "params": {"n": 1}
        }
    }));
    json!({})
}
fn serve(request: &Value, lines: &mut Lines) {
    send(json!({
        "jsonrpc": "2.0",
        "id": request["id"],
        "result": std::env::var("SAUSAGE_PID").unwrap()
    }));
    for _ in lines {}
}
fn network(params: &Value) -> Value {
    let mut socket = std::net::TcpStream::connect(params["tcp"].as_str().unwrap()).unwrap();
    let mut text = String::new();
    socket.read_to_string(&mut text).unwrap();
    let udp = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
    udp.send_to(b"udp", params["udp"].as_str().unwrap())
        .unwrap();
    json!(text)
}
fn read_file(params: &Value) -> Value {
    json!(std::fs::read_to_string(params["path"].as_str().unwrap()).unwrap())
}
fn write_file(params: &Value) -> Value {
    std::fs::write(
        params["path"].as_str().unwrap(),
        params["text"].as_str().unwrap(),
    )
    .unwrap();
    json!({})
}
fn outgoing_call(params: &Value, lines: &mut Lines) -> Value {
    send(json!({
        "jsonrpc": "2.0",
        "id": {"nested": 7},
        "method": params["method"],
        "params": params["params"]
    }));
    loop {
        let response = receive(lines);
        if response["id"] == json!({"nested":7}) {
            return response["result"].clone();
        }
    }
}
fn lock(params: &Value, lines: &mut Lines) -> Value {
    outgoing_call(&json!({"method":"core.lock", "params":params}), lines)
}
fn filesystem(params: &Value) -> Value {
    use std::fs;
    let pid = std::env::var("SAUSAGE_PID").unwrap();
    match params["topic"].as_str().unwrap() {
        "root" => {
            let preopens = wasip2::filesystem::preopens::get_directories();
            assert_eq!(preopens.len(), 1);
            assert_eq!(preopens[0].1, "/");
            let names: Vec<_> = fs::read_dir("/")
                .unwrap()
                .map(|e| e.unwrap().file_name().into_string().unwrap())
                .collect();
            for name in ["proc", "self", "run"] {
                assert!(names.iter().any(|n| n == name));
            }
            fs::write("/state/root", "world").unwrap();
            fs::write("/run/shared", "shared").unwrap();
            assert!(fs::rename("/run/shared", "/state/moved").is_err());
            assert!(fs::rename("/state/root", "/self/moved").is_err());
            assert_eq!(fs::read_to_string("/run/../state/root").unwrap(), "world");
        }
        "reserved" => {
            for name in ["info.json", "log", "fs"] {
                assert!(fs::write(format!("/self/{name}"), "bad").is_err());
                assert!(fs::create_dir(format!("/self/{name}")).is_err());
                fs::write("/self/temp", "ok").unwrap();
                assert!(fs::rename("/self/temp", format!("/self/{name}")).is_err());
                fs::remove_file("/self/temp").unwrap();
            }
        }
        "self" => {
            fs::write("/self/endpoint.tmp", "endpoint").unwrap();
            fs::rename("/self/endpoint.tmp", "/self/endpoint").unwrap();
            assert_eq!(
                fs::read_to_string(format!("/proc/{pid}/endpoint")).unwrap(),
                "endpoint"
            );
            assert!(fs::write(format!("/proc/{pid}/endpoint"), "bad").is_err());
            assert!(fs::remove_file(format!("/proc/{pid}/endpoint")).is_err());
            let info: Value =
                serde_json::from_slice(&fs::read("/self/info.json").unwrap()).unwrap();
            assert_eq!(info["plugin"], "probe");
            assert_eq!(info["method"], "filesystem");
            return json!(pid);
        }
        "isolated" => {
            assert!(!fs::exists("/self/endpoint").unwrap());
        }
        "snapshot" => {
            eprintln!("before snapshot");
            let mut file = fs::File::open("/self/log").unwrap();
            eprintln!("after snapshot");
            let mut contents = String::new();
            file.read_to_string(&mut contents).unwrap();
            assert!(contents.contains("start filesystem"));
            assert!(!contents.contains("after snapshot"));
            assert_eq!(file.read(&mut [0; 1]).unwrap(), 0);
        }
        "core" => {
            assert!(
                fs::read_to_string("/proc/core/PROTOCOL.md")
                    .unwrap()
                    .contains("## 4. The filesystem")
            );
            for name in ["graph.json", "locks.json"] {
                let _: Value =
                    serde_json::from_slice(&fs::read(format!("/proc/core/{name}")).unwrap())
                        .unwrap();
            }
            assert!(
                fs::read_to_string("/proc/core/graph.dot")
                    .unwrap()
                    .contains("digraph")
            );
            assert!(fs::read_dir("/proc/core/plugins").is_ok());
            assert!(fs::write("/proc/core/graph.json", "bad").is_err());
        }
        _ => panic!("unknown filesystem topic"),
    }
    json!({})
}
fn write_loop(request: &Value, lines: &mut Lines) {
    assert_eq!(
        lock(&json!({"path":"/state/loop"}), lines),
        json!({"held":true})
    );
    std::fs::write("/state/loop", "old").unwrap();
    send(
        json!({"jsonrpc":"2.0", "id":request["id"], "result":std::env::var("SAUSAGE_PID").unwrap()}),
    );
    loop {
        std::fs::write("/state/loop", "old").unwrap();
    }
}
fn locked_write(lines: &mut Lines) -> Value {
    assert_eq!(
        lock(&json!({"path":"/state/loop"}), lines),
        json!({"held":true})
    );
    std::fs::write("/state/loop", "new").unwrap();
    outgoing_call(&json!({"method":"probe.echo","params":{}}), lines);
    assert_eq!(std::fs::read_to_string("/state/loop").unwrap(), "new");
    json!({})
}
fn unlock_answer(request: &Value, lines: &mut Lines) {
    lock(&json!({"path":"ordered-unlock"}), lines);
    let unlock = json!({"jsonrpc":"2.0","method":"core.unlock","params":{"path":"ordered-unlock"}});
    let answer =
        json!({"jsonrpc":"2.0","id":request["id"],"result":std::env::var("SAUSAGE_PID").unwrap()});
    let mut stdout = io::stdout().lock();
    writeln!(stdout, "{unlock}\n{answer}").unwrap();
    stdout.flush().unwrap();
    drop(stdout);
    for _ in lines {}
}
fn notification_child(lines: &mut Lines) -> Value {
    send(json!({"jsonrpc":"2.0","method":"probe.orphan_notify","params":{}}));
    let before = std::time::Instant::now();
    let pid = loop {
        if let Ok(pid) = std::fs::read_to_string("/state/orphan") {
            break pid;
        }
        assert!(before.elapsed() < std::time::Duration::from_secs(10));
    };
    send(json!({"jsonrpc":"2.0","id":7,"method":"core.stop","params":{"pid":pid}}));
    let mut count = 0;
    loop {
        let response = receive(lines);
        if response["method"] == "core.notify" {
            count += 1;
        }
        if response["id"] == 7 {
            break;
        }
    }
    json!(count)
}
fn batch(lines: &mut Lines) -> Value {
    send(
        json!([{"jsonrpc":"2.0","method":"probe.write","params":{"path":"/state/batch","text":"bad"}}]),
    );
    outgoing_call(
        &json!({"method":"probe.echo","params":{"barrier":true}}),
        lines,
    );
    json!(!std::path::Path::new("/state/batch").exists())
}
fn late_answer(request: &Value, lines: &mut Lines) {
    send(json!({"jsonrpc":"2.0","id":request["id"],"result":42}));
    send(json!({"jsonrpc":"2.0","id":request["id"],"result":99}));
    for _ in lines {}
}
fn nested_environment(lines: &mut Lines) -> Value {
    let child = outgoing_call(&json!({"method":"probe.env","params":{}}), lines);
    json!({"caller":std::env::var("SAUSAGE_PID").unwrap(),"child":child})
}
fn forgotten(request: &Value, lines: &mut Lines) {
    std::fs::write("/state/forgotten", environment().to_string()).unwrap();
    serve(request, lines);
}
fn forget(lines: &mut Lines) -> Value {
    send(
        json!({"jsonrpc":"2.0","method":"core.call","params":{"method":"probe.forgotten","params":{},"detach":true}}),
    );
    let before = std::time::Instant::now();
    while !std::path::Path::new("/state/forgotten").exists() {
        assert!(before.elapsed() < std::time::Duration::from_secs(10));
    }
    send(json!({"jsonrpc":"2.0","id":"barrier","method":"probe.echo","params":{}}));
    let response = receive(lines);
    assert_eq!(
        response["id"], "barrier",
        "a forgotten call must not send a result"
    );
    environment()
}
fn tree(request: &Value, lines: &mut Lines) {
    let depth = request["params"]["depth"].as_u64().unwrap();
    if depth > 0 {
        outgoing_call(
            &json!({"method":"probe.tree","params":{"depth":depth-1}}),
            lines,
        );
    }
    send(
        json!({"jsonrpc":"2.0","id":request["id"],"result":std::env::var("SAUSAGE_PID").unwrap()}),
    );
    busy();
}
fn dispatch(request: &Value, lines: &mut Lines) -> Option<Value> {
    let params = &request["params"];
    let result = match request["method"].as_str().unwrap() {
        "batch" => batch(lines),
        "trap" => std::arch::wasm32::unreachable(),
        "nested_environment" => nested_environment(lines),
        "forget" => forget(lines),
        "forgotten" => {
            forgotten(request, lines);
            return None;
        }
        "tree" => {
            tree(request, lines);
            return None;
        }
        "late_answer" => {
            late_answer(request, lines);
            return None;
        }
        "invalid_stderr" => {
            io::stderr().write_all(&[255, 10]).unwrap();
            eprintln!("after invalid UTF-8");
            json!(std::env::var("SAUSAGE_PID").unwrap())
        }
        "unlock_answer" => {
            unlock_answer(request, lines);
            return None;
        }
        "null_error" => {
            send(json!({"jsonrpc":"2.0","id":request["id"],"error":null,"result":42}));
            return None;
        }
        "notification_child" => notification_child(lines),
        "orphan_notify" => {
            std::fs::write("/state/orphan", std::env::var("SAUSAGE_PID").unwrap()).unwrap();
            notify()
        }
        "self_serve" => {
            std::fs::write("/self/endpoint", "endpoint").unwrap();
            serve(request, lines);
            return None;
        }
        "filesystem" => filesystem(params),
        "write_loop" => {
            write_loop(request, lines);
            return None;
        }
        "locked_write" => locked_write(lines),
        "echo" => echo(params),
        "chat" => chat(params),
        "rpc.discover" => discover(params),
        "env" => environment(),
        "multicall" => multicall(lines),
        "notification" => notification(params),
        "unknownid" => {
            unknown_id();
            return None;
        }
        "request_id" => request["id"].clone(),
        "wait_eof" => {
            wait_eof(lines);
            return None;
        }
        "busy" => busy(),
        "exit" => {
            exit_without_answer();
            return None;
        }
        "junk" => junk(),
        "blank" => blank(),
        "invalid_utf8" => invalid_utf8(),
        "notify" => notify(),
        "burst" => {
            burst(request);
            return None;
        }
        "restart" => {
            assert_eq!(
                lock(&json!({"path":"/state/restart-stop"}), lines),
                json!({"held":true})
            );
            assert_eq!(
                outgoing_call(&json!({"method":"core.restart", "params":{}}), lines),
                json!({})
            );
            std::fs::write("/state/restart-answer", "received").unwrap();
            serve(request, lines);
            std::fs::write("/state/restart-stop", "persisted after EOF").unwrap();
            eprintln!("restart stop hook finished");
            return None;
        }
        "serve" => {
            serve(request, lines);
            return None;
        }
        "network" => network(params),
        "mutate" => mutate(params),
        "read" => read_file(params),
        "write" => write_file(params),
        "raw_call" => {
            send(params.clone());
            receive(lines)
        }
        "call" => outgoing_call(params, lines),
        "lock" => lock(params, lines),
        _ => {
            send(json!({
                "jsonrpc": "2.0",
                "id": request["id"],
                "error": {
                    "code": -32601,
                    "message": "method not found"
                }
            }));
            return None;
        }
    };
    Some(result)
}
fn main() {
    if let Ok(mode) = std::fs::read_to_string("/state/early") {
        if mode == "flood" {
            for index in 0..5000 {
                eprintln!("flood {index} {}", "x".repeat(128));
            }
        } else {
            eprintln!("last early log");
            if mode == "answer" {
                send(
                    json!({"jsonrpc":"2.0","id":std::env::var("SAUSAGE_PID").unwrap(),"result":"early"}),
                );
            }
            std::process::exit(23);
        }
    }
    let mut lines = io::BufReader::new(io::stdin()).lines();
    let request = receive(&mut lines);
    if let Some(result) = dispatch(&request, &mut lines)
        && let Some(id) = request.get("id")
    {
        send(json!({
            "jsonrpc": "2.0",
            "id": id,
            "result": result
        }));
    }
}

/// Separate mutation calls let the host await a scan before each subsequent change.
fn mutate(p: &Value) -> Value {
    use wasip2::filesystem::{
        preopens,
        types::{DescriptorFlags as D, OpenFlags as O, PathFlags as P},
    };
    let path = p["path"].as_str().unwrap();
    match p["op"].as_str().unwrap() {
        "mkdir" => std::fs::create_dir_all(path).unwrap(),
        "rename" => std::fs::rename(path, p["to"].as_str().unwrap()).unwrap(),
        "delete" => std::fs::remove_file(path).unwrap(),
        "rmdir" => std::fs::remove_dir(path).unwrap(),
        "atomic" => {
            let temp = format!("{path}.tmp");
            std::fs::write(&temp, "complete").unwrap();
            std::fs::rename(temp, path).unwrap();
        }
        "burst" => {
            for i in 0..5 {
                std::fs::write(format!("{path}/{i}"), "complete").unwrap();
            }
        }
        op => {
            let root = preopens::get_directories().remove(0).0;
            let file = root
                .open_at(
                    P::empty(),
                    path.trim_start_matches('/'),
                    O::empty(),
                    D::READ | D::WRITE,
                )
                .unwrap();
            match op {
                "descriptor" => {
                    file.write(b"descriptor", 0).unwrap();
                }
                "stream" => {
                    file.write_via_stream(0)
                        .unwrap()
                        .blocking_write_and_flush(b"stream")
                        .unwrap();
                }
                "append" => {
                    file.append_via_stream()
                        .unwrap()
                        .blocking_write_and_flush(b"append")
                        .unwrap();
                }
                "size" => file.set_size(0).unwrap(),
                "truncate" => {
                    root.open_at(
                        P::empty(),
                        path.trim_start_matches('/'),
                        O::TRUNCATE,
                        D::WRITE,
                    )
                    .unwrap();
                }
                _ => panic!("unknown mutation"),
            }
        }
    }
    json!({})
}