dyolo-kya 1.0.0

Know Your Agent (KYA): cryptographic chain-of-custody for recursive AI delegation with provable scope narrowing
use std::collections::HashSet;
use std::sync::{
    RwLock,
    atomic::{AtomicU64, Ordering},
};

use rand::{rngs::OsRng, RngCore};

// ── Nonce generation ──────────────────────────────────────────────────────────

/// Generate a cryptographically random 16-byte nonce using the OS entropy source.
///
/// The returned value is suitable for use as a [`DelegationCert`] nonce.
/// Each call produces an independent, uniformly distributed 128-bit value.
///
/// [`DelegationCert`]: crate::DelegationCert
#[must_use]
pub fn fresh_nonce() -> [u8; 16] {
    let mut nonce = [0u8; 16];
    OsRng.fill_bytes(&mut nonce);
    nonce
}

// ── Revocation ────────────────────────────────────────────────────────────────

/// Storage backend for revoked certificate fingerprints.
///
/// The default implementation ([`MemoryRevocationStore`]) is backed by an
/// in-process `HashSet`. Production deployments should implement this trait
/// over a persistent store (Redis, a signed CRL, or an on-chain accumulator)
/// so revocations survive process restarts.
pub trait RevocationStore: Send + Sync {
    /// Returns `Ok(true)` if `fingerprint` has been revoked, `Ok(false)` otherwise.
    ///
    /// Implementations must be constant-time with respect to the set contents
    /// to avoid timing side-channels. The default [`MemoryRevocationStore`]
    /// uses a Bloom filter fast-path before an exact-match shard lookup.
    fn is_revoked(&self, fingerprint: &[u8; 32]) -> Result<bool, String>;

    /// Record `fingerprint` as permanently revoked.
    ///
    /// After this call returns `Ok(())`, all subsequent calls to [`is_revoked`]
    /// for the same fingerprint must return `Ok(true)`. Revocations must be
    /// durable — implementations backed by volatile storage must flush to disk
    /// or a remote store before returning.
    ///
    /// [`is_revoked`]: RevocationStore::is_revoked
    fn revoke(&self, fingerprint: &[u8; 32]) -> Result<(), String>;
}

/// Storage backend for consumed nonces.
///
/// The default implementation ([`MemoryNonceStore`]) is backed by an in-process
/// append-only store and is suitable for development and single-process
/// deployments. Production systems that require replay protection across process
/// restarts or in distributed environments **must** implement this trait over a
/// persistent backend (e.g., Redis SETNX, a monotone database row, or an
/// append-only ledger).
///
/// Implementations must never evict or expire consumed nonces. The only safe
/// time to remove a nonce from the store is after every certificate that could
/// have produced it has also expired and been revoked — a condition that is
/// difficult to determine without additional context. The safest policy is
/// permanent retention.
pub trait NonceStore: Send + Sync {
    /// Returns `Ok(true)` if `nonce` has been consumed, `Ok(false)` otherwise.
    fn is_consumed(&self, nonce: &[u8; 16]) -> Result<bool, String>;

    /// Record `nonce` as permanently consumed.
    ///
    /// After this call returns `Ok(())`, all subsequent calls to [`is_consumed`]
    /// for the same nonce must return `Ok(true)`. Implementations must never
    /// evict consumed nonces, as doing so reopens a replay window.
    ///
    /// [`is_consumed`]: NonceStore::is_consumed
    fn mark_consumed(&self, nonce: &[u8; 16]) -> Result<(), String>;
}

// ── MemoryRevocationStore ─────────────────────────────────────────────────────

/// In-process revocation store backed by a 64-word atomic Bloom filter
/// for an O(1) lock-free fast path, with a 64-shard `HashSet` for exact-match
/// confirmation.
///
/// The Bloom filter eliminates read-lock contention for the overwhelmingly
/// common case (non-revoked certificates), scaling linearly with CPU cores.
pub struct MemoryRevocationStore {
    bloom:  Vec<AtomicU64>,
    shards: Vec<RwLock<HashSet<[u8; 32]>>>,
}

impl MemoryRevocationStore {
    /// Create a new empty revocation store.
    #[must_use]
    pub fn new() -> Self {
        let bloom  = (0..64).map(|_| AtomicU64::new(0)).collect();
        let shards = (0..64).map(|_| RwLock::new(HashSet::new())).collect();
        Self { bloom, shards }
    }

