clawhdf5/crates/clawhdf5-agent/src/signing.rs

//! Ed25519-signed checkpoints.
//!
//! When a signing key is set ([`crate::HDF5Memory::set_signing_key`]), every
//! checkpoint writes a signed manifest of the store: a SHA-256 per memory
//! record rolled into a Merkle root, plus hashes of the store's settings, its
//! sessions and its knowledge graph. [`verify_store`] recomputes all of it from
//! the file and checks the signature against a public key the caller trusts,
//! so any change to the checkpointed file — a record's text or embedding, a
//! setting, a session, a graph edge, made through this crate or any other HDF5
//! tool — is detected, and the per-record hashes say which records changed.
//!
//! What it does not cover: saves still only in the WAL (made since the last
//! checkpoint). [`VerifyReport::wal_entries_unsigned`] counts them.
//!
//! The hashes cover exactly what the file persists, in the form the loader
//! returns it, so a store verifies after any number of reopen/checkpoint
//! cycles. Derived data (L2 norms, the vector index) is not covered; it is
//! recomputed from covered data.

use ed25519_dalek::{Signature, Signer, Verifier};
pub use ed25519_dalek::{SigningKey, VerifyingKey};
use sha2::{Digest, Sha256};

use crate::MemoryConfig;
use crate::cache::MemoryCache;
use crate::knowledge::KnowledgeCache;
use crate::session::SessionCache;
use crate::wal::WalMark;

/// Version of the manifest encoding; part of what is signed.
pub const MANIFEST_VERSION: i64 = 1;

type Hash = [u8; 32];

/// The hashes a signature covers.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Manifest {
    pub record_count: u64,
    /// Merkle root over the per-record hashes.
    pub records_root: Hash,
    /// Settings persisted in `/meta`, plus the checkpoint's WAL mark.
    pub settings: Hash,
    pub sessions: Hash,
    pub graph: Hash,
}

impl Manifest {
    /// The exact bytes that are signed.
    pub fn signed_bytes(&self) -> Vec<u8> {
        let mut m = Vec::with_capacity(160);
        m.extend_from_slice(b"clawhdf5-agent signed checkpoint\0");
        m.extend_from_slice(&MANIFEST_VERSION.to_le_bytes());
        m.extend_from_slice(&self.record_count.to_le_bytes());
        m.extend_from_slice(&self.records_root);
        m.extend_from_slice(&self.settings);
        m.extend_from_slice(&self.sessions);
        m.extend_from_slice(&self.graph);
        m
    }
}

/// A signature as stored in a checkpoint.
#[derive(Debug, Clone)]
pub struct StoredSignature {
    pub manifest: Manifest,
    pub record_hashes: Vec<Hash>,
    pub public_key: [u8; 32],
    pub signature: [u8; 64],
}

/// Build the manifest (and per-record hashes) for the state about to be
/// checkpointed, and sign it.
pub fn sign(
    key: &SigningKey,
    config: &MemoryConfig,
    cache: &MemoryCache,
    sessions: &SessionCache,
    knowledge: &KnowledgeCache,
    wal_applied: Option<WalMark>,
) -> StoredSignature {
    let (manifest, record_hashes) = manifest(config, cache, sessions, knowledge, wal_applied);
    let signature = key.sign(&manifest.signed_bytes()).to_bytes();
    StoredSignature {
        manifest,
        record_hashes,
        public_key: key.verifying_key().to_bytes(),
        signature,
    }
}

/// Compute the manifest of a store's state.
pub fn manifest(
    config: &MemoryConfig,
    cache: &MemoryCache,
    sessions: &SessionCache,
    knowledge: &KnowledgeCache,
    wal_applied: Option<WalMark>,
) -> (Manifest, Vec<Hash>) {
    let record_hashes: Vec<Hash> = (0..cache.len()).map(|i| record_hash(cache, i)).collect();
    let manifest = Manifest {
        record_count: cache.len() as u64,
        records_root: merkle_root(&record_hashes),
        settings: settings_hash(config, wal_applied),
        sessions: sessions_hash(sessions),
        graph: graph_hash(knowledge),
    };
    (manifest, record_hashes)
}

// ---------------------------------------------------------------------------
// Canonical encoding
// ---------------------------------------------------------------------------

/// A SHA-256 over length-prefixed fields, so no two different field lists
/// hash the same bytes.
struct Fields(Sha256);

impl Fields {
    fn new(domain: &str) -> Self {
        let mut h = Sha256::new();
        h.update((domain.len() as u64).to_le_bytes());
        h.update(domain.as_bytes());
        Self(h)
    }
    fn bytes(&mut self, b: &[u8]) -> &mut Self {
        self.0.update((b.len() as u64).to_le_bytes());
        self.0.update(b);
        self
    }
    /// Strings as the loader returns them: stored null-padded, so a trailing
    /// NUL cannot survive a round trip and must not be part of the hash.
    fn str(&mut self, s: &str) -> &mut Self {
        self.bytes(s.trim_end_matches('\0').as_bytes())
    }
    fn u64(&mut self, v: u64) -> &mut Self {
        self.0.update(v.to_le_bytes());
        self
    }
    fn f64(&mut self, v: f64) -> &mut Self {
        self.0.update(v.to_bits().to_le_bytes());
        self
    }
    fn f32(&mut self, v: f32) -> &mut Self {
        self.0.update(v.to_bits().to_le_bytes());
        self
    }
    fn finish(self) -> Hash {
        self.0.finalize().into()
    }
}

