klieo-core 2.2.0

//! Per-stream ownership claim backed by a dedicated KV bucket.
//!
//! Cluster 0.22 binds each A2A/MCP stream to the authenticated
//! principal (typically OAuth `sub`) at invoke time and rejects
//! mismatched resume attempts as `stream not found` (deny-as-
//! NotFound; no existence leak per OWASP IDOR best practice).
//! See ADR-022 for the full design rationale + threat model.
//!
//! Mirrors the cluster-0.20 [`crate::LeaderRegistry`] shape but
//! ownership is durable per-invoke, not liveness-bounded — no
//! heartbeat task, no TTL on the entry. Lifecycle: claim on
//! invoke start; drop on stream end via [`OwnershipHandle`]'s
//! `Drop` impl which spawns a best-effort `kv.delete`.

use crate::bus::KvStore;
use crate::error::BusError;
use bytes::Bytes;
use std::sync::Arc;

/// Per-replica ownership registry. Stores `{stream_id} ->
/// principal` mappings in a dedicated KV bucket (typically
/// `klieo-tenants`).
#[derive(Clone)]
pub struct OwnershipRegistry {
    kv: Arc<dyn KvStore>,
    bucket: String,
    strict: bool,
}

/// Verdict of [`OwnershipRegistry::check_owner`] for an enforcement gate, with
/// the fail-open/closed policy already applied so every transport resolves a
/// store error the same way.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum OwnershipCheck {
    /// The caller owns the key, or it is unclaimed (legacy / no authenticator)
    /// — proceed.
    Allowed,
    /// A different principal owns the key — deny (surface as leak-safe
    /// not-found, per the IDOR posture).
    Denied,
    /// The store was unreachable and this registry is **strict**, so the gate
    /// fails closed rather than risk a cross-tenant resume on a KV blip. The
    /// transport surfaces a retryable "unavailable" error, never an allow. A
    /// lenient registry never returns this — it fails open to `Allowed`.
    Unavailable,
}

/// Outcome of [`OwnershipRegistry::claim_guarded`] at a stream-open gate, with
/// the fail-open/closed policy already applied so neither transport re-derives
/// it. (No `Debug` — `OwnershipHandle` wraps a non-`Debug` `Arc<dyn KvStore>`.)
#[non_exhaustive]
pub enum OwnershipClaim {
    /// Ownership was recorded; hold the handle for the stream's lifetime so its
    /// `Drop` removes the entry on stream end.
    Claimed(OwnershipHandle),
    /// The store rejected the write but this registry is lenient, so the stream
    /// proceeds unbound (fail-open, ADR-022 partial-deployment posture).
    Proceed,
    /// The store rejected the write and this registry is **strict**, so the
    /// invoke must be denied rather than started unprotected. (A lenient
    /// registry never returns this.)
    Unavailable,
}

impl OwnershipRegistry {
    /// Build a **lenient** (fail-open) registry bound to `bucket` on `kv`: a
    /// store error during a check proceeds (ADR-022 partial-deployment posture).
    pub fn new(kv: Arc<dyn KvStore>, bucket: String) -> Self {
        Self {
            kv,
            bucket,
            strict: false,
        }
    }

    /// Build a **strict** (fail-closed) registry: a store error during a check
    /// denies the operation as `Unavailable` rather than proceed, closing the
    /// transient-KV-blip cross-tenant-resume window for regulated multi-tenant
    /// deployments. An unclaimed key still proceeds — that is governed by
    /// cluster membership, not store reachability (out of scope here).
    pub fn new_strict(kv: Arc<dyn KvStore>, bucket: String) -> Self {
        Self {
            kv,
            bucket,
            strict: true,
        }
    }

    /// Borrow the backing [`KvStore`]. Cluster 0.25 reaper wiring
    /// reads ownership entries from the same bucket this registry
    /// writes to.
    pub fn kv(&self) -> &Arc<dyn KvStore> {
        &self.kv
    }

