boxmaker/docs/plans/M5c1/files/crates/brokerd/tests/state_socket.rs

//! One `state.sock` connection: `memoryd` asks for a session's state, and `brokerd` answers from
//! its state store, changing nothing (M5c spec, section 7). Every error is an error frame, never
//! an answer. Do not edit.

#[path = "support/client.rs"]
mod client;
#[path = "support/rig.rs"]
mod rig;
#[path = "support/runtime.rs"]
mod runtime;
#[path = "support/sink.rs"]
mod sink;
#[path = "support/tmp.rs"]
mod tmp;

use std::io::Write;
use std::os::unix::net::UnixStream;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;

use brokerd::broker::Broker;
use brokerd::policy::{Label, SessionState};
use brokerd::state_socket;
use bxcore::{
    CallId, DataClass, Empty, Envelope, ErrorCode, FrameError, Message, PROTOCOL_VERSION,
    SessionId, SessionState as StateQuestion, SessionStateReply, SessionUpdated, ToolRequest,
    WireError,
};
use client::{Serve, next, open};
use rig::Rig;
use runtime::Recording;

fn broker(rig: &Rig) -> Arc<Broker> {
    rig.broker(&Recording::answering("x"))
}

fn question(session: &str) -> Message {
    Message::SessionState(StateQuestion {
        session: SessionId::new(session).unwrap(),
    })
}

