klieo-ops 0.41.2

//! `KvHandoff`: default `Handoff` impl persisting envelopes in `KvStore`.
//!
//! Envelopes are stored content-addressed under
//! `ops.handoff.envelopes/<content_hash>`. Delivery pointers land under
//! `ops.handoff.delivery/<agent_id>:<content_hash>`. Verification performs
//! registry trust check, expiry, and signature-math checks — in that order.

use super::envelope::{
    verify_signature, CanonicalPayload, EnvelopeSignature, HandoffEnvelope, HandoffProof,
    MerklePrefixProof,
};
use super::trait_::{Handoff, HandoffError, HandoffState};
use crate::audit::OpsAuditSink;
use crate::ops_event::OpsEvent;
use crate::signer::SignerProvider;
use crate::source_identity_registry::SourceIdentityRegistry;
use crate::types::AgentId;
use async_trait::async_trait;
use bytes::Bytes;
use chrono::Utc;
use klieo_core::ids::RunId;
use klieo_core::memory::EpisodicMemory;
use klieo_core::KvStore;
use std::sync::Arc;

const BUCKET_ENVELOPES: &str = "ops.handoff.envelopes";
const BUCKET_DELIVERY: &str = "ops.handoff.delivery";

/// KV-backed `Handoff` implementation.
pub struct KvHandoff {
    kv: Arc<dyn KvStore>,
    trust_registry: Arc<dyn SourceIdentityRegistry>,
    audit: Option<OpsAuditSink>,
}

impl KvHandoff {
    /// Construct a `KvHandoff` over the given `KvStore` and source-identity
    /// trust registry. `verify()` rejects any envelope whose `source_identity`
    /// is absent from the registry or whose embedded `verifying_key_hex`
    /// disagrees with the registered key — there is no fail-open path.
    #[must_use]
    pub fn new(kv: Arc<dyn KvStore>, trust_registry: Arc<dyn SourceIdentityRegistry>) -> Self {
        Self {
            kv,
            trust_registry,
            audit: None,
        }
    }

    /// Wire an audit sink. When set, every state-changing call emits the
    /// corresponding `OpsEvent` into `episodic` under `run_id`.
    #[must_use]
    pub fn with_audit(mut self, episodic: Arc<dyn EpisodicMemory>, run_id: RunId) -> Self {
        self.audit = Some(OpsAuditSink::new(
            episodic,
            run_id,
            "klieo.ops.handoff.audit",
        ));
        self
    }

    async fn emit(&self, event: OpsEvent) {
        if let Some(sink) = &self.audit {
            sink.emit(event).await;
        }
    }

    fn content_hash(env: &HandoffEnvelope) -> String {
        let canonical_bytes = CanonicalPayload::from_envelope(env).to_bytes();
        blake3::hash(&canonical_bytes).to_hex().to_string()
    }

    /// Resolve `envelope.source_identity` against the trust registry and
    /// confirm the embedded `verifying_key_hex` matches the registered key.
    ///
    /// Returns `Ok(())` only when both conditions hold. Any failure yields
    /// `HandoffError::UntrustedSource` — there is no fail-open path.
    fn check_source_identity(&self, envelope: &HandoffEnvelope) -> Result<(), HandoffError> {
        let id = &envelope.source_identity;
        let registered_vk = self.trust_registry.lookup(id).ok_or_else(|| {
            HandoffError::UntrustedSource(format!("source identity '{id}' not in registry"))
        })?;
        let envelope_vk_bytes = hex::decode(&envelope.source_signature.verifying_key_hex)
            .map_err(|e| HandoffError::UntrustedSource(format!("verifying_key_hex decode: {e}")))?;
        if registered_vk.to_bytes().as_slice() != envelope_vk_bytes.as_slice() {
            return Err(HandoffError::UntrustedSource(format!(
                "source identity '{id}' verifying-key mismatch"
            )));
        }
        Ok(())
    }
}

#[async_trait]
impl Handoff for KvHandoff {
    async fn package(
        &self,
        state: HandoffState,
        signer: &dyn SignerProvider,
        ttl: chrono::Duration,
    ) -> Result<HandoffEnvelope, HandoffError> {
        let now = Utc::now();
        let expires = now + ttl;
        let issued_at = now.to_rfc3339();
        let expires_at = expires.to_rfc3339();

        // Build a stub envelope to derive the canonical payload for signing.
        let stub = HandoffEnvelope {
            from_run: state.from_run.clone(),
            merkle_prefix_proof: MerklePrefixProof {
                root: state.merkle_root.clone(),
                at_seq: state.at_seq,
            },
            redacted_state: state.redacted_state.clone(),
            source_signature: EnvelopeSignature {
                signature_hex: String::new(),
                verifying_key_hex: String::new(),
            },
            source_identity: state.source_identity.clone(),
            tenant: state.tenant.clone(),
            issued_at: issued_at.clone(),
            expires_at: expires_at.clone(),
        };

        let canonical_bytes = CanonicalPayload::from_envelope(&stub).to_bytes();
        let signed = signer
            .sign_handoff(&canonical_bytes)
            .await
            .map_err(|e| HandoffError::Internal(format!("signer: {e}")))?;
        let source_signature = EnvelopeSignature {
            signature_hex: hex::encode(signed.signature),
            verifying_key_hex: hex::encode(signed.verifying_key),
        };

        let envelope = HandoffEnvelope {
            source_signature,
            ..stub
        };
        self.emit(OpsEvent::HandoffPackaged {
            tenant: envelope.tenant.clone(),
            from_run: envelope.from_run.clone(),
            merkle_root: envelope.merkle_prefix_proof.root.clone(),
        })
        .await;
        Ok(envelope)
    }

