klieo-core 0.41.3

//! 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::Drop`]
//! 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,
}

impl OwnershipRegistry {
    /// Build a registry bound to `bucket` on `kv`.
    pub fn new(kv: Arc<dyn KvStore>, bucket: String) -> Self {
        Self { kv, bucket }
    }

    /// 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
    }

    /// 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,
        })
    }

    /// 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)
    }
}

/// 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());
    }
}