/// Asks once, under `id`, and returns the one frame that comes back.
fn ask(broker: &Arc<Broker>, id: u64, msg: Message) -> Envelope {
    let mut stream = open(broker, state_socket::handle, id, msg);
    let env = next(&mut stream);
    assert_eq!((env.id, env.r#final), (id, true), "{env:?}");
    env
}

fn reply(env: &Envelope) -> SessionStateReply {
    match &env.msg {
        Message::SessionStateReply(reply) => reply.clone(),
        other => panic!("wanted a session_state_reply, got {other:?}"),
    }
}

fn error(env: &Envelope) -> WireError {
    match &env.msg {
        Message::Error(e) => e.clone(),
        other => panic!("wanted an error, got {other:?}"),
    }
}

/// Raises a session's state the way the ledger does, through the store's own `raise`.
fn raise(rig: &Rig, session: &str, class: DataClass, untrusted: bool) {
    rig.state()
        .raise(
            &SessionId::new(session).unwrap(),
            SessionState::default(),
            Label { class, untrusted },
        )
        .unwrap();
}

/// Every file under `dir`, with its bytes, sorted; empty when `dir` does not exist.
fn files(dir: &Path) -> Vec<(String, Vec<u8>)> {
    let mut all = Vec::new();
    if let Ok(entries) = std::fs::read_dir(dir) {
        for entry in entries {
            let entry = entry.unwrap();
            all.push((
                entry.file_name().into_string().unwrap(),
                std::fs::read(entry.path()).unwrap_or_default(),
            ));
        }
    }
    all.sort();
    all
}

#[test]
fn the_answer_is_the_stored_class_and_flag() {
    let rig = Rig::new("state-sock-stored");
    raise(&rig, "a", DataClass::Private, false);
    raise(&rig, "b", DataClass::Private, true);
    raise(&rig, "c", DataClass::Secret, false);
    raise(&rig, "d", DataClass::Secret, true);
    // A public result leaves a session at private: a session is never below private.
    raise(&rig, "e", DataClass::Public, true);
    let broker = broker(&rig);
    for (i, (session, class, untrusted)) in [
        ("a", DataClass::Private, false),
        ("b", DataClass::Private, true),
        ("c", DataClass::Secret, false),
        ("d", DataClass::Secret, true),
        ("e", DataClass::Private, true),
    ]
    .into_iter()
    .enumerate()
    {
        let id = 40 + i as u64;
        assert_eq!(
            reply(&ask(&broker, id, question(session))),
            SessionStateReply { class, untrusted },
            "session {session}"
        );
    }
    assert!(rig.lines.all().is_empty(), "{:?}", rig.lines.all());
}

/// A file written by hand, as the runbook tells the owner to write one, is read the same way.
#[test]
fn a_state_file_the_owner_wrote_is_answered_as_it_says() {
    let rig = Rig::new("state-sock-owner");
    std::fs::create_dir_all(rig.cfg.state_dir()).unwrap();
    std::fs::write(
        rig.state_file("s1"),
        "{\"taint\":\"secret\",\"untrusted\":true}\n",
    )
    .unwrap();
    let env = ask(&broker(&rig), 7, question("s1"));
    assert_eq!(
        reply(&env),
        SessionStateReply {
            class: DataClass::Secret,
            untrusted: true,
        }
    );
}

/// A session `brokerd` has not seen has made no tool call: `private`, trusted (M5c spec, section
/// 7). It is answered without making its file, or the directory.
#[test]
fn a_session_brokerd_has_not_seen_is_private_and_trusted() {
    let rig = Rig::new("state-sock-unseen");
    let broker = broker(&rig);
    let env = ask(&broker, 3, question("never-seen"));
    assert_eq!(
        reply(&env),
        SessionStateReply {
            class: DataClass::Private,
            untrusted: false,
        }
    );
    assert!(!rig.cfg.state_dir().exists(), "asking makes no directory");
    // With the directory there and another session in it, still no file for this one.
    raise(&rig, "other", DataClass::Secret, true);
    let env = ask(&broker, 4, question("never-seen"));
    assert_eq!(
        reply(&env),
        SessionStateReply {
            class: DataClass::Private,
            untrusted: false,
        }
    );
    assert!(!rig.state_file("never-seen").exists());
}

#[test]
fn asking_changes_nothing() {
    let rig = Rig::new("state-sock-nothing");
    raise(&rig, "s1", DataClass::Secret, true);
    raise(&rig, "s2", DataClass::Private, false);
    let before = files(&rig.cfg.state_dir());
    let broker = broker(&rig);
    for id in 0..3 {
        for session in ["s1", "s2", "s3"] {
            ask(&broker, id, question(session));
        }
    }
    assert_eq!(files(&rig.cfg.state_dir()), before);
    assert!(
        rig.events().is_empty(),
        "nothing is recorded for a question"
    );
    assert!(rig.lines.all().is_empty(), "{:?}", rig.lines.all());
}

/// A state file that cannot be read is an `internal` error with the store's reason and runbook
/// pointer, logged once, never an answer: `memoryd` then leaves the turn undecided.
#[test]
fn a_damaged_state_is_an_error_never_an_answer() {
    let cases: [(&str, Option<&str>); 5] = [
        ("not json", Some("{\"taint\":")),
        (
            "public",
            Some("{\"taint\":\"public\",\"untrusted\":false}\n"),
        ),
        (
            "unknown key",
            Some("{\"taint\":\"private\",\"untrusted\":false,\"x\":1}\n"),
        ),
        (
            "unknown class",
            Some("{\"taint\":\"loud\",\"untrusted\":false}\n"),
        ),
        ("a directory", None),
    ];
    for (i, (name, text)) in cases.into_iter().enumerate() {
        let rig = Rig::new("state-sock-damaged");
        std::fs::create_dir_all(rig.cfg.state_dir()).unwrap();
        match text {
            Some(text) => std::fs::write(rig.state_file("s1"), text).unwrap(),
            None => std::fs::create_dir(rig.state_file("s1")).unwrap(),
        }
        let before = files(&rig.cfg.state_dir());
        let env = ask(&broker(&rig), 60 + i as u64, question("s1"));
        let e = error(&env);
        assert_eq!(e.code, ErrorCode::Internal, "{name}: {e:?}");
        assert!(e.detail.contains("s1.json"), "{name}: {}", e.detail);
        assert!(
            e.detail
                .ends_with("see docs/runbook.md#broker-state-damaged"),
            "{name}: {}",
            e.detail
        );
        let lines = rig.lines.all();
        assert_eq!(lines.len(), 1, "{name}: {lines:?}");
        assert!(lines[0].starts_with("brokerd: "), "{name}: {lines:?}");
        assert!(lines[0].contains("s1.json"), "{name}: {lines:?}");
        assert!(
            lines[0].ends_with("see docs/runbook.md#broker-state-damaged"),
            "{name}: {lines:?}"
        );
        assert_eq!(files(&rig.cfg.state_dir()), before, "{name}");
    }
}

#[test]
fn every_other_kind_is_forbidden_and_logged() {
    let rig = Rig::new("state-sock-forbidden");
    raise(&rig, "s1", DataClass::Secret, true);
    let broker = broker(&rig);
    let session = SessionId::new("s1").unwrap();
    let kinds = vec![
        Message::ToolRequest(ToolRequest {
            session: session.clone(),
            call: CallId(1),
            tool: "read_file".to_string(),
            arguments: "{}".to_string(),
        }),
        Message::Approvals(Empty {}),
        Message::CheckGrants(Empty {}),
        Message::Ok(Empty {}),
        Message::SessionUpdated(SessionUpdated {
            session: session.clone(),
            turn: 1,
        }),
        // An answer sent as a question is not a question.
        Message::SessionStateReply(SessionStateReply {
            class: DataClass::Private,
            untrusted: false,
        }),
    ];
    let count = kinds.len();
    for (i, msg) in kinds.into_iter().enumerate() {
        let id = 100 + i as u64;
        let e = error(&ask(&broker, id, msg));
        assert_eq!(e.code, ErrorCode::Forbidden, "{e:?}");
        assert!(e.detail.ends_with("is not accepted on state.sock"), "{e:?}");
    }
    let lines = rig.lines.with("on state.sock");
    assert_eq!(lines.len(), count, "{:?}", rig.lines.all());
    for line in &lines {
        assert!(
            line.ends_with("\nsee docs/runbook.md#socket-forbidden"),
            "{line}"
        );
    }
    assert_eq!(
        rig.lines
            .with("refused the message kind session_state_reply on state.sock")
            .len(),
        1
    );
    assert!(rig.events().is_empty(), "nothing is recorded");
}

/// A frame that is not a well-formed question is answered with the frame error, id 0, as on the
/// other sockets. A session id that is not one (a path, say) is refused when the frame is read.
#[test]
fn a_malformed_frame_is_answered_with_an_error() {
    let rig = Rig::new("state-sock-malformed");
    raise(&rig, "s1", DataClass::Secret, true);
    let broker = broker(&rig);
    let body = |session: &str, v: u32, extra: &str| {
        format!(
            "{{\"v\":{v},\"id\":5,\"final\":true,\"msg\":{{\"kind\":\"session_state\",\
             \"body\":{{\"session\":\"{session}\"{extra}}}}}}}"
        )
    };
    let cases = [
        ("not json", "{{{".to_string(), ErrorCode::BadMessage),
        ("a path", body("../s1", 1, ""), ErrorCode::BadMessage),
        ("upper case", body("S1", 1, ""), ErrorCode::BadMessage),
        (
            "unknown key",
            body("s1", 1, ",\"x\":1"),
            ErrorCode::BadMessage,
        ),
        ("version 2", body("s1", 2, ""), ErrorCode::BadVersion),
    ];
    for (name, text, code) in cases {
        let (mut client, server) = UnixStream::pair().unwrap();
        let b = Arc::clone(&broker);
        std::thread::spawn(move || state_socket::handle(server, &b));
        client
            .write_all(&(text.len() as u32).to_be_bytes())
            .unwrap();
        client.write_all(text.as_bytes()).unwrap();
        let env = next(&mut client);
        assert_eq!((env.id, env.r#final), (0, true), "{name}");
        assert_eq!(error(&env).code, code, "{name}");
    }
    // The well-formed question, for contrast.
    let env = ask(&broker, 5, question("s1"));
    assert_eq!(
        reply(&env),
        SessionStateReply {
            class: DataClass::Secret,
            untrusted: true,
        }
    );
}

/// One question per connection: after its answer the connection is closed, and a second question
/// on it is never answered.
#[test]
fn a_second_question_on_one_connection_is_not_answered() {
    let rig = Rig::new("state-sock-once");
    raise(&rig, "s1", DataClass::Secret, true);
    let broker = broker(&rig);
    let mut stream = open(&broker, state_socket::handle, 1, question("s1"));
    assert!(matches!(
        next(&mut stream).msg,
        Message::SessionStateReply(_)
    ));
    // The handler may have closed its end already, so the second frame may not be written.
    let _ = bxcore::write_frame(
        &mut stream,
        &Envelope {
            v: PROTOCOL_VERSION,
            id: 2,
            r#final: true,
            msg: question("s1"),
        },
    );
    // Once the peer has closed, macOS refuses the timeout with EINVAL (22); the read cannot block
    // then.
    if let Err(e) = stream.set_read_timeout(Some(Duration::from_secs(10))) {
        assert_eq!(e.raw_os_error(), Some(22), "set_read_timeout: {e}");
    }
    match bxcore::read_frame(&mut stream) {
        Err(FrameError::Closed) => {}
        // Linux resets a Unix socket closed with unread data in it: closed as well.
        Err(FrameError::Io(e)) if e.kind() == std::io::ErrorKind::ConnectionReset => {}
        other => panic!("wanted the connection closed, got {other:?}"),
    }
}