brokerd: state.sock answers memoryd's question about a session's state
Co-Authored-By: ornith-1.5-35b-a3b <agent@wntrmute.dev>
This commit is contained in:
@@ -0,0 +1,371 @@
|
||||
//! 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:?}"),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user