canon_core/

proof.rs

//! Proof receipts for Canon Protocol
//!
//! Generates and verifies cryptographic proof receipts that demonstrate
//! exactly what context an AI had access to at a given point in time.

use ed25519_dalek::{Signature, Verifier, VerifyingKey};
use serde::{Deserialize, Serialize};

use crate::state::{compute_merkle_root, generate_merkle_proof, verify_merkle_proof};

// ============================================================================
// Hex serde helpers for byte arrays
// ============================================================================

mod hex_bytes {
    use serde::{Deserialize, Deserializer, Serializer};

    pub mod b16 {
        use super::{Deserialize, Deserializer, Serializer};
        pub fn serialize<S: Serializer>(bytes: &[u8; 16], s: S) -> Result<S::Ok, S::Error> {
            use std::fmt::Write;
            let mut hex = String::with_capacity(32);
            for b in bytes {
                let _ = write!(hex, "{b:02x}");
            }
            s.serialize_str(&hex)
        }
        pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<[u8; 16], D::Error> {
            let s = String::deserialize(d)?;
            let bytes = (0..s.len())
                .step_by(2)
                .map(|i| u8::from_str_radix(&s[i..i + 2], 16).map_err(serde::de::Error::custom))
                .collect::<Result<Vec<u8>, _>>()?;
            bytes
                .try_into()
                .map_err(|_| serde::de::Error::custom("expected 16 bytes"))
        }
    }

    pub mod b32 {
        use super::{Deserialize, Deserializer, Serializer};
        pub fn serialize<S: Serializer>(bytes: &[u8; 32], s: S) -> Result<S::Ok, S::Error> {
            use std::fmt::Write;
            let mut hex = String::with_capacity(64);
            for b in bytes {
                let _ = write!(hex, "{b:02x}");
            }
            s.serialize_str(&hex)
        }
        pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<[u8; 32], D::Error> {
            let s = String::deserialize(d)?;
            let bytes = (0..s.len())
                .step_by(2)
                .map(|i| u8::from_str_radix(&s[i..i + 2], 16).map_err(serde::de::Error::custom))
                .collect::<Result<Vec<u8>, _>>()?;
            bytes
                .try_into()
                .map_err(|_| serde::de::Error::custom("expected 32 bytes"))
        }
    }

    pub mod b64 {
        use super::{Deserialize, Deserializer, Serializer};
        pub fn serialize<S: Serializer>(bytes: &[u8; 64], s: S) -> Result<S::Ok, S::Error> {
            use std::fmt::Write;
            let mut hex = String::with_capacity(128);
            for b in bytes {
                let _ = write!(hex, "{b:02x}");
            }
            s.serialize_str(&hex)
        }
        pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<[u8; 64], D::Error> {
            let s = String::deserialize(d)?;
            let bytes = (0..s.len())
                .step_by(2)
                .map(|i| u8::from_str_radix(&s[i..i + 2], 16).map_err(serde::de::Error::custom))
                .collect::<Result<Vec<u8>, _>>()?;
            bytes
                .try_into()
                .map_err(|_| serde::de::Error::custom("expected 64 bytes"))
        }
    }
}

fn hex_encode(bytes: &[u8]) -> String {
    use std::fmt::Write;
    let mut s = String::with_capacity(bytes.len() * 2);
    for b in bytes {
        let _ = write!(s, "{b:02x}");
    }
    s
}

// ============================================================================
// Types
// ============================================================================

/// A single step in a Merkle proof path.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MerkleSibling {
    #[serde(with = "hex_bytes::b32")]
    pub hash: [u8; 32],
    pub is_left: bool,
}

/// Merkle inclusion proof for a single chunk.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChunkProof {
    #[serde(with = "hex_bytes::b16")]
    pub chunk_id: [u8; 16],
    #[serde(with = "hex_bytes::b32")]
    pub chunk_text_hash: [u8; 32],
    pub index: usize,
    pub siblings: Vec<MerkleSibling>,
}

/// Human-readable source information for a proof receipt.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SourceRef {
    pub document_path: String,
    #[serde(with = "hex_bytes::b16")]
    pub chunk_id: [u8; 16],
    pub chunk_text: String,
    pub chunk_sequence: u32,
    pub relevance_score: f32,
}