    /// Name of the KV bucket this registry writes ownership entries
    /// to (typically `klieo-tenants`).
    pub fn bucket(&self) -> &str {
        &self.bucket
    }

    /// Whether this registry fails closed on store error (strict) rather than
    /// proceeding (lenient). Lets a transport assert the posture a
    /// [`DeploymentProfile`](crate::DeploymentProfile) forced.
    pub fn is_strict(&self) -> bool {
        self.strict
    }

    /// Claim ownership of `key` for `principal`. Writes the
    /// entry once; no heartbeat (unlike `LeaderRegistry`).
    /// Returns a guard whose `Drop` triggers a best-effort
    /// `kv.delete` on stream end.
    #[tracing::instrument(
        skip_all,
        fields(
            klieo.tenants.key = %key,
            klieo.principal_hash = %crate::principal_hash(&principal),
            db.system = "klieo-kv",
            db.namespace = %self.bucket,
            db.operation = "put",
        ),
        err,
    )]
    pub async fn claim(&self, key: String, principal: String) -> Result<OwnershipHandle, BusError> {
        self.kv
            .put(&self.bucket, &key, Bytes::from(principal))
            .await?;
        Ok(OwnershipHandle {
            kv: self.kv.clone(),
            bucket: self.bucket.clone(),
            key,
        })
    }

    /// Claim ownership applying this registry's fail-open/closed policy on store
    /// error, so the stream-open gate in each transport never re-derives it.
    /// Success → [`OwnershipClaim::Claimed`]; a store error → [`OwnershipClaim::Proceed`]
    /// when lenient (start unbound) or [`OwnershipClaim::Unavailable`] when strict
    /// (deny). The underlying error is logged by [`claim`](Self::claim)'s
    /// instrumentation; this records only the resulting decision.
    pub async fn claim_guarded(&self, key: String, principal: String) -> OwnershipClaim {
        match self.claim(key, principal).await {
            Ok(handle) => OwnershipClaim::Claimed(handle),
            Err(_) if self.strict => {
                tracing::warn!(
                    target: "klieo.tenants",
                    "ownership claim failed; denying invoke (fail-closed, strict tenant binding)"
                );
                OwnershipClaim::Unavailable
            }
            Err(_) => {
                tracing::warn!(
                    target: "klieo.tenants",
                    "ownership claim failed; proceeding unbound (fail-open per ADR-022)"
                );
                OwnershipClaim::Proceed
            }
        }
    }

    /// Look up the owner of `key`. `Ok(None)` means no
    /// ownership entry — legacy pre-0.22 stream OR no
    /// authenticator wired at invoke time. Callers fail open
    /// on `None` (proceed) per ADR-022 partial-deployment
    /// posture.
    #[tracing::instrument(
        skip_all,
        fields(
            klieo.tenants.key = %key,
            klieo.tenants.found = tracing::field::Empty,
            db.system = "klieo-kv",
            db.namespace = %self.bucket,
            db.operation = "get",
        ),
        err,
        level = "debug",
    )]
    pub async fn lookup(&self, key: &str) -> Result<Option<String>, BusError> {
        let result = match self.kv.get(&self.bucket, key).await? {
            Some(entry) => {
                let s = String::from_utf8(entry.value.to_vec())
                    .map_err(|e| BusError::Permanent(format!("non-utf8 owner: {e}")))?;
                Some(s)
            }
            None => None,
        };
        tracing::Span::current().record("klieo.tenants.found", result.is_some());
        Ok(result)
    }

    /// Authorize `principal` against the owner recorded for `key`, applying this
    /// registry's fail-open/closed policy on store error so the decision lives
    /// in one place rather than being re-derived per transport. Owner match or
    /// an unclaimed key → [`OwnershipCheck::Allowed`]; a different owner →
    /// [`OwnershipCheck::Denied`]; a store error → [`OwnershipCheck::Unavailable`]
    /// when strict, else `Allowed` (fail-open). The store error is logged here
    /// either way, so the transport never has to.
    pub async fn check_owner(&self, key: &str, principal: &str) -> OwnershipCheck {
        match self.lookup(key).await {
            Ok(Some(owner)) if owner == principal => OwnershipCheck::Allowed,
            Ok(Some(_)) => OwnershipCheck::Denied,
            Ok(None) => OwnershipCheck::Allowed,
            Err(e) if self.strict => {
                tracing::warn!(
                    target: "klieo.tenants",
                    key = %key,
                    error = %e,
                    "ownership lookup failed; denying (fail-closed, strict tenant binding)"
                );
                OwnershipCheck::Unavailable
            }
            Err(e) => {
                tracing::warn!(
                    target: "klieo.tenants",
                    key = %key,
                    error = %e,
                    "ownership lookup failed; proceeding (fail-open per ADR-022)"
                );
                OwnershipCheck::Allowed
            }
        }
    }
}

