use std::collections::BTreeSet;
use std::fs;
use std::path::PathBuf;
use agent_toolprint::{
canonical, chain, countersign_tool, ed25519_sign, pae_encode, sha256_hash, sign_agent, verify,
DidKeyResolver, VerifyOptions,
};
use base64::{engine::general_purpose::STANDARD, Engine};
use serde_json::{json, Map, Value};
fn vectors_dir() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("vectors")
}
fn collect_vectors() -> Vec<(PathBuf, Value)> {
let mut out = Vec::new();
walk(&vectors_dir(), &mut out);
out.sort_by(|a, b| a.0.cmp(&b.0));
out
}
fn walk(dir: &PathBuf, out: &mut Vec<(PathBuf, Value)>) {
for entry in fs::read_dir(dir).unwrap() {
let entry = entry.unwrap();
let path = entry.path();
if path.is_dir() {
walk(&path, out);
} else if path.extension().and_then(|s| s.to_str()) == Some("json") {
let bytes = fs::read(&path).unwrap();
let v: Value = serde_json::from_slice(&bytes).unwrap();
out.push((path, v));
}
}
}
fn seed_sk(seed: u64) -> [u8; 32] {
[seed as u8; 32]
}
async fn build_envelope(receipt: &Value, agent_seed: u64, tool_seed: u64) -> Value {
let env = sign_agent(receipt, &seed_sk(agent_seed)).unwrap();
countersign_tool(&env, &seed_sk(tool_seed)).unwrap()
}
fn opts<'a>(resolver: &'a DidKeyResolver) -> VerifyOptions<'a> {
let mut o = VerifyOptions::new(resolver);
o.skip_timestamp_check = true;
o
}
fn apply_mutation(env: &Value, mut_v: &Value, agent_seed: u64, tool_seed: u64) -> Value {
let target = mut_v["target"].as_str().unwrap();
if target == "swap-sigs" {
let mut out = env.clone();
let sigs = env["signatures"].as_array().unwrap();
out["signatures"] = json!([sigs[1].clone(), sigs[0].clone()]);
return out;
}
if target == "payload" {
let mut out = env.clone();
let payload = env["payload"].as_str().unwrap();
let first = if payload.starts_with('A') { 'B' } else { 'A' };
let mut new_payload = String::new();
new_payload.push(first);
new_payload.push_str(&payload[1..]);
out["payload"] = json!(new_payload);
return out;
}
if target == "keyid-mismatch" {
let mut out = env.clone();
let i = mut_v.get("index").and_then(Value::as_u64).unwrap_or(0) as usize;
out["signatures"][i]["keyid"] = json!("WRONG-KEYID");
return out;
}
if target == "non-canonical-payload" {
let payload_bytes = STANDARD.decode(env["payload"].as_str().unwrap()).unwrap();
let obj: Map<String, Value> = serde_json::from_slice(&payload_bytes).unwrap();
let mut reordered: Map<String, Value> = Map::new();
for k in obj.keys().rev() {
reordered.insert(k.clone(), obj.get(k).unwrap().clone());
}
let mut non_canon = String::from("{");
let entries: Vec<_> = reordered.iter().collect();
for (i, (k, v)) in entries.iter().enumerate() {
if i > 0 {
non_canon.push_str(", ");
}
non_canon.push_str(&serde_json::to_string(k).unwrap());
non_canon.push_str(": ");
non_canon.push_str(&serde_json::to_string(v).unwrap());
}
non_canon.push('}');
let non_canon_bytes = non_canon.into_bytes();
let pae = pae_encode(env["payloadType"].as_str().unwrap(), &non_canon_bytes);
let agent_sig = ed25519_sign(&seed_sk(agent_seed), &pae);
let tool_sig = ed25519_sign(&seed_sk(tool_seed), &pae);
let sigs = env["signatures"].as_array().unwrap();
return json!({
"payloadType": env["payloadType"],
"payload": STANDARD.encode(&non_canon_bytes),
"signatures": [
{"keyid": sigs[0]["keyid"], "sig": STANDARD.encode(&agent_sig)},
{"keyid": sigs[1]["keyid"], "sig": STANDARD.encode(&tool_sig)},
],
});
}
let mut out = env.clone();
let i = mut_v.get("index").and_then(Value::as_u64).unwrap_or(0) as usize;
let sig = out["signatures"][i]["sig"].as_str().unwrap().to_string();
let first = if sig.starts_with('A') { 'B' } else { 'A' };
let mut new_sig = String::new();
new_sig.push(first);
new_sig.push_str(&sig[1..]);
out["signatures"][i]["sig"] = json!(new_sig);
out
}
#[tokio::test]
async fn clauses_present() {
let vectors = collect_vectors();
assert!(!vectors.is_empty(), "no vectors loaded");
let seen: BTreeSet<String> = vectors
.iter()
.map(|(_, v)| v["clause"].as_str().unwrap().to_string())
.collect();
for c in ["C1", "C2", "C3", "C4"] {
assert!(seen.contains(c), "missing conformance clause {c}");
}
}
#[tokio::test]
async fn all_vectors_pass() {
let resolver = DidKeyResolver;
let vectors = collect_vectors();
let mut total = 0;
let mut passed = 0;
for (path, v) in &vectors {
total += 1;
let clause = v["clause"].as_str().unwrap();
let name = v["name"].as_str().unwrap();
let result = run_vector(&resolver, v).await;
match result {
Ok(_) => passed += 1,
Err(e) => panic!("[{clause}-{name}] {} → {e}", path.display()),
}
}
assert_eq!(passed, total, "{passed}/{total} vectors passed");
}
async fn run_vector(resolver: &DidKeyResolver, v: &Value) -> Result<(), String> {
let clause = v["clause"].as_str().unwrap();
let inp = &v["input"];
let exp = &v["expected"];
match clause {
"C1" => {
let receipt = inp["receipt"].clone();
agent_toolprint::validate_receipt(&receipt).map_err(|e| e.to_string())?;
let canonical_sha = sha256_hash(&receipt).map_err(|e| e.to_string())?;
if let Some(exp_sha) = exp.get("canonical_payload_sha256").and_then(Value::as_str) {
if canonical_sha != exp_sha {
return Err(format!(
"canonical sha mismatch: got {canonical_sha}, expected {exp_sha}"
));
}
}
if exp.get("verify_ok").and_then(Value::as_bool) == Some(true) {
let env = build_envelope(
&receipt,
inp["agent_sk_seed"].as_u64().unwrap(),
inp["tool_sk_seed"].as_u64().unwrap(),
)
.await;
let res = verify(&env, &opts(resolver)).await;
if !res.ok {
return Err(format!("verify failed: {:?}", res.error));
}
}
Ok(())
}
"C2" => {
let receipt = inp["receipt"].clone();
agent_toolprint::validate_receipt(&receipt).map_err(|e| e.to_string())?;
let env = build_envelope(
&receipt,
inp["agent_sk_seed"].as_u64().unwrap(),
inp["tool_sk_seed"].as_u64().unwrap(),
)
.await;
let mutated = apply_mutation(
&env,
&inp["mutation"],
inp["agent_sk_seed"].as_u64().unwrap(),
inp["tool_sk_seed"].as_u64().unwrap(),
);
let res = verify(&mutated, &opts(resolver)).await;
if res.ok {
return Err("mutation was accepted by verify".into());
}
Ok(())
}
"C3" => {
let receipt = inp["receipt"].clone();
agent_toolprint::validate_receipt(&receipt).map_err(|e| e.to_string())?;
let role = inp
.get("role")
.and_then(Value::as_str)
.unwrap_or("agent-only");
if role == "agent-only" {
let partial =
sign_agent(&receipt, &seed_sk(inp["agent_sk_seed"].as_u64().unwrap()))
.map_err(|e| e.to_string())?;
let res = verify(&partial, &opts(resolver)).await;
if res.ok {
return Err("agent-only envelope accepted".into());
}
} else {
let full = build_envelope(
&receipt,
inp["agent_sk_seed"].as_u64().unwrap(),
inp["tool_sk_seed"].as_u64().unwrap(),
)
.await;
let mut tool_only = full.clone();
let sigs = full["signatures"].as_array().unwrap();
tool_only["signatures"] = json!([sigs[1].clone()]);
let res = verify(&tool_only, &opts(resolver)).await;
if res.ok {
return Err("tool-only envelope accepted".into());
}
}
Ok(())
}
"C4" => {
let parent = inp["parent"].clone();
let child = inp["child"].clone();
let parent_env = build_envelope(
&parent,
inp["agent_sk_seed"].as_u64().unwrap(),
inp["tool_sk_seed"].as_u64().unwrap(),
)
.await;
let child_env = build_envelope(
&child,
inp["agent_sk_seed"].as_u64().unwrap(),
inp["tool_sk_seed"].as_u64().unwrap(),
)
.await;
let linked = chain(&parent_env, &child_env);
let expected = exp["chain_ok"].as_bool().unwrap();
if linked != expected {
return Err(format!("chain() returned {linked}, expected {expected}"));
}
Ok(())
}
other => Err(format!("unknown clause {other}")),
}
}
#[tokio::test]
async fn end_to_end_roundtrip() {
use ed25519_dalek::SigningKey;
use rand::rngs::OsRng;
let agent_sk = SigningKey::generate(&mut OsRng);
let tool_sk = SigningKey::generate(&mut OsRng);
let agent_pk = agent_sk.verifying_key().to_bytes();
let tool_pk = tool_sk.verifying_key().to_bytes();
let agent_did = agent_toolprint::did_key_from_ed25519_pubkey(&agent_pk).unwrap();
let tool_did = agent_toolprint::did_key_from_ed25519_pubkey(&tool_pk).unwrap();
let args = json!({"query": "bun docs"});
let response = json!({"results": ["https://bun.sh/docs"]});
let receipt = json!({
"v": "tp/0.1",
"id": "0192b6c7-4e58-7a2d-ae9e-6c77f3e20c44",
"ts": chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
"agent": {"did": agent_did, "key_id": "agent"},
"tool": {"did": tool_did, "key_id": "tool"},
"call": {"name": "search", "args_hash": sha256_hash(&args).unwrap()},
"result": {"status": "ok", "response_hash": sha256_hash(&response).unwrap()},
"nonce": STANDARD.encode([0u8; 32]),
});
let env = sign_agent(&receipt, &agent_sk.to_bytes()).unwrap();
let env = countersign_tool(&env, &tool_sk.to_bytes()).unwrap();
let resolver = DidKeyResolver;
let mut opts = VerifyOptions::new(&resolver);
opts.plaintext = Some(agent_toolprint::Plaintext {
args: Some(args),
response: Some(response),
});
let res = verify(&env, &opts).await;
assert!(res.ok, "verify failed: {:?}", res.error);
let _ = canonical(&receipt).unwrap();
}