/// A complete cryptographic proof receipt.
///
/// Contains everything needed to independently verify what context
/// an AI had access to when it made a decision or generated code.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProofReceipt {
    /// Schema version (starts at 1)
    pub version: u32,
    /// The original query text
    pub query: String,
    /// BLAKE3 hash of the query
    #[serde(with = "hex_bytes::b32")]
    pub query_hash: [u8; 32],
    /// ISO 8601 timestamp
    pub timestamp: String,
    /// BLAKE3 hash of the assembled context string
    #[serde(with = "hex_bytes::b32")]
    pub context_hash: [u8; 32],
    /// Merkle root of the entire graph state
    #[serde(with = "hex_bytes::b32")]
    pub state_root: [u8; 32],
    /// Merkle root of all chunk text hashes (binary tree)
    #[serde(with = "hex_bytes::b32")]
    pub chunk_tree_root: [u8; 32],
    /// Per-chunk Merkle inclusion proofs
    pub chunk_proofs: Vec<ChunkProof>,
    /// Human-readable source references
    pub sources: Vec<SourceRef>,
    /// Ed25519 signature over `signing_bytes()`
    #[serde(with = "hex_bytes::b64")]
    pub signature: [u8; 64],
    /// Signer's Ed25519 public key
    #[serde(with = "hex_bytes::b32")]
    pub signer_public_key: [u8; 32],
    /// Signer's device ID
    #[serde(with = "hex_bytes::b16")]
    pub device_id: [u8; 16],
}

impl ProofReceipt {
    /// Compute deterministic bytes for signing.
    ///
    /// Layout: version (4) || `query_hash` (32) || `context_hash` (32) || `state_root` (32) || `chunk_tree_root` (32)
    /// Total: 132 bytes
    pub fn signing_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(132);
        bytes.extend_from_slice(&self.version.to_le_bytes());
        bytes.extend_from_slice(&self.query_hash);
        bytes.extend_from_slice(&self.context_hash);
        bytes.extend_from_slice(&self.state_root);
        bytes.extend_from_slice(&self.chunk_tree_root);
        bytes
    }

    /// Verify the Ed25519 signature over `signing_bytes()`.
    pub fn verify_signature(&self) -> crate::Result<()> {
        let verifying_key = VerifyingKey::from_bytes(&self.signer_public_key)
            .map_err(|e| crate::CPError::Verification(format!("Invalid public key: {e}")))?;
        let signature = Signature::from_bytes(&self.signature);
        verifying_key
            .verify(&self.signing_bytes(), &signature)
            .map_err(|e| {
                crate::CPError::Verification(format!("Signature verification failed: {e}"))
            })
    }

    /// Verify all chunk Merkle proofs against `chunk_tree_root`.
    pub fn verify_chunk_proofs(&self) -> crate::Result<()> {
        for proof in &self.chunk_proofs {
            let siblings: Vec<([u8; 32], bool)> =
                proof.siblings.iter().map(|s| (s.hash, s.is_left)).collect();
            if !verify_merkle_proof(
                &proof.chunk_text_hash,
                proof.index,
                &siblings,
                &self.chunk_tree_root,
            ) {
                return Err(crate::CPError::Verification(format!(
                    "Chunk proof failed for chunk {}",
                    hex_encode(&proof.chunk_id)
                )));
            }
        }
        Ok(())
    }

    /// Verify that `context_hash` matches the BLAKE3 hash of the given context string.
    pub fn verify_context_hash(&self, context: &str) -> bool {
        let computed = *blake3::hash(context.as_bytes()).as_bytes();
        computed == self.context_hash
    }

    /// Run all verification checks: signature + chunk proofs.
    pub fn verify_all(&self) -> crate::Result<()> {
        self.verify_signature()?;
        self.verify_chunk_proofs()?;
        Ok(())
    }

    /// Get a short summary string for display.
    pub fn summary(&self) -> String {
        format!(
            "ProofReceipt v{}: {} chunks from {} sources, state_root={}, signed by {}",
            self.version,
            self.chunk_proofs.len(),
            self.sources.len(),
            &hex_encode(&self.state_root)[..8],
            &hex_encode(&self.device_id)[..8],
        )
    }
}

