cedar-policy-symcc 0.3.2

/*
 * Copyright Cedar Contributors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! This module includes functions to convert
//! literal Term/SymRequest/SymEntities to their
//! concrete versions

use std::borrow::Borrow;
use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::sync::Arc;

use cedar_policy::{Entities, EntityId, EntityTypeName, EntityUid, Request};
use cedar_policy_core::ast::{
    Context, Entity, EntityAttrEvaluationError, Expr, ExprVisitor, Literal, Set, Value, ValueKind,
};
use cedar_policy_core::entities::{NoEntitiesSchema, TCComputation};
use cedar_policy_core::extensions::Extensions;
use cedar_policy_core::parser::Loc;
use miette::Diagnostic;
use num_bigint::{BigInt, TryFromBigIntError};
use smol_str::SmolStr;
use thiserror::Error;

use crate::symcc::factory;

use super::env::{SymEntities, SymEntityData, SymRequest};
use super::ext::ExtError;
use super::function::{Udf, UnaryFunction};
use super::term::{Term, TermPrim};
use super::SymEnv;

/// Errors that happen during concretization, i.e., the process
/// of converting literal [`Term`]s back to representatinos in
/// [`cedar_policy`]/[`cedar_policy_core`].
#[derive(Debug, Diagnostic, Error)]
pub enum ConcretizeError {
    /// Got no policies, expected to have at least one.
    #[error("expected to have at least one policy")]
    NoPolicies,
    /// Expecting a literal entity.
    #[error("Not a literal entity: {0:?}")]
    NotLiteralEntity(Term),
    /// Expecting a literal string.
    #[error("Not a literal string: {0:?}")]
    NotLiteralString(Term),
    /// Cannot convert a term to a value.
    #[error("Unable to convert {0:?} to a value")]
    UnableToConvertToValue(Term),
    /// Cannot convert a term to a context.
    #[error("Unable to convert {0:?} to a context")]
    UnableToConvertToContext(Term),
    /// Unable to construct entity.
    #[error("Unable to construct a valid entity")]
    UnableToConstructEntity(#[from] EntityAttrEvaluationError),
    /// Entity type not found.
    #[error("Entity type not found: {0}")]
    EntityTypeNotFound(EntityTypeName),
    /// Unable to falidate request.
    #[error("Request validation error")]
    RequestValidationError(#[from] cedar_policy::RequestValidationError),
    /// Errors when constructing entity.
    #[error("Unable to construct entities")]
    EntitiesError(#[from] cedar_policy::entities_errors::EntitiesError),
    /// Fail to convert from big integer.
    #[error("Unable to convert BitVec to integer")]
    TryFromBigIntError(#[from] TryFromBigIntError<BigInt>),
    /// Extension error.
    #[error("extension error")]
    ExtError(#[from] ExtError),
    /// Unsupported expression.
    #[error("unsupported expression: {0}")]
    UnsupportedExpr(Expr),
}

/// A concrete environment recovered from a [`SymEnv`].
#[derive(Debug, Clone, PartialEq)]
pub struct Env {
    /// Concrete request
    pub request: Request,
    /// Concrete entities
    pub entities: Entities,
}

impl std::fmt::Display for Env {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let req = &self.request;
        let principal = req
            .principal()
            .map_or_else(|| "unknown".to_string(), |p| p.to_string());
        let action = req
            .action()
            .map_or_else(|| "unknown".to_string(), |a| a.to_string());
        let resource = req
            .resource()
            .map_or_else(|| "unknown".to_string(), |r| r.to_string());
        write!(
            f,
            "principal: {principal}, action: {action}, resource: {resource}"
        )?;
        if let Some(ctx) = req.context() {
            write!(f, "\ncontext: {ctx}")?;
        }
        let entities = &self.entities;
        if !entities.is_empty() {
            writeln!(f, "\nentities: [")?;
            for entity in entities.iter() {
                let uid = entity.uid();
                write!(f, "  {uid}")?;
                // ancestors
                if let Some(ancestors) = entities.ancestors(&uid) {
                    let ancs: Vec<_> = ancestors.map(|a| a.to_string()).collect();
                    if !ancs.is_empty() {
                        write!(f, " in [{}]", ancs.join(", "))?;
                    }
                }
                // attrs
                let attrs: Vec<_> = entity.attrs().collect();
                if !attrs.is_empty() {
                    writeln!(f, " {{")?;
                    for (k, v) in &attrs {
                        match v {
                            Ok(val) => writeln!(f, "    {k}: {val},")?,
                            Err(_) => writeln!(f, "    {k}: <unknown>,")?,
                        }
                    }
                    write!(f, "  }}")?;
                }
                // tags
                let tags: Vec<_> = entity.tags().collect();
                if !tags.is_empty() {
                    writeln!(f, " tags {{")?;
                    for (k, v) in &tags {
                        match v {
                            Ok(val) => writeln!(f, "    {k}: {val},")?,
                            Err(_) => writeln!(f, "    {k}: <unknown>,")?,
                        }
                    }
                    write!(f, "  }}")?;
                }
                writeln!(f, ",")?;
            }
            write!(f, "]")?;
        }
        Ok(())
    }
}

/// Tries to extract an `EntityUid` from a `Term`.
/// Corresponds to `Term.entityUID?` in `Concretize.lean`
impl TryFrom<&Term> for EntityUid {
    type Error = ConcretizeError;

    fn try_from(term: &Term) -> Result<Self, Self::Error> {
        if let Term::Prim(TermPrim::Entity(uid)) = term {
            Ok(uid.clone())
        } else {
            Err(ConcretizeError::NotLiteralEntity(term.clone()))
        }
    }
}

/// Tries to extract a set of `EntityUid`'s from a `Term`.
/// Corresponds `Term.setOfEntityUIDs?` in `Concretize.lean`
impl TryFrom<&Term> for BTreeSet<EntityUid> {
    type Error = ConcretizeError;

    fn try_from(term: &Term) -> Result<Self, Self::Error> {
        if let Term::Set { elts, .. } = term {
            Ok(elts
                .iter()
                .map(|t| t.try_into())
                .collect::<Result<_, _>>()?)
        } else {
            Err(ConcretizeError::NotLiteralEntity(term.clone()))
        }
    }
}

/// Tries to convert a `Term` to a string.
impl TryFrom<&Term> for SmolStr {
    type Error = ConcretizeError;

    fn try_from(term: &Term) -> Result<Self, Self::Error> {
        if let Term::Prim(TermPrim::String(s)) = term {
            Ok(s.clone())
        } else {
            Err(ConcretizeError::NotLiteralString(term.clone()))
        }
    }
}

/// Tries to extract a set of `Strings`'s from a `Term`.
impl TryFrom<&Term> for BTreeSet<SmolStr> {
    type Error = ConcretizeError;

    fn try_from(term: &Term) -> Result<Self, Self::Error> {
        if let Term::Set { elts, .. } = term {
            Ok(elts
                .iter()
                .map(|t| t.try_into())
                .collect::<Result<_, _>>()?)
        } else {
            Err(ConcretizeError::NotLiteralEntity(term.clone()))
        }
    }
}

impl TryFrom<&Term> for Value {
    type Error = ConcretizeError;

    fn try_from(term: &Term) -> Result<Self, Self::Error> {
        match term {
            Term::Prim(TermPrim::Bool(b)) => {
                Ok(Value::new(ValueKind::Lit(Literal::Bool(*b)), None))
            }

            Term::Prim(TermPrim::Bitvec(v)) => Ok(Value::new(
                ValueKind::Lit(Literal::Long(v.to_int().try_into()?)),
                None,
            )),

            Term::Prim(TermPrim::String(s)) => {
                Ok(Value::new(ValueKind::Lit(Literal::String(s.clone())), None))
            }

            Term::Prim(TermPrim::Entity(uid)) => Ok(Value::new(
                ValueKind::Lit(Literal::EntityUID(Arc::new(uid.clone().into()))),
                None,
            )),

            Term::Prim(TermPrim::Ext(ext)) => Ok(Self::try_from(ext)?),

            Term::Set { elts, .. } => Ok(Value::new(
                ValueKind::Set(Set::new(
                    elts.iter()
                        .map(|t| t.try_into())
                        .collect::<Result<Vec<_>, _>>()?,
                )),
                None,
            )),

            Term::Record(rec) => Ok(Value::new(
                ValueKind::Record(Arc::new(
                    rec.iter()
                        .map(|(k, v)| {
                            if let Term::Some(t) = v {
                                Ok(Some((k.clone(), t.as_ref().try_into()?)))
                            } else if let Term::None(_) = v {
                                // None fields are simply ignored
                                Ok(None)
                            } else {
                                Ok(Some((k.clone(), v.try_into()?)))
                            }
                        })
                        .collect::<Result<Vec<Option<_>>, ConcretizeError>>()?
                        .into_iter()
                        .flatten()
                        .collect(),
                )),
                None,
            )),

            // Otherwise it's not convertable
            _ => Err(ConcretizeError::UnableToConvertToValue(term.clone())),
        }
    }
}

impl SymRequest {
    pub fn concretize(&self) -> Result<Request, ConcretizeError> {
        Ok(Request::new(
            (&self.principal).try_into()?,
            (&self.action).try_into()?,
            (&self.resource).try_into()?,
            Context::Value(self.context.try_into_record()?).into(),
            None, // TODO: schema == None disables request validation
        )?)
    }

    fn get_all_entity_uids(&self, uids: &mut BTreeSet<EntityUid>) {
        self.context.get_all_entity_uids(uids);
        self.principal.get_all_entity_uids(uids);
        self.action.get_all_entity_uids(uids);
        self.resource.get_all_entity_uids(uids);
    }
}

impl Term {
    /// Tries to convert a term into a record
    ///
    /// Corresponds to `Term.recordValue?` in `Concretize.lean`
    fn try_into_record(&self) -> Result<Arc<BTreeMap<SmolStr, Value>>, ConcretizeError> {
        if let Value {
            value: ValueKind::Record(record),
            ..
        } = self.try_into()?
        {
            Ok(record)
        } else {
            Err(ConcretizeError::UnableToConvertToContext(self.clone()))
        }
    }

    /// Collect all entity UIDs occurring in the term
    ///
    /// Corresponds to `Term.entityUIDs` in `Concretizer.lean`
    pub(crate) fn get_all_entity_uids(&self, uids: &mut BTreeSet<EntityUid>) {
        match self {
            Term::Prim(TermPrim::Entity(uid)) => {
                uids.insert(uid.clone());
            }

            Term::Some(t) => {
                t.get_all_entity_uids(uids);
            }

            Term::Set { elts, .. } => {
                for t in elts.iter() {
                    t.get_all_entity_uids(uids);
                }
            }

            Term::Record(rec) => {
                for t in rec.values() {
                    t.get_all_entity_uids(uids);
                }
            }

            Term::App { args, .. } => {
                for t in args.iter() {
                    t.get_all_entity_uids(uids);
                }
            }

            _ => {}
        }
    }
}

impl Udf {
    fn get_all_entity_uids(&self, uids: &mut BTreeSet<EntityUid>) {
        self.default.get_all_entity_uids(uids);
        for (k, v) in self.table.iter() {
            k.get_all_entity_uids(uids);
            v.get_all_entity_uids(uids);
        }
    }
}

impl UnaryFunction {
    /// Corresponds to `UnaryFunction.entityUIDs` in `Concretize.lean`
    fn get_all_entity_uids(&self, uids: &mut BTreeSet<EntityUid>) {
        match self {
            UnaryFunction::Udf(udf) => udf.get_all_entity_uids(uids),
            UnaryFunction::Uuf(_) => {}
        }
    }
}

impl SymEntityData {
    /// Concretizes a particular entity.
    pub fn concretize(&self, euid: &EntityUid) -> Result<Entity, ConcretizeError> {
        let tuid = Term::Prim(TermPrim::Entity(euid.clone()));

        let concrete_attrs = factory::app(self.attrs.clone(), tuid.clone()).try_into_record()?;

        // For each ancestor entity type, apply the suitable ancestor function
        // to obtain a concrete set of ancestor EUIDs
        let concrete_ancestors = self
            .ancestors
            .values()
            .map(|ancestor| {
                let euids: BTreeSet<EntityUid> =
                    (&factory::app(ancestor.clone(), tuid.clone())).try_into()?;

                Ok(euids.into_iter().map(|euid| euid.as_ref().clone()))
            })
            .collect::<Result<Vec<_>, ConcretizeError>>()?
            .into_iter()
            .flatten()
            .collect::<HashSet<_>>();

        // Read tags from the model
        let tags = if let Some(tags) = &self.tags {
            // Get all valid tag keys first
            let keys: BTreeSet<SmolStr> =
                (&factory::app(tags.keys.clone(), tuid.clone())).try_into()?;

            keys.into_iter()
                .map(|k| {
                    // Using get_tag_unchecked here since we know already that k is in the key set
                    let val: Value = (&tags
                        .get_tag_unchecked(tuid.clone(), Term::Prim(TermPrim::String(k.clone()))))
                        .try_into()?;

                    Ok((k, val.into()))
                })
                .collect::<Result<_, ConcretizeError>>()?
        } else {
            BTreeMap::new()
        };

        Ok(Entity::new(
            euid.as_ref().clone(),
            concrete_attrs
                .as_ref()
                .clone()
                .into_iter()
                .map(|(k, v)| (k, v.into())),
            HashSet::new(),
            concrete_ancestors,
            tags,
            Extensions::all_available(),
        )?)
    }

    /// Corresponds to `SymEntityData.entityUIDs` in `Concretize.lean`
    fn get_all_entity_uids(&self, ety: &EntityTypeName, uids: &mut BTreeSet<EntityUid>) {
        if let Some(members) = &self.members {
            for member in members {
                uids.insert(EntityUid::from_type_name_and_id(
                    ety.clone(),
                    EntityId::new(member),
                ));
            }
        }

        self.attrs.get_all_entity_uids(uids);

        for ancestor in self.ancestors.values() {
            ancestor.get_all_entity_uids(uids);
        }

        if let Some(tags) = &self.tags {
            // tags.keys.get_all_entity_uids(uids);
            tags.vals.get_all_entity_uids(uids);
        }
    }
}

impl SymEntities {
    /// Concretizes a literal SymEntities to Entities
    pub fn concretize(&self, all_euids: &BTreeSet<EntityUid>) -> Result<Entities, ConcretizeError> {
        let mut entities = Vec::with_capacity(all_euids.len());

        for euid in all_euids {
            let sym_entity_data =
                self.0
                    .get(euid.type_name())
                    .ok_or(ConcretizeError::EntityTypeNotFound(
                        euid.type_name().clone(),
                    ))?;

            entities.push(sym_entity_data.concretize(euid)?);
        }

        // As the internal cedar_policy_core::entities::Entities
        let internal_entities = cedar_policy_core::entities::Entities::from_entities(
            entities.into_iter(),
            None::<&NoEntitiesSchema>,
            #[cfg(debug_assertions)]
            TCComputation::EnforceAlreadyComputed,
            #[cfg(not(debug_assertions))]
            TCComputation::AssumeAlreadyComputed,
            Extensions::all_available(),
        )?;

        Ok(Entities::from(internal_entities))
    }

    /// Corresponds to `SymEntities.entityUIDs` in `Concretize.lean`
    fn get_all_entity_uids(&self, uids: &mut BTreeSet<EntityUid>) {
        for (ety, data) in self.0.iter() {
            data.get_all_entity_uids(ety, uids);
        }
    }
}

/// An [`ExprVisitor`] to collect all entity UIDs occurring in an expression.
///
/// Corresponds to `Expr.entityUIDs` in `Concretize.lean`.
struct EntityUIDCollector<'a>(&'a mut BTreeSet<EntityUid>);

impl ExprVisitor for EntityUIDCollector<'_> {
    type Output = ();

    fn visit_literal(&mut self, lit: &Literal, _: Option<&Loc>) -> Option<Self::Output> {
        if let Literal::EntityUID(euid) = lit {
            self.0.insert(euid.as_ref().clone().into());
        }
        None
    }
}

impl SymEnv {
    /// Concretizes a literal [`SymEnv`] to a concrete [`Env`].
    ///
    /// In most cases, one should use [`SymEnv::extract`] instead
    /// to ensure well-formed output [`Env`].
    pub(crate) fn concretize<E: Borrow<Expr>>(
        &self,
        exprs: impl IntoIterator<Item = E>,
    ) -> Result<Env, ConcretizeError> {
        let mut uids = BTreeSet::new();
        self.request.get_all_entity_uids(&mut uids);
        self.entities.get_all_entity_uids(&mut uids);

        // Instead of using `footprint` and `Term::get_all_entity_uids`,
        // we collect EUIDs in expressions directly to avoid incorrect
        // short-circuiting in an incomplete entity store.
        let mut visitor = EntityUIDCollector(&mut uids);
        for expr in exprs {
            visitor.visit_expr(expr.borrow());
        }

        Ok(Env {
            request: self.request.concretize()?,
            entities: self.entities.concretize(&uids)?,
        })
    }
}