/// Everything persisted about record `i`, including its position. The
/// embedding is hashed as the cache holds it — for a `float16` store that is
/// the half-rounded value the file holds.
fn record_hash(cache: &MemoryCache, i: usize) -> Hash {
    let mut f = Fields::new("clawhdf5-agent/record");
    f.u64(i as u64).str(&cache.chunks[i]);
    let emb: Vec<u8> = cache.embeddings[i]
        .iter()
        .flat_map(|v| v.to_bits().to_le_bytes())
        .collect();
    f.bytes(&emb)
        .str(&cache.source_channels[i])
        .f64(cache.timestamps[i])
        .str(&cache.session_ids[i])
        .str(&cache.tags[i])
        .u64(u64::from(cache.tombstones[i]))
        .f32(cache.activation_weights[i]);
    f.finish()
}

/// Binary Merkle tree: leaves are the record hashes; a parent hashes its two
/// children with a node prefix; an odd node is carried up unchanged.
fn merkle_root(leaves: &[Hash]) -> Hash {
    if leaves.is_empty() {
        return Fields::new("clawhdf5-agent/merkle-empty").finish();
    }
    let mut level: Vec<Hash> = leaves.to_vec();
    while level.len() > 1 {
        level = level
            .chunks(2)
            .map(|pair| match pair {
                [l, r] => {
                    let mut h = Sha256::new();
                    h.update([1u8]);
                    h.update(l);
                    h.update(r);
                    h.finalize().into()
                }
                [only] => *only,
                _ => unreachable!(),
            })
            .collect();
    }
    level[0]
}

fn settings_hash(c: &MemoryConfig, wal_applied: Option<WalMark>) -> Hash {
    let mut f = Fields::new("clawhdf5-agent/settings");
    f.str(crate::schema::SCHEMA_VERSION)
        .str(&c.created_at)
        .str(&c.agent_id)
        .str(&c.embedder)
        .u64(c.embedding_dim as u64)
        .u64(c.chunk_size as u64)
        .u64(c.overlap as u64)
        .u64(u64::from(c.float16))
        .u64(u64::from(c.compression))
        .u64(u64::from(c.compression_level))
        .f32(c.compact_threshold)
        .f32(c.hebbian_boost)
        .f32(c.decay_factor)
        .u64(u64::from(c.wal_enabled))
        .u64(c.wal_max_entries as u64)
        .u64(u64::from(c.quantized_index))
        .u64(c.hnsw_m as u64)
        .u64(c.hnsw_ef_construction as u64)
        .u64(c.hnsw_ef_search as u64);
    // An empty mark is not written to the file, so it must hash as none.
    match wal_applied.filter(|m| m.len > 0) {
        Some(m) => f.u64(1).u64(m.len).u64(u64::from(m.crc)),
        None => f.u64(0),
    };
    f.finish()
}

fn sessions_hash(s: &SessionCache) -> Hash {
    let mut f = Fields::new("clawhdf5-agent/sessions");
    f.u64(s.entries.len() as u64);
    for (i, e) in s.entries.iter().enumerate() {
        f.str(&e.id)
            .u64(e.start_idx)
            .u64(e.end_idx)
            .str(&e.channel)
            .f64(e.ts)
            .str(s.summaries.get(i).map(String::as_str).unwrap_or(""));
    }
    f.finish()
}

fn graph_hash(k: &KnowledgeCache) -> Hash {
    let mut f = Fields::new("clawhdf5-agent/graph");
    f.u64(k.entities.len() as u64);
    for e in &k.entities {
        f.u64(e.id)
            .str(&e.name)
            .str(&e.entity_type)
            .u64(e.embedding_idx as u64);
    }
    f.u64(k.relations.len() as u64);
    for r in &k.relations {
        f.u64(r.src)
            .u64(r.tgt)
            .str(&r.relation)
            .f32(r.weight)
            .f64(r.ts);
    }
    f.u64(k.alias_strings.len() as u64);
    for (s, id) in k.alias_strings.iter().zip(&k.alias_entity_ids) {
        f.str(s).u64(*id as u64);
    }
    f.finish()
}

// ---------------------------------------------------------------------------
// Verification
// ---------------------------------------------------------------------------

