kcode-jsonrpc-stdio 0.3.0

#![cfg(unix)]

use kcode_jsonrpc_stdio::{
    Config, Error, MAX_INBOUND_CAPACITY, MAX_OUTBOUND_LINE_BYTES, MAX_STDERR_CAPTURE_BYTES,
    ProcessExit, StdioRpc,
};
use kcode_jsonrpc_wire::{Message, PeerId};
use serde_json::{Value, json};
use tokio::{
    process::Command,
    time::{Duration, sleep, timeout},
};

fn shell(script: &str) -> Command {
    let mut command = Command::new("sh");
    command.arg("-c").arg(script);
    command
}

fn spawn(script: &str, capacity: usize) -> StdioRpc {
    StdioRpc::spawn(
        shell(script),
        Config {
            inbound_capacity: capacity,
        },
    )
    .unwrap()
}

async fn finish(rpc: StdioRpc) -> ProcessExit {
    timeout(Duration::from_secs(2), rpc.shutdown())
        .await
        .unwrap()
        .unwrap()
}

async fn terminal(script: &str, expected: impl FnOnce(&Error) -> bool) {
    let mut rpc = spawn(script, 1);
    let error = timeout(Duration::from_secs(1), rpc.next())
        .await
        .unwrap()
        .expect_err("expected terminal read error");
    assert!(expected(&error));
    assert!(matches!(
        timeout(Duration::from_secs(1), rpc.next()).await.unwrap(),
        Err(Error::ReaderUnavailable)
    ));
    let _ = finish(rpc).await;
}

#[tokio::test]
async fn ordered_compact_headerless_lines_round_trip() {
    let script = r#"read a; read b; [ "$a" = '{"id":1,"method":"go","params":null}' ] || exit 9; [ "$b" = '{"method":"ping","params":[]}' ] || exit 10; printf '%s\n' '{"id":1,"result":true}' '{"method":"ready"}'"#;
    let mut rpc = spawn(script, 2);
    rpc.send(json!({"id":1,"method":"go","params":null}))
        .await
        .unwrap();
    rpc.send(json!({"method":"ping","params":[]}))
        .await
        .unwrap();
    assert_eq!(
        rpc.next().await.unwrap(),
        Message::Response {
            id: PeerId::Signed(1),
            outcome: Ok(Value::Bool(true))
        }
    );
    assert_eq!(
        rpc.next().await.unwrap(),
        Message::Notification {
            method: "ready".into(),
            params: Value::Null
        }
    );
    let _ = finish(rpc).await;
}

#[tokio::test]
async fn capacity_and_line_limits_are_enforced() {
    assert!(matches!(
        StdioRpc::spawn(
            shell("exit 0"),
            Config {
                inbound_capacity: 0
            }
        ),
        Err(Error::InvalidInboundCapacity(0))
    ));
    assert!(matches!(
        StdioRpc::spawn(
            shell("exit 0"),
            Config {
                inbound_capacity: MAX_INBOUND_CAPACITY + 1
            }
        ),
        Err(Error::InvalidInboundCapacity(_))
    ));
    let mut rpc = spawn("cat", 1);
    assert!(matches!(
        rpc.send(json!("x".repeat(MAX_OUTBOUND_LINE_BYTES))).await,
        Err(Error::OutboundLineTooLong)
    ));
    let _ = finish(rpc).await;
    let mut rpc = spawn("head -c 8388609 /dev/zero; printf '\\n'; exec sleep 5", 1);
    assert!(matches!(
        timeout(Duration::from_secs(3), rpc.next()).await.unwrap(),
        Err(Error::InboundLineTooLong)
    ));
    let _ = finish(rpc).await;
}

#[tokio::test]
async fn malformed_input_is_terminal_and_sticky() {
    terminal("printf '%s\\n' '{bad'; exec sleep 5", |error| {
        matches!(error, Error::Json(_))
    })
    .await;
    terminal(
        r#"printf '%s\n' '{"jsonrpc":"2.0","method":"x"}'; exec sleep 5"#,
        |error| matches!(error, Error::Wire(_)),
    )
    .await;
}

#[tokio::test]
async fn live_stdout_close_returns_promptly() {
    let mut rpc = spawn("exec 1>&-; exec sleep 5", 1);
    assert!(matches!(
        timeout(Duration::from_secs(1), rpc.next()).await.unwrap(),
        Err(Error::StdoutClosed)
    ));
    assert!(matches!(rpc.next().await, Err(Error::ReaderUnavailable)));
    let _ = finish(rpc).await;
}

#[tokio::test]
async fn partial_eof_distinguishes_utf8() {
    terminal(r#"printf '%s' '{"method":"x"}'"#, |error| {
        matches!(error, Error::IncompleteFrame)
    })
    .await;
    terminal(r"printf '\377'", |error| matches!(error, Error::Utf8(_))).await;
}

#[tokio::test]
async fn broken_pipe_permanently_closes_stdin() {
    let mut rpc = spawn("exec 0<&-; exec sleep 5", 1);
    let payload = json!("x".repeat(1024 * 1024));
    assert!(matches!(
        timeout(Duration::from_secs(2), rpc.send(payload))
            .await
            .unwrap(),
        Err(Error::Io(_))
    ));
    assert!(matches!(
        rpc.send(Value::Null).await,
        Err(Error::StdinClosed)
    ));
    let _ = finish(rpc).await;
}

#[tokio::test]
async fn stderr_is_bounded_sanitized_and_reaped() {
    let rpc = spawn(
        "printf 'bad\\001\\n' >&2; head -c 70000 /dev/zero >&2; exit 7",
        1,
    );
    sleep(Duration::from_millis(100)).await;
    let exit = finish(rpc).await;
    assert_eq!(exit.status.code(), Some(7));
    assert_eq!(exit.stderr.len(), MAX_STDERR_CAPTURE_BYTES);
    assert!(exit.stderr.starts_with("bad?\n"));
}

#[tokio::test]
async fn shutdown_does_not_wait_for_a_full_queue() {
    let rpc = spawn(
        r#"printf '%s\n' '{"method":"one"}' '{"method":"two"}'; exec sleep 5"#,
        1,
    );
    sleep(Duration::from_millis(50)).await;
    let _ = finish(rpc).await;
}