/// Helper to build a `ChunkProof` from sorted chunk hashes.
///
/// Given all chunk hashes (sorted by chunk ID) and the target chunk's index,
/// generates the Merkle inclusion proof.
pub fn build_chunk_proof(
    chunk_id: [u8; 16],
    chunk_text_hash: [u8; 32],
    index: usize,
    all_chunk_hashes: &[[u8; 32]],
) -> ChunkProof {
    let raw_siblings = generate_merkle_proof(all_chunk_hashes, index);
    let siblings = raw_siblings
        .into_iter()
        .map(|(hash, is_left)| MerkleSibling { hash, is_left })
        .collect();

    ChunkProof {
        chunk_id,
        chunk_text_hash,
        index,
        siblings,
    }
}

/// Compute the Merkle root of chunk text hashes (binary tree).
pub fn compute_chunk_tree_root(chunk_hashes: &[[u8; 32]]) -> [u8; 32] {
    compute_merkle_root(chunk_hashes)
}

// ============================================================================
// Session Proof — captures what the AI touched during a turn
// ============================================================================

/// A user prompt captured from the `UserPromptSubmit` hook.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserPrompt {
    pub timestamp: String,
    pub prompt: String,
}

/// A single tool call event in a session.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionEvent {
    pub timestamp: String,
    pub tool: String,
    /// "read", "write", "edit", "search", "bash", etc.
    pub action: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub file_path: Option<String>,
    /// What was sent to the tool (command, content written, edit diff, search pattern)
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub input_preview: Option<String>,
    /// What the tool returned (file content read, command output, search results)
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub output_preview: Option<String>,
}

/// A snapshot of a file the AI touched, with content hash.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FileSnapshot {
    pub path: String,
    /// BLAKE3 hash of file content at the time
    #[serde(with = "hex_bytes::b32")]
    pub content_hash: [u8; 32],
    /// "read" or "write"
    pub action: String,
    pub size_bytes: u64,
    /// First ~500 chars of the file for human readability
    pub snippet: String,
}

/// Semantic search result captured in a proof — what the substrate returned for a query.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SemanticSearchResult {
    /// The query that was run against the substrate
    pub query: String,
    /// Hash of the query
    #[serde(with = "hex_bytes::b32")]
    pub query_hash: [u8; 32],
    /// State root of the substrate at search time
    #[serde(with = "hex_bytes::b32")]
    pub state_root: [u8; 32],
    /// Substrate stats at search time
    pub substrate_docs: usize,
    pub substrate_chunks: usize,
    pub substrate_embeddings: usize,
    /// Top-k results with scores and content
    pub results: Vec<SemanticHit>,
}

/// A single hit from semantic search — a ranked chunk with its score and content.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SemanticHit {
    /// Source document path
    pub doc_path: String,
    /// Chunk sequence number within the document
    pub chunk_sequence: u32,
    /// Hybrid relevance score (0.0-1.0) — how relevant the substrate ranked this chunk
    pub relevance_score: f32,
    /// Full chunk text — exactly what the AI would see as context
    pub chunk_text: String,
    /// BLAKE3 hash of the chunk text
    #[serde(with = "hex_bytes::b32")]
    pub chunk_hash: [u8; 32],
}

/// A session proof captures everything the AI touched in a single turn.
///
/// Unlike a search `ProofReceipt` (which proves what search results the AI got),
/// a `SessionProof` proves which files the AI read and wrote during a response,
/// what the user asked, and what the semantic substrate contained.
///
/// This is the thing git CAN'T do — git tracks diffs, Canon tracks cognition:
/// what the AI's understanding of the codebase looked like at decision time.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionProof {
    pub version: u32,
    pub session_id: String,
    pub timestamp: String,
    /// What the user asked the AI
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub user_prompts: Vec<UserPrompt>,
    /// Semantic search results for each user prompt — the AI's "knowledge state"
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub semantic_context: Vec<SemanticSearchResult>,
    /// Tool call events in chronological order (with input/output previews)
    pub events: Vec<SessionEvent>,
    /// Files the AI read (with content hashes)
    pub files_read: Vec<FileSnapshot>,
    /// Files the AI wrote or edited (with content hashes)
    pub files_written: Vec<FileSnapshot>,
    /// Merkle root of all file hashes (reads + writes, sorted by path)
    #[serde(with = "hex_bytes::b32")]
    pub files_root: [u8; 32],
    /// Ed25519 signature
    #[serde(with = "hex_bytes::b64")]
    pub signature: [u8; 64],
    /// Signer's public key
    #[serde(with = "hex_bytes::b32")]
    pub signer_public_key: [u8; 32],
    /// Signer's device ID
    #[serde(with = "hex_bytes::b16")]
    pub device_id: [u8; 16],
}

