clawhdf5/crates/clawhdf5-agent/tests/signed_store.rs

//! Ed25519-signed checkpoints: `HDF5Memory::set_signing_key` and
//! `HDF5Memory::verify`.

use std::path::Path;

use clawhdf5_agent::signing::{SigningKey, VerifyReport, VerifyingKey};
use clawhdf5_agent::storage;
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, MemoryError, schema};
use tempfile::TempDir;

const DIM: usize = 16;

fn key(seed: u8) -> SigningKey {
    SigningKey::from_bytes(&[seed; 32])
}

fn entry(i: usize, chunk: &str) -> MemoryEntry {
    MemoryEntry {
        chunk: chunk.to_string(),
        embedding: (0..DIM)
            .map(|j| ((i * DIM + j) as f32 * 0.37).sin())
            .collect(),
        source_channel: "chat".into(),
        timestamp: 1_700_000_000.0 + i as f64,
        session_id: format!("s{}", i % 3),
        tags: format!("t{i}"),
    }
}

/// Awkward strings on purpose: they must hash the same after a round trip.
const TEXTS: [&str; 6] = [
    "plain text",
    "ünïcödé — 日本語 🙂",
    "",
    "trailing spaces   ",
    "tab\tand\nnewline",
    "x",
];

fn signed_store(dir: &TempDir, float16: bool, k: &SigningKey) -> std::path::PathBuf {
    let mut cfg = MemoryConfig::new(dir.path().join("s.h5"), "agent", DIM);
    cfg.float16 = float16;
    let path = cfg.path.clone();
    let mut m = HDF5Memory::create(cfg).unwrap();
    m.set_signing_key(k.clone());
    let entries = (0..30).map(|i| entry(i, TEXTS[i % TEXTS.len()])).collect();
    m.save_batch(entries).unwrap();
    // Some graph and a deleted record, so every part of the manifest is used.
    let a = m.knowledge_mut().add_entity("Alice", "person", 0);
    let b = m.knowledge_mut().add_entity("Acme", "org", -1);
    m.knowledge_mut().add_relation(a, b, "works_at", 0.75);
    m.sessions_mut()
        .add_at("s0", 0, 9, "chat", "first session", 1_700_000_000.0);
    m.delete(4).unwrap();
    m.flush_wal().unwrap();
    path
}

fn verify(path: &Path, k: &SigningKey) -> VerifyReport {
    HDF5Memory::verify(path, &k.verifying_key()).unwrap()
}

#[test]
fn a_signed_store_verifies_through_reopen_and_checkpoint_cycles() {
    for float16 in [true, false] {
        let dir = TempDir::new().unwrap();
        let k = key(7);
        let path = signed_store(&dir, float16, &k);
        let r = verify(&path, &k);
        assert!(r.is_valid(), "float16={float16}: {r:?}");
        assert_eq!(r.public_key, Some(k.verifying_key().to_bytes()));
        assert_eq!(r.record_count, 30);
        assert!(r.changed_records.is_empty());

        // Reopen, change nothing, checkpoint again (with the key): still valid.
        for _ in 0..3 {
            let mut m = HDF5Memory::open(&path).unwrap();
            assert!(m.is_signed());
            m.set_signing_key(k.clone());
            m.flush_wal().unwrap();
            drop(m);
            assert!(verify(&path, &k).is_valid());
        }
        // And after real changes, re-signed.
        let mut m = HDF5Memory::open(&path).unwrap();
        m.set_signing_key(k.clone());
        m.save(entry(99, "added later")).unwrap();
        m.hybrid_search(&entry(1, "").embedding, "text", 0.4, 0.6, 5);
        m.flush_wal().unwrap();
        drop(m);
        let r = verify(&path, &k);
        assert!(r.is_valid(), "{r:?}");
        assert_eq!(r.record_count, 31);
    }
}

