cedar_policy_symcc/symcc/

concretizer.rs

/*
 * 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, EvalResult, Request};
use cedar_policy_core::ast::{
    is_normalized_ident, 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,
}

/// Write an entity's attributes or tags
fn fmt_attrs_or_tags<'a, E>(
    f: &mut std::fmt::Formatter<'_>,
    label: Option<&str>,
    values: impl IntoIterator<Item = (&'a str, Result<EvalResult, E>)>,
) -> std::fmt::Result {
    let mut values = values.into_iter().peekable();
    if values.peek().is_none() {
        return Ok(());
    }
    if let Some(label) = label {
        write!(f, " {label}")?;
    }
    writeln!(f, " {{")?;
    for (k, v) in values {
        write!(f, "    ")?;
        if is_normalized_ident(k) {
            write!(f, "{k}")?;
        } else {
            write!(f, "\"{}\"", k.escape_debug())?;
        }
        match v {
            Ok(val) => writeln!(f, ": {val},")?,
            Err(_) => writeln!(f, ": <unknown>,")?,
        }
    }
    write!(f, "  }}")
}

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(", "))?;
                    }
                }
                fmt_attrs_or_tags(f, None, entity.attrs())?;
                fmt_attrs_or_tags(f, Some("tags"), entity.tags())?;
                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,
            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)?,
        })
    }
}

#[cfg(test)]
mod test {
    use super::Env;
    use cedar_policy::{Context, Entities, Entity, Request, RestrictedExpression};
    use insta::assert_snapshot;

    fn display_env(
        attrs: impl IntoIterator<Item = (String, RestrictedExpression)>,
        tags: impl IntoIterator<Item = (String, RestrictedExpression)>,
    ) -> String {
        let entity =
            Entity::new_with_tags(r#"User::"alice""#.parse().unwrap(), attrs, [], tags).unwrap();
        let request = Request::new(
            r#"User::"alice""#.parse().unwrap(),
            r#"Action::"view""#.parse().unwrap(),
            r#"Photo::"vacation""#.parse().unwrap(),
            Context::empty(),
            None,
        )
        .unwrap();
        Env {
            request,
            entities: Entities::from_entities([entity], None).unwrap(),
        }
        .to_string()
    }

    #[test]
    fn test_display_env() {
        assert_snapshot!(display_env(
            [("a".to_string(), RestrictedExpression::new_long(0))],
            [("b".to_string(), RestrictedExpression::new_long(1))],
        ), @r#"
        principal: User::"alice", action: Action::"view", resource: Photo::"vacation"
        context: {}
        entities: [
          User::"alice" {
            a: 0,
          } tags {
            b: 1,
          },
        ]
        "#);

        assert_snapshot!(display_env([("a".to_string(), RestrictedExpression::new_long(0))], [],), @r#"
        principal: User::"alice", action: Action::"view", resource: Photo::"vacation"
        context: {}
        entities: [
          User::"alice" {
            a: 0,
          },
        ]
        "#);

        assert_snapshot!(display_env([], [("b".to_string(), RestrictedExpression::new_long(1))],), @r#"
        principal: User::"alice", action: Action::"view", resource: Photo::"vacation"
        context: {}
        entities: [
          User::"alice" tags {
            b: 1,
          },
        ]
        "#);

        assert_snapshot!(display_env([], []), @r#"
        principal: User::"alice", action: Action::"view", resource: Photo::"vacation"
        context: {}
        entities: [
          User::"alice",
        ]
        "#);
    }

    #[test]
    fn display_attr_and_tag_names() {
        // the whole `Env`, for context on the per-name snapshots below
        let display_attr = |attr: &str| {
            display_env([(attr.to_string(), RestrictedExpression::new_long(0))], [])
                .lines()
                .find(|l| l.starts_with("    "))
                .expect("no attr line")
                .trim()
                .to_string()
        };
        assert_snapshot!(display_attr("plain"), @"plain: 0,");
        assert_snapshot!(display_attr("new\nline"), @r#""new\nline": 0,"#);

        let display_tag = |tag: &str| {
            display_env([], [(tag.to_string(), RestrictedExpression::new_long(0))])
                .lines()
                .find(|l| l.starts_with("    "))
                .expect("no tag line")
                .trim()
                .to_string()
        };
        assert_snapshot!(display_tag("plain"), @"plain: 0,");
        assert_snapshot!(display_tag("spa ce"), @r#""spa ce": 0,"#);
    }
}