impl SessionProof {
    /// Compute deterministic bytes for signing.
    ///
    /// Layout: version (4) || `session_id_hash` (32) || `prompts_hash` (32) || `semantic_hash` (32) || `files_root` (32)
    /// Total: 164 bytes
    pub fn signing_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(164);
        bytes.extend_from_slice(&self.version.to_le_bytes());
        bytes.extend_from_slice(blake3::hash(self.session_id.as_bytes()).as_bytes());
        // Hash all user prompts so they can't be tampered
        let mut prompts_hasher = blake3::Hasher::new();
        for p in &self.user_prompts {
            prompts_hasher.update(p.prompt.as_bytes());
            prompts_hasher.update(&[0xFF]);
        }
        bytes.extend_from_slice(prompts_hasher.finalize().as_bytes());
        // Hash all semantic search results — the AI's knowledge state can't be tampered
        let mut semantic_hasher = blake3::Hasher::new();
        for sc in &self.semantic_context {
            semantic_hasher.update(&sc.query_hash);
            semantic_hasher.update(&sc.state_root);
            for hit in &sc.results {
                semantic_hasher.update(&hit.chunk_hash);
                semantic_hasher.update(&hit.relevance_score.to_le_bytes());
            }
            semantic_hasher.update(&[0xFF]);
        }
        bytes.extend_from_slice(semantic_hasher.finalize().as_bytes());
        bytes.extend_from_slice(&self.files_root);
        bytes
    }

    /// Verify the Ed25519 signature.
    pub fn verify_signature(&self) -> crate::Result<()> {
        let verifying_key = VerifyingKey::from_bytes(&self.signer_public_key)
            .map_err(|e| crate::CPError::Verification(format!("Invalid public key: {e}")))?;
        let signature = Signature::from_bytes(&self.signature);
        verifying_key
            .verify(&self.signing_bytes(), &signature)
            .map_err(|e| {
                crate::CPError::Verification(format!("Signature verification failed: {e}"))
            })
    }

    /// Verify the files Merkle root matches the file snapshots.
    pub fn verify_files_root(&self) -> crate::Result<()> {
        let mut all_hashes: Vec<[u8; 32]> = self
            .files_read
            .iter()
            .chain(self.files_written.iter())
            .map(|f| f.content_hash)
            .collect();
        all_hashes.sort_unstable();

        let computed = if all_hashes.is_empty() {
            *blake3::hash(b"empty").as_bytes()
        } else {
            compute_merkle_root(&all_hashes)
        };

        if computed != self.files_root {
            return Err(crate::CPError::Verification(
                "Files root mismatch".to_string(),
            ));
        }
        Ok(())
    }

    /// Verify everything: signature + files root.
    pub fn verify_all(&self) -> crate::Result<()> {
        self.verify_signature()?;
        self.verify_files_root()?;
        Ok(())
    }
}

