klieo-ops 0.41.2

//! Cedar-backed policy gate with optional per-tenant bundles.
//!
//! `PolicyGate` holds one `PolicySet` per tenant plus an optional default
//! bundle (keyed by `None`). At evaluate time:
//!   1. Read `current_tenant()` task-local.
//!   2. Look up the matching bundle. If absent, fall back to the `None`-keyed
//!      default.
//!   3. If neither exists: return `GateDecision::Deny { code:
//!      "policy.no_bundle_for_tenant" }` (fail-CLOSED).
//!
//! JSON args are lifted into the Cedar context as follows: `bool` → Cedar
//! bool; `i64`-compatible numbers → Cedar long; floats → Cedar decimal
//! (4-fractional-digit precision); strings → Cedar string; arrays → Cedar
//! set (recursively lifted); objects → Cedar record (recursively lifted);
//! `null` and out-of-range `u64` → unrepresentable, treated as
//! `GateDecision::Deny { code: "cedar.context.unrepresentable" }`.

use super::trait_::{Gate, GateDecision, GateRequest};
use crate::tenant::current_tenant;
use crate::types::TenantId;
use async_trait::async_trait;
use cedar_policy::{Authorizer, Context, Entities, EntityUid, PolicySet, Request};
use std::collections::HashMap;
use std::str::FromStr;
use thiserror::Error;

/// Errors raised by `PolicyGate` construction.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum PolicyGateError {
    /// Cedar policy source did not parse.
    #[error("cedar policy parse for {tenant_label}: {message}")]
    Parse {
        /// Human label of the offending tenant slot ("default" or the TenantId).
        tenant_label: String,
        /// Underlying parse error message.
        message: String,
    },
}

/// Cedar policy bundle gate. Holds one `PolicySet` per tenant plus an
/// optional default bundle (keyed by `None`).
pub struct PolicyGate {
    bundles: HashMap<Option<TenantId>, PolicySet>,
    authorizer: Authorizer,
}

impl PolicyGate {
    /// Construct from a single Cedar source — installs as the default
    /// (no-tenant-override) bundle. Backward-compat constructor.
    pub fn from_cedar_source(src: &str) -> Result<Self, PolicyGateError> {
        let policies = PolicySet::from_str(src).map_err(|e| PolicyGateError::Parse {
            tenant_label: "default".into(),
            message: e.to_string(),
        })?;
        let mut bundles = HashMap::new();
        bundles.insert(None, policies);
        Ok(Self {
            bundles,
            authorizer: Authorizer::new(),
        })
    }

    /// Construct from a per-tenant bundle map. No default fallback.
    pub fn per_tenant(per_tenant: HashMap<TenantId, String>) -> Result<Self, PolicyGateError> {
        let mut bundles = HashMap::with_capacity(per_tenant.len());
        for (tenant, src) in per_tenant {
            let policies = PolicySet::from_str(&src).map_err(|e| PolicyGateError::Parse {
                tenant_label: tenant.0.clone(),
                message: e.to_string(),
            })?;
            bundles.insert(Some(tenant), policies);
        }
        Ok(Self {
            bundles,
            authorizer: Authorizer::new(),
        })
    }

    /// Construct with both a default fallback and per-tenant overrides.
    pub fn with_default_and_per_tenant(
        default_src: &str,
        per_tenant: HashMap<TenantId, String>,
    ) -> Result<Self, PolicyGateError> {
        let mut gate = Self::per_tenant(per_tenant)?;
        let default_policies =
            PolicySet::from_str(default_src).map_err(|e| PolicyGateError::Parse {
                tenant_label: "default".into(),
                message: e.to_string(),
            })?;
        gate.bundles.insert(None, default_policies);
        Ok(gate)
    }

    /// Builder method: add or replace a tenant-specific bundle.
    pub fn with_tenant_bundle(
        mut self,
        tenant: TenantId,
        src: &str,
    ) -> Result<Self, PolicyGateError> {
        let policies = PolicySet::from_str(src).map_err(|e| PolicyGateError::Parse {
            tenant_label: tenant.0.clone(),
            message: e.to_string(),
        })?;
        self.bundles.insert(Some(tenant), policies);
        Ok(self)
    }

    /// Prefer the tenant-specific bundle; fall back to the default.
    fn resolve_bundle(&self, tenant: &Option<TenantId>) -> Option<&PolicySet> {
        if let Some(t) = tenant {
            if let Some(p) = self.bundles.get(&Some(t.clone())) {
                return Some(p);
            }
        }
        self.bundles.get(&None)
    }