/// Drop-guard for an ownership claim. Spawns a best-effort
/// `kv.delete` on Drop; runtime guard mirrors
/// [`crate::LeaderHandle`] from cluster 0.20 — degrades to
/// warn-log when no tokio runtime is active during Drop
/// (synchronous shutdown paths).
pub struct OwnershipHandle {
    kv: Arc<dyn KvStore>,
    bucket: String,
    key: String,
}

impl Drop for OwnershipHandle {
    fn drop(&mut self) {
        let kv = self.kv.clone();
        let bucket = std::mem::take(&mut self.bucket);
        let key = std::mem::take(&mut self.key);
        match tokio::runtime::Handle::try_current() {
            Ok(handle) => {
                handle.spawn(async move {
                    if let Err(e) = kv.delete(&bucket, &key).await {
                        tracing::warn!(
                            target: "klieo.tenants",
                            bucket = %bucket,
                            key = %key,
                            error = %e,
                            "ownership delete failed; entry will linger until bucket TTL",
                        );
                    }
                });
            }
            Err(_) => {
                tracing::warn!(
                    target: "klieo.tenants",
                    bucket = %bucket,
                    key = %key,
                    "ownership delete skipped: no tokio runtime active during Drop",
                );
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::fake_kv;
    use std::time::Duration;

    const BUCKET: &str = "klieo-tenants";

    #[tokio::test]
    async fn claim_writes_entry_and_lookup_returns_principal() {
        let kv = fake_kv();
        let reg = OwnershipRegistry::new(kv.clone(), BUCKET.into());
        let _handle = reg.claim("a2a.t-1".into(), "alice".into()).await.unwrap();
        assert_eq!(
            reg.lookup("a2a.t-1").await.unwrap().as_deref(),
            Some("alice")
        );
    }

    #[tokio::test]
    async fn drop_handle_deletes_entry() {
        let kv = fake_kv();
        let reg = OwnershipRegistry::new(kv.clone(), BUCKET.into());
        {
            let _handle = reg.claim("a2a.t-2".into(), "bob".into()).await.unwrap();
            assert!(reg.lookup("a2a.t-2").await.unwrap().is_some());
        }
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert!(reg.lookup("a2a.t-2").await.unwrap().is_none());
    }

    #[tokio::test]
    async fn lookup_returns_none_for_unknown_key() {
        let kv = fake_kv();
        let reg = OwnershipRegistry::new(kv, BUCKET.into());
        assert!(reg.lookup("never").await.unwrap().is_none());
    }

    /// Every read fails, to drive the store-error arm of `check_owner`.
    struct GetFailsKv;

    #[async_trait::async_trait]
    impl KvStore for GetFailsKv {
        async fn get(&self, _: &str, _: &str) -> Result<Option<crate::bus::KvEntry>, BusError> {
            Err(BusError::Connection("kv down".into()))
        }
        async fn put(&self, _: &str, _: &str, _: Bytes) -> Result<crate::bus::Revision, BusError> {
            Err(BusError::Connection("kv down".into()))
        }
        async fn cas(
            &self,
            _: &str,
            _: &str,
            _: Bytes,
            _: Option<crate::bus::Revision>,
        ) -> Result<crate::bus::Revision, BusError> {
            Err(BusError::Connection("kv down".into()))
        }
        async fn delete(&self, _: &str, _: &str) -> Result<(), BusError> {
            Err(BusError::Connection("kv down".into()))
        }
        async fn lease(
            &self,
            _: &str,
            _: &str,
            _: Duration,
        ) -> Result<crate::bus::Lease, BusError> {
            Err(BusError::Connection("kv down".into()))
        }
        async fn keys(&self, _: &str) -> Result<Vec<String>, BusError> {
            Err(BusError::Connection("kv down".into()))
        }
    }

    #[tokio::test]
    async fn check_owner_allows_matching_principal_and_denies_other() {
        let kv = fake_kv();
        let reg = OwnershipRegistry::new(kv, BUCKET.into());
        let _h = reg.claim("a2a.t-c".into(), "alice".into()).await.unwrap();
        assert_eq!(
            reg.check_owner("a2a.t-c", "alice").await,
            OwnershipCheck::Allowed
        );
        assert_eq!(
            reg.check_owner("a2a.t-c", "bob").await,
            OwnershipCheck::Denied
        );
    }

    #[tokio::test]
    async fn check_owner_allows_unclaimed_key() {
        let kv = fake_kv();
        let reg = OwnershipRegistry::new(kv, BUCKET.into());
        assert_eq!(
            reg.check_owner("a2a.legacy", "alice").await,
            OwnershipCheck::Allowed
        );
    }

    #[tokio::test]
    async fn check_owner_lenient_proceeds_on_store_error() {
        let reg = OwnershipRegistry::new(Arc::new(GetFailsKv), BUCKET.into());
        assert_eq!(
            reg.check_owner("a2a.t", "alice").await,
            OwnershipCheck::Allowed,
            "a lenient registry fails open on a KV error (ADR-022)"
        );
    }

    #[tokio::test]
    async fn claim_guarded_returns_handle_then_policy_on_store_error() {
        let ok = OwnershipRegistry::new(fake_kv(), BUCKET.into());
        assert!(matches!(
            ok.claim_guarded("a2a.t".into(), "alice".into()).await,
            OwnershipClaim::Claimed(_)
        ));

        let lenient = OwnershipRegistry::new(Arc::new(GetFailsKv), BUCKET.into());
        assert!(
            matches!(
                lenient.claim_guarded("a2a.t".into(), "alice".into()).await,
                OwnershipClaim::Proceed
            ),
            "a lenient registry proceeds unbound when the claim write fails"
        );

        let strict = OwnershipRegistry::new_strict(Arc::new(GetFailsKv), BUCKET.into());
        assert!(
            matches!(
                strict.claim_guarded("a2a.t".into(), "alice".into()).await,
                OwnershipClaim::Unavailable
            ),
            "a strict registry denies (Unavailable) when the claim write fails"
        );
    }

    #[tokio::test]
    async fn check_owner_strict_denies_on_store_error() {
        let reg = OwnershipRegistry::new_strict(Arc::new(GetFailsKv), BUCKET.into());
        assert_eq!(
            reg.check_owner("a2a.t", "alice").await,
            OwnershipCheck::Unavailable,
            "a strict registry fails closed on a KV error rather than risk a cross-tenant resume"
        );
    }

    #[test]
    fn is_strict_reflects_constructor_choice() {
        let kv = fake_kv();
        assert!(
            !OwnershipRegistry::new(kv.clone(), BUCKET.into()).is_strict(),
            "new() builds a lenient registry"
        );
        assert!(
            OwnershipRegistry::new_strict(kv, BUCKET.into()).is_strict(),
            "new_strict() builds a strict registry"
        );
    }
}