#[test]
fn a_signed_store_refuses_to_checkpoint_without_its_key() {
    let dir = TempDir::new().unwrap();
    let k = key(1);
    let path = signed_store(&dir, true, &k);

    let mut m = HDF5Memory::open(&path).unwrap();
    m.save(entry(50, "pending")).unwrap();
    match m.flush_wal() {
        Err(MemoryError::SigningKeyRequired(msg)) => assert!(msg.contains("signed"), "{msg}"),
        other => panic!("expected SigningKeyRequired, got {other:?}"),
    }
    // The file is untouched and still valid; the save is still in the WAL.
    let r = verify(&path, &k);
    assert!(r.is_valid());
    assert_eq!(r.wal_entries_unsigned, 1);

    // Supplying the key lets the checkpoint through, signed.
    m.set_signing_key(k.clone());
    m.flush_wal().unwrap();
    drop(m);
    let r = verify(&path, &k);
    assert!(r.is_valid());
    assert_eq!((r.record_count, r.wal_entries_unsigned), (31, 0));

    // Removing the signature on purpose writes it unsigned.
    let mut m = HDF5Memory::open(&path).unwrap();
    m.remove_signature();
    m.flush_wal().unwrap();
    drop(m);
    let r = verify(&path, &k);
    assert!(!r.signed && !r.is_valid());
    assert!(!HDF5Memory::open(&path).unwrap().is_signed());
}

#[test]
fn the_wrong_key_does_not_verify_and_a_new_key_re_signs() {
    let dir = TempDir::new().unwrap();
    let (a, b) = (key(1), key(2));
    let path = signed_store(&dir, true, &a);
    let r = verify(&path, &b);
    assert!(r.signed && !r.key_matches && !r.signature_valid && !r.is_valid());

    let mut m = HDF5Memory::open(&path).unwrap();
    m.set_signing_key(b.clone());
    m.flush_wal().unwrap();
    drop(m);
    assert!(verify(&path, &b).is_valid());
    assert!(!verify(&path, &a).is_valid());
}

/// Rewrite the store with changed contents but the *old* signature — what
/// someone with write access to the file, but not the key, can do.
fn tamper(path: &Path, change: impl FnOnce(&mut Tampered)) {
    let file = clawhdf5::File::open(path).unwrap();
    let (config, cache, sessions, knowledge) = schema::validate_and_load(&file).unwrap();
    let checkpoint = schema::read_checkpoint_meta(&file);
    let signature = schema::read_signature(&file).unwrap().unwrap();
    drop(file);
    let mut t = Tampered {
        config,
        cache,
        sessions,
        knowledge,
    };
    change(&mut t);
    storage::write_to_disk_signed(
        path,
        &t.config,
        &t.cache,
        &t.sessions,
        &t.knowledge,
        &checkpoint,
        Some(&signature),
    )
    .unwrap();
}

struct Tampered {
    config: MemoryConfig,
    cache: clawhdf5_agent::cache::MemoryCache,
    sessions: clawhdf5_agent::SessionCache,
    knowledge: clawhdf5_agent::knowledge::KnowledgeCache,
}