    #[inline(always)]
    fn bloom_indices(fp: &[u8; 32]) -> (usize, u64) {
        let mut entropy = 0u64;
        for chunk in fp.chunks_exact(8) {
            entropy ^= u64::from_le_bytes(chunk.try_into().unwrap());
        }
        let word = (entropy % 64) as usize;
        let bit  = 1u64 << (entropy.rotate_right(13) % 64);
        (word, bit)
    }

    #[inline(always)]
    fn shard_index(fp: &[u8; 32]) -> usize {
        (u64::from_le_bytes(fp[24..32].try_into().unwrap()) % 64) as usize
    }
}

impl Default for MemoryRevocationStore {
    fn default() -> Self { Self::new() }
}

impl RevocationStore for MemoryRevocationStore {
    fn is_revoked(&self, fingerprint: &[u8; 32]) -> Result<bool, String> {
        let (word, bit) = Self::bloom_indices(fingerprint);
        if (self.bloom[word].load(Ordering::Relaxed) & bit) == 0 {
            return Ok(false);
        }
        let shard = self.shards[Self::shard_index(fingerprint)]
            .read()
            .map_err(|_| "revocation shard lock poisoned".to_owned())?;
        Ok(shard.contains(fingerprint))
    }

    fn revoke(&self, fingerprint: &[u8; 32]) -> Result<(), String> {
        let (word, bit) = Self::bloom_indices(fingerprint);
        self.bloom[word].fetch_or(bit, Ordering::SeqCst);
        self.shards[Self::shard_index(fingerprint)]
            .write()
            .map_err(|_| "revocation shard lock poisoned".to_owned())?
            .insert(*fingerprint);
        Ok(())
    }
}

// ── MemoryNonceStore ──────────────────────────────────────────────────────────

/// In-process nonce store backed by a 1024-word atomic Bloom filter for an
/// O(1) lock-free fast path, with a 256-shard `HashSet` for exact-match
/// confirmation.
///
/// Consumed nonces are **never evicted**. This is correct by design: removing
/// a nonce before every certificate that could have produced it has expired
/// opens a replay window. Production deployments that need memory-bounded
/// nonce tracking should implement [`NonceStore`] over a persistent backend
/// with TTL expiry tied to the maximum certificate lifetime in their system.
pub struct MemoryNonceStore {
    bloom:  Vec<AtomicU64>,
    shards: Vec<RwLock<HashSet<[u8; 16]>>>,
}

impl MemoryNonceStore {
    const BLOOM_WORDS: usize = 1024;
    const SHARD_COUNT: usize = 256;

    /// Create a new empty nonce store.
    #[must_use]
    pub fn new() -> Self {
        let bloom  = (0..Self::BLOOM_WORDS).map(|_| AtomicU64::new(0)).collect();
        let shards = (0..Self::SHARD_COUNT).map(|_| RwLock::new(HashSet::new())).collect();
        Self { bloom, shards }
    }

    #[inline(always)]
    fn indices(nonce: &[u8; 16]) -> (usize, u64, usize) {
        let e1 = u64::from_le_bytes(nonce[0..8].try_into().unwrap());
        let e2 = u64::from_le_bytes(nonce[8..16].try_into().unwrap());
        let h  = e1.wrapping_mul(0x9E3779B185EBCA87).wrapping_add(e2.rotate_left(23));

        let word  = (h as usize) % Self::BLOOM_WORDS;
        let bit   = 1u64 << (h.rotate_right(11) % 64);
        let shard = (h.rotate_right(31) as usize) % Self::SHARD_COUNT;
        (word, bit, shard)
    }
}

impl Default for MemoryNonceStore {
    fn default() -> Self { Self::new() }
}

impl NonceStore for MemoryNonceStore {
    fn is_consumed(&self, nonce: &[u8; 16]) -> Result<bool, String> {
        let (word, bit, shard) = Self::indices(nonce);
        if (self.bloom[word].load(Ordering::Acquire) & bit) == 0 {
            return Ok(false);
        }
        let s = self.shards[shard]
            .read()
            .map_err(|_| "nonce shard lock poisoned".to_owned())?;
        Ok(s.contains(nonce))
    }

    fn mark_consumed(&self, nonce: &[u8; 16]) -> Result<(), String> {
        let (word, bit, shard) = Self::indices(nonce);
        self.bloom[word].fetch_or(bit, Ordering::Release);
        self.shards[shard]
            .write()
            .map_err(|_| "nonce shard lock poisoned".to_owned())?
            .insert(*nonce);
        Ok(())
    }
}