    async fn deliver(
        &self,
        envelope: HandoffEnvelope,
        to: AgentId,
    ) -> Result<String, HandoffError> {
        let body = serde_json::to_vec(&envelope)
            .map_err(|e| HandoffError::Serialisation(format!("envelope serialise: {e}")))?;
        let hash = Self::content_hash(&envelope);
        self.kv
            .put(BUCKET_ENVELOPES, &hash, Bytes::from(body))
            .await
            .map_err(|e| HandoffError::Storage(e.to_string()))?;
        let delivery_key = format!("{}:{}", to.0, hash);
        self.kv
            .put(
                BUCKET_DELIVERY,
                &delivery_key,
                Bytes::from(envelope.from_run.clone()),
            )
            .await
            .map_err(|e| HandoffError::Storage(e.to_string()))?;
        self.emit(OpsEvent::HandoffDelivered {
            tenant: envelope.tenant.clone(),
            from_run: envelope.from_run.clone(),
            to_agent: to.0.clone(),
        })
        .await;
        Ok(envelope.from_run)
    }

    /// Verify a [`HandoffEnvelope`] and return a [`HandoffProof`] on success.
    ///
    /// ## Verification order
    ///
    /// 1. **Expiry** — fast-reject on obviously stale envelopes.
    /// 2. **Signature math** — verifies the Ed25519 signature over the
    ///    canonical payload using the key embedded in the envelope itself.
    /// 3. **Identity trust** — confirms the embedded key matches the registry
    ///    entry for `source_identity`.
    ///
    /// ## Identity-oracle mitigation
    ///
    /// Returning distinct errors for `SignatureInvalid` vs `UntrustedSource`
    /// leaks registry membership: an attacker can probe arbitrary identities
    /// and distinguish "identity exists but sig is wrong" from "identity not
    /// in registry" even with signature-first ordering, because the signature
    /// math step uses the envelope's *self-reported* key, not the registry key.
    ///
    /// The distinct error variants are preserved for internal observability
    /// (audit logs, metrics, operator tooling). **Remote-facing services must
    /// call `verify_opaque` instead**, which collapses both `SignatureInvalid`
    /// and `UntrustedSource` into `HandoffError::Unverifiable { reason: "" }`
    /// so no identity-membership information escapes the network boundary.
    async fn verify(&self, envelope: &HandoffEnvelope) -> Result<HandoffProof, HandoffError> {
        // Order: expiry → signature math → identity.
        //
        // Signature-first-then-identity ensures an attacker probing arbitrary
        // source identities receives only `SignatureInvalid` (the same error
        // regardless of registry membership) rather than `UntrustedSource`
        // which would confirm whether an identity exists in the registry.
        let expires = chrono::DateTime::parse_from_rfc3339(&envelope.expires_at)
            .map_err(|e| HandoffError::Internal(format!("parse expires_at: {e}")))?;
        if Utc::now() > expires {
            let err = HandoffError::Expired {
                expired_at: envelope.expires_at.clone(),
            };
            self.emit(OpsEvent::HandoffVerificationFailed {
                tenant: envelope.tenant.clone(),
                from_run: Some(envelope.from_run.clone()),
                reason: err.to_string(),
            })
            .await;
            return Err(err);
        }
        if let Err(msg) = verify_signature(envelope) {
            self.emit(OpsEvent::HandoffVerificationFailed {
                tenant: envelope.tenant.clone(),
                from_run: Some(envelope.from_run.clone()),
                reason: msg.clone(),
            })
            .await;
            return Err(HandoffError::SignatureInvalid(msg));
        }
        if let Err(trust_err) = self.check_source_identity(envelope) {
            self.emit(OpsEvent::HandoffVerificationFailed {
                tenant: envelope.tenant.clone(),
                from_run: Some(envelope.from_run.clone()),
                reason: trust_err.to_string(),
            })
            .await;
            return Err(trust_err);
        }
        self.emit(OpsEvent::HandoffVerified {
            tenant: envelope.tenant.clone(),
            from_run: envelope.from_run.clone(),
        })
        .await;
        Ok(HandoffProof {
            from_run: envelope.from_run.clone(),
            source_identity: envelope.source_identity.clone(),
            tenant: envelope.tenant.clone(),
            redacted_state: envelope.redacted_state.clone(),
        })
    }
}