    fn build_uid(prefix: &str, name: &str) -> Result<EntityUid, GateDecision> {
        EntityUid::from_str(&format!("{prefix}::\"{name}\"")).map_err(|e| GateDecision::Deny {
            code: format!("cedar.{}.invalid", prefix.to_ascii_lowercase()),
            reason: format!("could not form {prefix} uid: {e}"),
        })
    }
}

#[async_trait]
impl Gate for PolicyGate {
    async fn evaluate(&self, req: GateRequest) -> GateDecision {
        let tenant = current_tenant();
        let Some(policies) = self.resolve_bundle(&tenant) else {
            return GateDecision::Deny {
                code: "policy.no_bundle_for_tenant".into(),
                reason: format!(
                    "no Cedar bundle registered for tenant {tenant:?} and no default fallback present"
                ),
            };
        };

        let principal = match Self::build_uid("Tool", &req.tool_name) {
            Ok(u) => u,
            Err(d) => return d,
        };
        let action = match EntityUid::from_str("Action::\"invoke\"") {
            Ok(u) => u,
            Err(e) => {
                return GateDecision::Deny {
                    code: "cedar.action.invalid".into(),
                    reason: format!("could not form action uid: {e}"),
                };
            }
        };
        let resource = match Self::build_uid("Args", &req.tool_name) {
            Ok(u) => u,
            Err(d) => return d,
        };

        // Lift all top-level args into Cedar context (recursive for nested
        // objects and arrays). Any unrepresentable value (null, out-of-range
        // u64) causes an immediate Deny so we fail-CLOSED rather than
        // silently dropping an attribute a forbid rule may reference.
        let ctx_pairs: Vec<(String, cedar_policy::RestrictedExpression)> = match &req.args {
            serde_json::Value::Object(map) => {
                let mut pairs = Vec::with_capacity(map.len());
                for (k, v) in map {
                    match json_to_cedar_expr(v) {
                        Some(expr) => pairs.push((k.clone(), expr)),
                        None => {
                            return GateDecision::Deny {
                                code: "cedar.context.unrepresentable".into(),
                                reason: format!("arg `{k}` is not Cedar-representable"),
                            };
                        }
                    }
                }
                pairs
            }
            _ => Vec::new(),
        };

        let context = match Context::from_pairs(ctx_pairs) {
            Ok(c) => c,
            Err(e) => {
                return GateDecision::Deny {
                    code: "cedar.context.invalid".into(),
                    reason: format!("context build: {e}"),
                };
            }
        };

        let request = match Request::new(principal, action, resource, context, None) {
            Ok(r) => r,
            Err(e) => {
                return GateDecision::Deny {
                    code: "cedar.request.invalid".into(),
                    reason: format!("request build: {e}"),
                };
            }
        };

        let entities = Entities::empty();
        let response = self.authorizer.is_authorized(&request, policies, &entities);
        match response.decision() {
            cedar_policy::Decision::Allow => GateDecision::Allow,
            cedar_policy::Decision::Deny => GateDecision::Deny {
                code: "cedar.deny".into(),
                reason: "cedar deny".to_string(),
            },
        }
    }

    fn name(&self) -> &'static str {
        "PolicyGate"
    }
}

fn json_to_cedar_expr(v: &serde_json::Value) -> Option<cedar_policy::RestrictedExpression> {
    use cedar_policy::RestrictedExpression as RE;
    match v {
        serde_json::Value::Bool(b) => Some(RE::new_bool(*b)),
        serde_json::Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                Some(RE::new_long(i))
            } else if let Some(u) = n.as_u64() {
                // u64 values >= 2^63 cannot fit in Cedar long; treat as
                // unrepresentable so the caller can deny fail-CLOSED.
                i64::try_from(u).ok().map(RE::new_long)
            } else {
                // Cedar decimal extension stores fixed-point with up to 4
                // fractional digits. Format to exactly that precision so
                // `decimal("0.5000") == context.field` evaluations match.
                n.as_f64().map(|f| RE::new_decimal(format!("{f:.4}")))
            }
        }
        serde_json::Value::String(s) => Some(RE::new_string(s.clone())),
        serde_json::Value::Array(xs) => {
            let elems: Option<Vec<RE>> = xs.iter().map(json_to_cedar_expr).collect();
            elems.map(RE::new_set)
        }
        serde_json::Value::Object(map) => {
            let pairs: Option<Vec<(String, RE)>> = map
                .iter()
                .map(|(k, v)| json_to_cedar_expr(v).map(|expr| (k.clone(), expr)))
                .collect();
            pairs.and_then(|p| RE::new_record(p).ok())
        }
        serde_json::Value::Null => None,
    }
}