use std::io::{self, BufRead, BufReader, Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::thread::JoinHandle;
use std::time::Duration;
use super::{LysAgent, PassSource, Proof, RENEW_WITHIN_SECS};
use crate::mcp::adapter::ServiceRefusal;
type TestResult = Result<(), Box<dyn std::error::Error>>;
const CREDENTIAL: &str = "grant-credential-never-shown";
const START: u64 = 1_000_000;
const VALID: u64 = 600;
struct Seen {
head: String,
body: Vec<u8>,
}
struct FakeLys {
address: SocketAddr,
server: JoinHandle<io::Result<Vec<Seen>>>,
}
fn fake_lys(answers: Vec<(u16, String)>) -> io::Result<FakeLys> {
let listener = TcpListener::bind("127.0.0.1:0")?;
let address = listener.local_addr()?;
let server = std::thread::spawn(move || {
let mut seen = Vec::new();
let mut answers = answers.into_iter();
for stream in listener.incoming() {
let stream = stream?;
let Some(request) = read_request(&stream)? else {
break;
};
seen.push(request);
let (status, body) = answers
.next()
.unwrap_or_else(|| (500, r#"{"refusal":"TestAnswersSpent"}"#.to_owned()));
let response = format!(
"HTTP/1.1 {status} Canned\r\ncontent-type: application/json\r\n\
content-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
);
let mut out = &stream;
out.write_all(response.as_bytes())?;
out.flush()?;
}
Ok(seen)
});
Ok(FakeLys { address, server })
}
fn read_request(stream: &TcpStream) -> io::Result<Option<Seen>> {
let mut reader = BufReader::new(stream);
let mut head = String::new();
loop {
let mut line = String::new();
if reader.read_line(&mut line)? == 0 {
return Ok(None);
}
if line == "\r\n" {
break;
}
head.push_str(&line);
}
let length = head
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
let mut body = vec![0; length];
reader.read_exact(&mut body)?;
Ok(Some(Seen { head, body }))
}
impl FakeLys {
fn issuer(&self) -> String {
format!("http://{}", self.address)
}
fn finish(self) -> Result<Vec<Seen>, Box<dyn std::error::Error>> {
drop(TcpStream::connect(self.address)?);
let seen = self
.server
.join()
.map_err(|_| "the fake Lys thread panicked")??;
Ok(seen)
}
}
fn pass_answer(pass: &str, expires_at: u64) -> (u16, String) {
(
200,
format!(r#"{{"pass":"{pass}","expires_at":{expires_at},"audience":"haematite"}}"#),
)
}
struct Fixture {
fake: FakeLys,
source: PassSource,
now: Arc<AtomicU64>,
_directory: tempfile::TempDir,
}
fn fixture(answers: Vec<(u16, String)>) -> Result<Fixture, Box<dyn std::error::Error>> {
let directory = tempfile::tempdir()?;
let credential = directory.path().join("credential");
std::fs::write(&credential, format!("{CREDENTIAL}\n"))?;
let fake = fake_lys(answers)?;
let mut agent = LysAgent::new(
&fake.issuer(),
"agent-1",
Proof::Grant(credential),
Duration::from_secs(5),
)?;
let now = Arc::new(AtomicU64::new(START));
let clock = Arc::clone(&now);
agent.clock = Box::new(move || clock.load(Ordering::SeqCst));
Ok(Fixture {
fake,
source: PassSource::Lys(agent),
now,
_directory: directory,
})
}
fn accept(presented: &mut Vec<String>) -> impl FnMut(&str) -> Result<(), ServiceRefusal> + '_ {
move |pass| {
presented.push(pass.to_owned());
Ok(())
}
}
fn refuse<'a>(
presented: &'a mut Vec<String>,
code: &'static str,
refusals: usize,
) -> impl FnMut(&str) -> Result<(), ServiceRefusal> + 'a {
move |pass| {
presented.push(pass.to_owned());
if presented.len() <= refusals {
Err(ServiceRefusal {
code: code.into(),
message: "the Lys pass was refused".into(),
receipt: None,
})
} else {
Ok(())
}
}
}
fn accepted(outcome: Result<(), ServiceRefusal>) -> Result<(), Box<dyn std::error::Error>> {
outcome.map_err(|refusal| format!("{}: {}", refusal.code, refusal.message).into())
}
#[test]
fn the_first_call_asks_lys_and_presents_its_pass() -> TestResult {
let mut test = fixture(vec![pass_answer("pass-1", START + VALID)])?;
let mut presented = Vec::new();
accepted(test.source.present(accept(&mut presented)))?;
assert_eq!(presented, ["pass-1"]);
let seen = test.fake.finish()?;
assert_eq!(seen.len(), 1);
let head = seen[0].head.to_ascii_lowercase();
assert!(
head.starts_with("post /api/agents/agent-1/pass http/1.1\r\n"),
"{head}"
);
assert!(head.contains(&format!("\r\n{}: {CREDENTIAL}\r\n", super::GRANT_HEADER)));
assert!(head.contains("\r\ncontent-type: application/json\r\n"));
let body: serde_json::Value = serde_json::from_slice(&seen[0].body)?;
assert_eq!(body, serde_json::json!({"audience": "haematite"}));
Ok(())
}
#[test]
fn a_second_call_within_the_window_reuses_the_held_pass() -> TestResult {
let mut test = fixture(vec![pass_answer("pass-1", START + VALID)])?;
let mut presented = Vec::new();
accepted(test.source.present(accept(&mut presented)))?;
test.now
.store(START + VALID - RENEW_WITHIN_SECS - 1, Ordering::SeqCst);
accepted(test.source.present(accept(&mut presented)))?;
assert_eq!(presented, ["pass-1", "pass-1"]);
assert_eq!(test.fake.finish()?.len(), 1);
Ok(())
}
#[test]
fn a_held_pass_within_sixty_seconds_of_expiry_is_renewed() -> TestResult {
let mut test = fixture(vec![
pass_answer("pass-1", START + VALID),
pass_answer("pass-2", START + 2 * VALID),
])?;
let mut presented = Vec::new();
accepted(test.source.present(accept(&mut presented)))?;
test.now
.store(START + VALID - RENEW_WITHIN_SECS, Ordering::SeqCst);
accepted(test.source.present(accept(&mut presented)))?;
assert_eq!(presented, ["pass-1", "pass-2"]);
assert_eq!(test.fake.finish()?.len(), 2);
Ok(())
}
#[test]
fn a_store_refusal_renews_once_and_presents_once_more() -> TestResult {
let mut test = fixture(vec![
pass_answer("pass-1", START + VALID),
pass_answer("pass-2", START + VALID),
])?;
let mut presented = Vec::new();
accepted(
test.source
.present(refuse(&mut presented, "pass_expired", 1)),
)?;
assert_eq!(presented, ["pass-1", "pass-2"]);
assert_eq!(test.fake.finish()?.len(), 2);
Ok(())
}
#[test]
fn a_store_refusing_the_renewed_pass_too_gets_no_third_try() -> TestResult {
let mut test = fixture(vec![
pass_answer("pass-1", START + VALID),
pass_answer("pass-2", START + VALID),
pass_answer("pass-3", START + VALID),
])?;
let mut presented = Vec::new();
let outcome = test
.source
.present(refuse(&mut presented, "signature_refused", usize::MAX));
let refused = outcome.err().ok_or("a twice-refused pass was accepted")?;
assert_eq!(refused.code, "signature_refused");
assert_eq!(presented, ["pass-1", "pass-2"]);
assert_eq!(test.fake.finish()?.len(), 2);
Ok(())
}
#[test]
fn a_store_refusal_a_new_pass_cannot_answer_is_not_renewed() -> TestResult {
let mut test = fixture(vec![pass_answer("pass-1", START + VALID)])?;
let mut presented = Vec::new();
let outcome = test
.source
.present(refuse(&mut presented, "permission_refused", usize::MAX));
let refused = outcome.err().ok_or("a refused pass was accepted")?;
assert_eq!(refused.code, "permission_refused");
assert_eq!(presented, ["pass-1"]);
assert_eq!(test.fake.finish()?.len(), 1);
Ok(())
}
#[test]
fn a_lys_refusal_is_named_and_never_carries_the_credential() -> TestResult {
let mut test = fixture(vec![(
401,
format!(
r#"{{"refusal":"AgentPassRefused","reason":"{CREDENTIAL} is not held","fields":[]}}"#
),
)])?;
let mut presented = Vec::new();
let outcome = test.source.present(accept(&mut presented));
let refused = outcome.err().ok_or("a Lys refusal was accepted")?;
assert_eq!(refused.code, "mcp_lys_pass_refused");
assert!(
refused.message.contains("AgentPassRefused"),
"{}",
refused.message
);
assert!(refused.message.contains("401"), "{}", refused.message);
assert!(!refused.message.contains(CREDENTIAL), "{}", refused.message);
assert!(presented.is_empty());
assert_eq!(test.fake.finish()?.len(), 1);
Ok(())
}
#[test]
fn an_unreachable_lys_is_named_and_never_carries_the_credential() -> TestResult {
let directory = tempfile::tempdir()?;
let credential = directory.path().join("credential");
std::fs::write(&credential, CREDENTIAL)?;
let closed = TcpListener::bind("127.0.0.1:0")?.local_addr()?;
let agent = LysAgent::new(
&format!("http://{closed}"),
"agent-1",
Proof::Grant(credential),
Duration::from_secs(5),
)?;
let mut source = PassSource::Lys(agent);
let mut presented = Vec::new();
let refused = source
.present(accept(&mut presented))
.err()
.ok_or("an unreachable Lys answered")?;
assert_eq!(refused.code, "mcp_lys_unreachable");
assert!(!refused.message.contains(CREDENTIAL), "{}", refused.message);
assert!(presented.is_empty());
Ok(())
}
#[test]
fn a_key_proof_signs_the_request_and_sends_no_grant_credential() -> TestResult {
use base64::Engine;
let directory = tempfile::tempdir()?;
let key = directory.path().join("seat.json");
let der = [
&[
0x30, 0x2e, 0x02, 0x01, 0x00, 0x30, 0x05, 0x06, 0x03, 0x2b, 0x65, 0x70, 0x04, 0x22,
0x04, 0x20,
][..],
b"lys-cose-conformance-test-seed01",
]
.concat();
let pem = format!(
"-----BEGIN PRIVATE KEY-----\n{}\n-----END PRIVATE KEY-----\n",
base64::engine::general_purpose::STANDARD.encode(der)
);
std::fs::write(
&key,
serde_json::json!({ "privateKeyPem": pem }).to_string(),
)?;
let fake = fake_lys(vec![pass_answer("pass-1", START + VALID)])?;
let mut agent = LysAgent::new(
&fake.issuer(),
"agent-1",
Proof::Key(key),
Duration::from_secs(5),
)?;
agent.clock = Box::new(|| START);
let mut presented = Vec::new();
accepted(PassSource::Lys(agent).present(accept(&mut presented)))?;
assert_eq!(presented, ["pass-1"]);
let seen = fake.finish()?;
let head = seen[0].head.to_ascii_lowercase();
let prefix = format!("\r\n{}: agent-1 {} ", super::SIGNATURE_HEADER, START * 1000);
assert!(head.contains(&prefix), "{head}");
assert!(!head.contains(super::GRANT_HEADER), "{head}");
assert!(!head.contains(&pem.to_ascii_lowercase()));
Ok(())
}
#[test]
fn a_run_pass_proof_sends_the_launched_agents_run_pass_and_nothing_else() -> TestResult {
let fake = fake_lys(vec![pass_answer("pass-1", START + VALID)])?;
let mut agent = LysAgent::new(
&fake.issuer(),
"agent-1",
Proof::RunPass("LYS_AGENT_PASS".to_owned()),
Duration::from_secs(5),
)?;
agent.clock = Box::new(|| START);
agent.environment =
Box::new(|name| (name == "LYS_AGENT_PASS").then(|| "run-pass-1\n".to_owned()));
assert_eq!(
agent.ask_now().map_err(|refusal| refusal.message)?,
"pass-1"
);
let seen = fake.finish()?;
let head = seen[0].head.to_ascii_lowercase();
assert!(
head.contains("\r\nlys-agent-pass: run-pass-1\r\n"),
"{head}"
);
assert!(!head.contains(super::GRANT_HEADER), "{head}");
assert!(!head.contains(super::SIGNATURE_HEADER), "{head}");
assert_eq!(seen[0].body, super::REQUEST);
Ok(())
}
#[test]
fn an_unset_run_pass_variable_is_named_and_lys_is_never_asked() -> TestResult {
let fake = fake_lys(Vec::new())?;
let mut agent = LysAgent::new(
&fake.issuer(),
"agent-1",
Proof::RunPass("LYS_AGENT_PASS".to_owned()),
Duration::from_secs(5),
)?;
agent.environment = Box::new(|_| None);
let refusal = agent.ask_now().err().ok_or("an unset variable was taken")?;
assert_eq!(refusal.code, "mcp_lys_credential_unreadable");
assert!(
refusal.message.contains("LYS_AGENT_PASS"),
"{}",
refusal.message
);
assert!(fake.finish()?.is_empty());
Ok(())
}
#[test]
fn a_run_pass_variable_name_must_be_plain_upper_case() {
for name in ["", "lys_agent_pass", "LYS-PASS", "LYS PASS"] {
assert!(
LysAgent::new(
"http://127.0.0.1:1",
"agent-1",
Proof::RunPass(name.to_owned()),
Duration::from_secs(1)
)
.is_err(),
"{name:?} was taken"
);
}
}