/// Compute the files Merkle root from file snapshots.
pub fn compute_files_root(files: &[FileSnapshot]) -> [u8; 32] {
    let mut hashes: Vec<[u8; 32]> = files.iter().map(|f| f.content_hash).collect();
    hashes.sort_unstable();
    if hashes.is_empty() {
        *blake3::hash(b"empty").as_bytes()
    } else {
        compute_merkle_root(&hashes)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_test_receipt() -> ProofReceipt {
        // Create a simple test receipt with known values
        let chunk_hashes: Vec<[u8; 32]> = vec![[1u8; 32], [2u8; 32], [3u8; 32], [4u8; 32]];
        let chunk_tree_root = compute_chunk_tree_root(&chunk_hashes);

        let query = "test query";
        let query_hash = *blake3::hash(query.as_bytes()).as_bytes();
        let context = "test context";
        let context_hash = *blake3::hash(context.as_bytes()).as_bytes();

        // Build chunk proofs for first two chunks
        let proof0 = build_chunk_proof([0u8; 16], [1u8; 32], 0, &chunk_hashes);
        let proof1 = build_chunk_proof([1u8; 16], [2u8; 32], 1, &chunk_hashes);

        // Generate Ed25519 key pair for signing
        let signing_key = ed25519_dalek::SigningKey::from_bytes(&[42u8; 32]);
        let public_key = signing_key.verifying_key().to_bytes();
        let device_id = {
            let h = blake3::hash(&public_key);
            let mut id = [0u8; 16];
            id.copy_from_slice(&h.as_bytes()[..16]);
            id
        };

        let mut receipt = ProofReceipt {
            version: 1,
            query: query.to_string(),
            query_hash,
            timestamp: "2026-02-19T12:00:00Z".to_string(),
            context_hash,
            state_root: [99u8; 32],
            chunk_tree_root,
            chunk_proofs: vec![proof0, proof1],
            sources: vec![SourceRef {
                document_path: "test.rs".to_string(),
                chunk_id: [0u8; 16],
                chunk_text: "test chunk text".to_string(),
                chunk_sequence: 0,
                relevance_score: 0.95,
            }],
            signature: [0u8; 64],
            signer_public_key: public_key,
            device_id,
        };

        // Sign the receipt
        use ed25519_dalek::Signer;
        let sig = signing_key.sign(&receipt.signing_bytes());
        receipt.signature = sig.to_bytes();

        receipt
    }

    #[test]
    fn test_signing_bytes_deterministic() {
        let receipt = make_test_receipt();
        let bytes1 = receipt.signing_bytes();
        let bytes2 = receipt.signing_bytes();
        assert_eq!(bytes1, bytes2);
        // version(4) + query_hash(32) + context_hash(32) + state_root(32) + chunk_tree_root(32) = 132
        assert_eq!(bytes1.len(), 132);
    }

    #[test]
    fn test_verify_signature_valid() {
        let receipt = make_test_receipt();
        assert!(receipt.verify_signature().is_ok());
    }

    #[test]
    fn test_verify_signature_tampered() {
        let mut receipt = make_test_receipt();
        // Tamper with the state root
        receipt.state_root[0] ^= 0xFF;
        assert!(receipt.verify_signature().is_err());
    }

    #[test]
    fn test_verify_chunk_proofs_valid() {
        let receipt = make_test_receipt();
        assert!(receipt.verify_chunk_proofs().is_ok());
    }

    #[test]
    fn test_verify_chunk_proofs_tampered() {
        let mut receipt = make_test_receipt();
        // Tamper with a chunk hash
        receipt.chunk_proofs[0].chunk_text_hash[0] ^= 0xFF;
        assert!(receipt.verify_chunk_proofs().is_err());
    }

    #[test]
    fn test_verify_context_hash() {
        let receipt = make_test_receipt();
        assert!(receipt.verify_context_hash("test context"));
        assert!(!receipt.verify_context_hash("wrong context"));
    }

    #[test]
    fn test_verify_all() {
        let receipt = make_test_receipt();
        assert!(receipt.verify_all().is_ok());
    }

    #[test]
    fn test_serde_roundtrip() {
        let receipt = make_test_receipt();
        let json = serde_json::to_string_pretty(&receipt).unwrap();
        let parsed: ProofReceipt = serde_json::from_str(&json).unwrap();

        assert_eq!(receipt.version, parsed.version);
        assert_eq!(receipt.query, parsed.query);
        assert_eq!(receipt.query_hash, parsed.query_hash);
        assert_eq!(receipt.state_root, parsed.state_root);
        assert_eq!(receipt.chunk_tree_root, parsed.chunk_tree_root);
        assert_eq!(receipt.signature, parsed.signature);
        assert_eq!(receipt.signer_public_key, parsed.signer_public_key);
        assert_eq!(receipt.device_id, parsed.device_id);

        // Verify the deserialized receipt
        assert!(parsed.verify_all().is_ok());
    }

    #[test]
    fn test_build_chunk_proof() {
        let hashes = [[1u8; 32], [2u8; 32], [3u8; 32], [4u8; 32]];
        let root = compute_chunk_tree_root(&hashes);

        for i in 0..4 {
            let proof = build_chunk_proof([i as u8; 16], hashes[i], i, &hashes);
            let siblings: Vec<([u8; 32], bool)> =
                proof.siblings.iter().map(|s| (s.hash, s.is_left)).collect();
            assert!(verify_merkle_proof(&hashes[i], i, &siblings, &root));
        }
    }

    #[test]
    fn test_summary() {
        let receipt = make_test_receipt();
        let summary = receipt.summary();
        assert!(summary.contains("v1"));
        assert!(summary.contains("2 chunks"));
    }
}