use buffa::Message as _;
use polyc_proto::proto::polychrome::subagent::v1::{SubagentResult, SubagentSpawn};
use crate::{Signer, verify};
fn spawn_canonical_bytes(spawn: &SubagentSpawn) -> Vec<u8> {
let mut canonical = spawn.clone();
canonical.signature_hex.clear();
canonical.encode_to_vec()
}
fn result_canonical_bytes(result: &SubagentResult) -> Vec<u8> {
let mut canonical = result.clone();
canonical.signature_hex.clear();
canonical.encode_to_vec()
}
pub fn sign_subagent_spawn_into(signer: &Signer, spawn: &mut SubagentSpawn) {
spawn.signed_by = signer.public_key_bytes();
spawn.signature_hex.clear();
let sig = signer.sign(&spawn_canonical_bytes(spawn));
spawn.signature_hex = crate::hex::lower(&sig);
}
pub fn sign_subagent_result_into(signer: &Signer, result: &mut SubagentResult) {
result.signed_by = signer.public_key_bytes();
result.signature_hex.clear();
let sig = signer.sign(&result_canonical_bytes(result));
result.signature_hex = crate::hex::lower(&sig);
}
#[must_use]
pub fn verify_subagent_spawn(public_key: &[u8], spawn: &SubagentSpawn) -> bool {
let Some(sig) = crate::hex::decode(&spawn.signature_hex) else {
return false;
};
verify(public_key, &spawn_canonical_bytes(spawn), &sig)
}
#[must_use]
pub fn verify_subagent_result(public_key: &[u8], result: &SubagentResult) -> bool {
let Some(sig) = crate::hex::decode(&result.signature_hex) else {
return false;
};
verify(public_key, &result_canonical_bytes(result), &sig)
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs)]
use super::*;
fn sample_spawn() -> SubagentSpawn {
SubagentSpawn {
sub_agent_id: "call-7".to_owned(),
target_agent_id: "researcher".to_owned(),
resolved_provider: "vertex".to_owned(),
resolved_model: "fable-pro".to_owned(),
task: "find prior art".to_owned(),
..Default::default()
}
}
fn sample_result() -> SubagentResult {
SubagentResult {
sub_agent_id: "call-7".to_owned(),
target_agent_id: "researcher".to_owned(),
succeeded: true,
error: String::new(),
input_tokens: 42,
output_tokens: 7,
first_party: true,
..Default::default()
}
}
#[test]
fn spawn_round_trips() {
let signer = Signer::from_seed(21);
let mut s = sample_spawn();
sign_subagent_spawn_into(&signer, &mut s);
assert!(!s.signature_hex.is_empty());
assert_eq!(s.signed_by, signer.public_key_bytes());
assert!(verify_subagent_spawn(&s.signed_by, &s));
}
#[test]
fn spawn_tampered_task_fails() {
let signer = Signer::from_seed(21);
let mut s = sample_spawn();
sign_subagent_spawn_into(&signer, &mut s);
s.task = "leaked task".to_owned();
assert!(!verify_subagent_spawn(&s.signed_by, &s));
}
#[test]
fn spawn_wrong_key_fails() {
let signer = Signer::from_seed(21);
let other = Signer::from_seed(22);
let mut s = sample_spawn();
sign_subagent_spawn_into(&signer, &mut s);
assert!(!verify_subagent_spawn(&other.public_key_bytes(), &s));
}
#[test]
fn spawn_unsigned_fails() {
let signer = Signer::from_seed(21);
let s = sample_spawn();
assert!(!verify_subagent_spawn(&signer.public_key_bytes(), &s));
}
#[test]
fn result_round_trips() {
let signer = Signer::from_seed(23);
let mut r = sample_result();
sign_subagent_result_into(&signer, &mut r);
assert!(!r.signature_hex.is_empty());
assert_eq!(r.signed_by, signer.public_key_bytes());
assert!(verify_subagent_result(&r.signed_by, &r));
}
#[test]
fn result_tampered_usage_fails() {
let signer = Signer::from_seed(23);
let mut r = sample_result();
sign_subagent_result_into(&signer, &mut r);
r.input_tokens = 999_999;
assert!(!verify_subagent_result(&r.signed_by, &r));
}
#[test]
fn result_tampered_succeeded_fails() {
let signer = Signer::from_seed(23);
let mut r = sample_result();
sign_subagent_result_into(&signer, &mut r);
r.succeeded = false;
r.error = "worker turn failed: forged".to_owned();
assert!(!verify_subagent_result(&r.signed_by, &r));
}
#[test]
fn tampered_first_party_fails_verification() {
let signer = Signer::from_seed(23);
let mut r = sample_result();
sign_subagent_result_into(&signer, &mut r);
r.first_party = false;
assert!(!verify_subagent_result(&r.signed_by, &r));
}
#[test]
fn result_wrong_key_fails() {
let signer = Signer::from_seed(23);
let other = Signer::from_seed(24);
let mut r = sample_result();
sign_subagent_result_into(&signer, &mut r);
assert!(!verify_subagent_result(&other.public_key_bytes(), &r));
}
#[test]
fn result_unsigned_fails() {
let signer = Signer::from_seed(23);
let r = sample_result();
assert!(!verify_subagent_result(&signer.public_key_bytes(), &r));
}
#[test]
fn garbage_signature_hex_returns_false() {
let mut s = sample_spawn();
s.signature_hex = "not-hex!".to_owned();
assert!(!verify_subagent_spawn(b"any-key", &s));
s.signature_hex = "abcd".to_owned();
assert!(!verify_subagent_spawn(b"any-key", &s));
}
}