#[test]
fn every_kind_of_edit_is_detected_and_located() {
    let k = key(3);
    type Edit = Box<dyn FnOnce(&mut Tampered)>;
    type Case = (&'static str, Edit, fn(&VerifyReport) -> bool);
    let cases: Vec<Case> = vec![
        (
            "record text",
            Box::new(|t: &mut Tampered| t.cache.chunks[7] = "rewritten".into()),
            |r| !r.records_match && r.changed_records == vec![7],
        ),
        (
            "one embedding value",
            Box::new(|t: &mut Tampered| {
                let mut e = t.cache.embeddings[12].to_vec();
                e[3] = 0.5;
                t.cache.embeddings.set(12, &e);
            }),
            |r| r.changed_records == vec![12],
        ),
        (
            "undelete",
            Box::new(|t: &mut Tampered| t.cache.tombstones[4] = 0),
            |r| r.changed_records == vec![4],
        ),
        (
            "timestamp",
            Box::new(|t: &mut Tampered| t.cache.timestamps[20] += 1.0),
            |r| r.changed_records == vec![20],
        ),
        (
            "record appended",
            Box::new(|t: &mut Tampered| {
                t.cache.push(
                    "new".into(),
                    vec![0.1; DIM],
                    "x".into(),
                    1.0,
                    "s".into(),
                    "".into(),
                );
            }),
            |r| !r.records_match && r.changed_records == vec![30] && r.record_count == 31,
        ),
        (
            "setting",
            Box::new(|t: &mut Tampered| t.config.agent_id = "someone-else".into()),
            |r| !r.settings_match && r.records_match,
        ),
        (
            "session summary",
            Box::new(|t: &mut Tampered| t.sessions.summaries[0] = "edited".into()),
            |r| !r.sessions_match && r.records_match,
        ),
        (
            "graph edge",
            Box::new(|t: &mut Tampered| t.knowledge.relations[0].weight = 1.0),
            |r| !r.graph_match && r.records_match,
        ),
    ];
    for (name, edit, check) in cases {
        let dir = TempDir::new().unwrap();
        let path = signed_store(&dir, true, &k);
        tamper(&path, edit);
        let r = verify(&path, &k);
        assert!(
            r.signed && r.key_matches && r.signature_valid,
            "{name}: {r:?}"
        );
        assert!(!r.is_valid(), "{name}: edit not detected: {r:?}");
        assert!(check(&r), "{name}: {r:?}");
    }
}

#[test]
fn a_forged_manifest_fails_the_signature() {
    // Recomputing the hashes for tampered contents does not help without the
    // key: the signature no longer matches the manifest.
    let dir = TempDir::new().unwrap();
    let k = key(5);
    let path = signed_store(&dir, true, &k);
    let file = clawhdf5::File::open(&path).unwrap();
    let (config, mut cache, sessions, knowledge) = schema::validate_and_load(&file).unwrap();
    let checkpoint = schema::read_checkpoint_meta(&file);
    let mut sig = schema::read_signature(&file).unwrap().unwrap();
    drop(file);
    cache.chunks[0] = "forged".into();
    // Re-sign with an attacker key, then splice the victim's public key back.
    let forged = clawhdf5_agent::signing::sign(
        &key(66),
        &config,
        &cache,
        &sessions,
        &knowledge,
        checkpoint.wal_applied,
    );
    sig.manifest = forged.manifest;
    sig.record_hashes = forged.record_hashes;
    storage::write_to_disk_signed(
        &path,
        &config,
        &cache,
        &sessions,
        &knowledge,
        &checkpoint,
        Some(&sig),
    )
    .unwrap();
    let r = verify(&path, &k);
    assert!(
        r.key_matches && !r.signature_valid && !r.is_valid(),
        "{r:?}"
    );
}

#[test]
fn an_unsigned_store_reports_unsigned() {
    let dir = TempDir::new().unwrap();
    let mut m = HDF5Memory::create(MemoryConfig::new(dir.path().join("u.h5"), "a", DIM)).unwrap();
    m.save_batch(vec![entry(0, "hello")]).unwrap();
    drop(m);
    let r = HDF5Memory::verify(&dir.path().join("u.h5"), &VerifyingKey::from(&key(1))).unwrap();
    assert!(!r.signed && !r.is_valid());
    assert_eq!(r.record_count, 1);
}

#[test]
fn nul_bytes_in_text_still_verify() {
    // Strings are stored null-padded; the hash must follow what a reopened
    // store actually holds, or an untouched store would fail to verify.
    let dir = TempDir::new().unwrap();
    let k = key(9);
    let mut m = HDF5Memory::create(MemoryConfig::new(dir.path().join("n.h5"), "a", DIM)).unwrap();
    m.set_signing_key(k.clone());
    m.save_batch(vec![
        entry(0, "inner\0nul"),
        entry(1, "trailing nul\0"),
        entry(2, "\0leading"),
    ])
    .unwrap();
    drop(m);
    let r = verify(&dir.path().join("n.h5"), &k);
    assert!(r.is_valid(), "{r:?}");
    let m = HDF5Memory::open(&dir.path().join("n.h5")).unwrap();
    eprintln!(
        "reloaded: {:?}",
        (0..3).map(|i| m.get_chunk(i)).collect::<Vec<_>>()
    );
}