polyc-crypto 0.1.3

//! Canonical signing for tool-call provenance.
//!
//! A signed tool call commits to its *canonical bytes*: the buffa encoding of
//! the message with its `signature` field cleared. Both the signer and the
//! verifier clear the field before encoding, so the signature covers every
//! other field (id, oneof body, ...) but not itself. buffa's encoding is
//! deterministic for a given message, making the round-trip reproducible.
//!
//! Use [`sign_tool_call`] / [`sign_tool_result`] to mint a signature, the
//! `*_into` helpers to fill the message's field in place, and
//! [`verify_tool_call`] / [`verify_tool_result`] to check provenance. The
//! verifiers never panic.

use buffa::Message as _;
use polyc_proto::proto::polychrome::agent::v1::{ToolCallContent, ToolResultContent};

use crate::{Signer, verify};

/// Canonical bytes for a [`ToolCallContent`]: the encoding with `signature`
/// cleared, so the signature commits to everything *except* itself.
fn tool_call_canonical_bytes(call: &ToolCallContent) -> Vec<u8> {
    let mut canonical = call.clone();
    canonical.signature.clear();
    canonical.encode_to_vec()
}

/// Canonical bytes for a [`ToolResultContent`]: the encoding with `signature`
/// cleared.
fn tool_result_canonical_bytes(result: &ToolResultContent) -> Vec<u8> {
    let mut canonical = result.clone();
    canonical.signature.clear();
    canonical.encode_to_vec()
}

/// Sign the canonical bytes of `call`; returns signature bytes suitable for
/// [`ToolCallContent::signature`].
#[must_use]
pub fn sign_tool_call(signer: &Signer, call: &ToolCallContent) -> Vec<u8> {
    signer.sign(&tool_call_canonical_bytes(call))
}

/// Sign `call` and store the signature in its `signature` field in place.
pub fn sign_tool_call_into(signer: &Signer, call: &mut ToolCallContent) {
    call.signature = sign_tool_call(signer, call);
}

/// Verify the provenance signature carried in `call.signature` against an
/// encoded `public_key`.
///
/// Returns `false` on any decode failure or signature mismatch — never panics.
#[must_use]
pub fn verify_tool_call(public_key: &[u8], call: &ToolCallContent) -> bool {
    verify(
        public_key,
        &tool_call_canonical_bytes(call),
        &call.signature,
    )
}

/// Sign the canonical bytes of `result`; returns signature bytes suitable for
/// [`ToolResultContent::signature`].
#[must_use]
pub fn sign_tool_result(signer: &Signer, result: &ToolResultContent) -> Vec<u8> {
    signer.sign(&tool_result_canonical_bytes(result))
}

/// Sign `result` and store the signature in its `signature` field in place.
pub fn sign_tool_result_into(signer: &Signer, result: &mut ToolResultContent) {
    result.signature = sign_tool_result(signer, result);
}

/// Verify the provenance signature carried in `result.signature` against an
/// encoded `public_key`.
///
/// Returns `false` on any decode failure or signature mismatch — never panics.
#[must_use]
pub fn verify_tool_result(public_key: &[u8], result: &ToolResultContent) -> bool {
    verify(
        public_key,
        &tool_result_canonical_bytes(result),
        &result.signature,
    )
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]

    use polyc_proto::proto::polychrome::agent::v1::{
        FunctionCallContent, FunctionResultContent, ToolCallContent, ToolResultContent,
    };

    use super::*;

    fn function_call(name: &str) -> FunctionCallContent {
        FunctionCallContent {
            name: name.to_string(),
            ..Default::default()
        }
    }

    fn sample_call() -> ToolCallContent {
        ToolCallContent {
            id: "call-1".to_string(),
            r#type: Some(function_call("web_search").into()),
            ..Default::default()
        }
    }

    fn sample_result() -> ToolResultContent {
        ToolResultContent {
            call_id: "call-1".to_string(),
            r#type: Some(
                FunctionResultContent {
                    name: "web_search".to_string(),
                    ..Default::default()
                }
                .into(),
            ),
            ..Default::default()
        }
    }

    #[test]
    fn tool_call_round_trips() {
        let signer = Signer::from_seed(7);
        let pk = signer.public_key_bytes();
        let mut call = sample_call();
        sign_tool_call_into(&signer, &mut call);
        assert!(!call.signature.is_empty());
        assert!(verify_tool_call(&pk, &call));
    }

    #[test]
    fn tool_call_tampered_id_fails() {
        let signer = Signer::from_seed(7);
        let pk = signer.public_key_bytes();
        let mut call = sample_call();
        sign_tool_call_into(&signer, &mut call);
        call.id = "call-2".to_string();
        assert!(!verify_tool_call(&pk, &call));
    }

    #[test]
    fn tool_call_tampered_args_fails() {
        let signer = Signer::from_seed(7);
        let pk = signer.public_key_bytes();
        let mut call = sample_call();
        sign_tool_call_into(&signer, &mut call);
        call.r#type = Some(function_call("rm_rf").into());
        assert!(!verify_tool_call(&pk, &call));
    }

    #[test]
    fn tool_call_wrong_key_fails() {
        let signer = Signer::from_seed(7);
        let other = Signer::from_seed(8);
        let mut call = sample_call();
        sign_tool_call_into(&signer, &mut call);
        assert!(!verify_tool_call(&other.public_key_bytes(), &call));
    }

    #[test]
    fn tool_call_unsigned_fails() {
        let signer = Signer::from_seed(7);
        let call = sample_call();
        assert!(!verify_tool_call(&signer.public_key_bytes(), &call));
    }

    #[test]
    fn tool_result_round_trips() {
        let signer = Signer::from_seed(9);
        let pk = signer.public_key_bytes();
        let mut result = sample_result();
        sign_tool_result_into(&signer, &mut result);
        assert!(!result.signature.is_empty());
        assert!(verify_tool_result(&pk, &result));
    }

    #[test]
    fn tool_result_tampered_call_id_fails() {
        let signer = Signer::from_seed(9);
        let pk = signer.public_key_bytes();
        let mut result = sample_result();
        sign_tool_result_into(&signer, &mut result);
        result.call_id = "call-99".to_string();
        assert!(!verify_tool_result(&pk, &result));
    }

    #[test]
    fn tool_result_wrong_key_fails() {
        let signer = Signer::from_seed(9);
        let other = Signer::from_seed(10);
        let mut result = sample_result();
        sign_tool_result_into(&signer, &mut result);
        assert!(!verify_tool_result(&other.public_key_bytes(), &result));
    }
}