/// The outcome of [`verify_store`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VerifyReport {
    /// The checkpoint carries a signature.
    pub signed: bool,
    /// The signature was made by the key the caller trusts.
    pub key_matches: bool,
    /// The signature over the stored manifest is valid.
    pub signature_valid: bool,
    /// The file's current contents match the signed manifest.
    pub records_match: bool,
    pub settings_match: bool,
    pub sessions_match: bool,
    pub graph_match: bool,
    /// Records whose contents differ from what was signed (by position),
    /// when the stored per-record hashes are themselves authentic.
    pub changed_records: Vec<usize>,
    /// Records in the file versus in the signed manifest.
    pub record_count: u64,
    pub signed_record_count: u64,
    /// The public key the checkpoint claims to be signed by.
    pub public_key: Option<[u8; 32]>,
    /// Saves in the WAL after the checkpoint: not covered by the signature.
    pub wal_entries_unsigned: usize,
}

impl VerifyReport {
    /// Signed by the trusted key, signature valid, and every part of the
    /// file unchanged since it was signed.
    pub fn is_valid(&self) -> bool {
        self.signed
            && self.key_matches
            && self.signature_valid
            && self.records_match
            && self.settings_match
            && self.sessions_match
            && self.graph_match
    }
}

/// Check a store file against the public key the caller trusts.
///
/// Reads the checkpoint (not the WAL), recomputes every hash from its
/// contents and checks the signature. Never writes.
pub fn verify_store(
    path: &std::path::Path,
    trusted: &VerifyingKey,
) -> Result<VerifyReport, crate::MemoryError> {
    let file = clawhdf5::File::open(path)
        .map_err(|e| crate::MemoryError::Hdf5(format!("cannot open {}: {e}", path.display())))?;
    let (config, cache, sessions, knowledge) = crate::schema::validate_and_load(&file)?;
    let checkpoint = crate::schema::read_checkpoint_meta(&file);
    let stored = crate::schema::read_signature(&file)?;
    let wal_entries_unsigned = count_wal_entries_after(path, checkpoint.wal_applied);

    let (current, current_hashes) = manifest(
        &config,
        &cache,
        &sessions,
        &knowledge,
        checkpoint.wal_applied,
    );

    let Some(stored) = stored else {
        return Ok(VerifyReport {
            signed: false,
            key_matches: false,
            signature_valid: false,
            records_match: false,
            settings_match: false,
            sessions_match: false,
            graph_match: false,
            changed_records: Vec::new(),
            record_count: current.record_count,
            signed_record_count: 0,
            public_key: None,
            wal_entries_unsigned,
        });
    };

    let key_matches = stored.public_key == trusted.to_bytes();
    let signature_valid = trusted
        .verify(
            &stored.manifest.signed_bytes(),
            &Signature::from_bytes(&stored.signature),
        )
        .is_ok();
    // The stored per-record hashes can localise a change only if they are
    // the ones that were signed.
    let hashes_authentic = signature_valid
        && stored.record_hashes.len() as u64 == stored.manifest.record_count
        && merkle_root(&stored.record_hashes) == stored.manifest.records_root;
    let changed_records = if hashes_authentic {
        let n = current_hashes.len().max(stored.record_hashes.len());
        (0..n)
            .filter(|&i| current_hashes.get(i) != stored.record_hashes.get(i))
            .collect()
    } else {
        Vec::new()
    };

    Ok(VerifyReport {
        signed: true,
        key_matches,
        signature_valid,
        records_match: signature_valid
            && current.record_count == stored.manifest.record_count
            && current.records_root == stored.manifest.records_root,
        settings_match: signature_valid && current.settings == stored.manifest.settings,
        sessions_match: signature_valid && current.sessions == stored.manifest.sessions,
        graph_match: signature_valid && current.graph == stored.manifest.graph,
        changed_records,
        record_count: current.record_count,
        signed_record_count: stored.manifest.record_count,
        public_key: Some(stored.public_key),
        wal_entries_unsigned,
    })
}

fn count_wal_entries_after(store: &std::path::Path, mark: Option<WalMark>) -> usize {
    let wal = store.with_extension("h5.wal");
    if !wal.exists() {
        return 0;
    }
    crate::wal::WalFile::read_entries_for_migration(&wal, mark)
        .map(|e| e.len())
        .unwrap_or(0)
}

/// A new random signing key from the operating system's RNG.
pub fn generate_key() -> SigningKey {
    SigningKey::generate(&mut rand_core::OsRng)
}

/// Hex encoding for keys and signatures in attributes and the CLI.
pub fn to_hex(bytes: &[u8]) -> String {
    bytes.iter().map(|b| format!("{b:02x}")).collect()
}

/// Parse hex into exactly `N` bytes.
pub fn from_hex<const N: usize>(s: &str) -> Option<[u8; N]> {
    let s = s.trim();
    if s.len() != 2 * N {
        return None;
    }
    let mut out = [0u8; N];
    for (i, byte) in out.iter_mut().enumerate() {
        *byte = u8::from_str_radix(&s[2 * i..2 * i + 2], 16).ok()?;
    }
    Some(out)
}