cedarling 0.0.58

// This software is available under the Apache-2.0 license.
// See https://www.apache.org/licenses/LICENSE-2.0.txt for full text.
//
// Copyright (c) 2024, Gluu, Inc.

use crate::entity_builder::schema::AttrSrc;
use cedar_policy::RestrictedExpression;
use serde_json::Value;
use smol_str::{SmolStr, ToSmolStr};
use std::collections::HashMap;

use super::built_entities::BuiltEntities;
use super::schema::AttrsShape;
use super::value_to_expr::value_to_expr;
use super::{BuildAttrsError, BuildAttrsErrorVec};

/// Builds Cedar entity attributes using the given attribute source
///
/// This function will *JOIN* the token attributes built from each token
pub(super) fn build_entity_attrs(
    attrs_src: &HashMap<String, Value>,
    entities: &BuiltEntities,
    attrs_shape: Option<&HashMap<SmolStr, AttrsShape>>,
) -> Result<HashMap<String, RestrictedExpression>, BuildAttrsErrorVec> {
    if let Some(attrs_shape) = attrs_shape {
        build_entity_attrs_with_shape(|name| attrs_src.get(name), entities, attrs_shape)
    } else {
        build_entity_attrs_without_schema(attrs_src)
    }
}

/// Schema-driven variant that accepts a borrow-based lookup closure instead of
/// a full `HashMap`.
pub(super) fn build_entity_attrs_with_shape_lookup<'a, F>(
    lookup: F,
    entities: &BuiltEntities,
    attrs_shape: &HashMap<SmolStr, AttrsShape>,
) -> Result<HashMap<String, RestrictedExpression>, BuildAttrsErrorVec>
where
    F: Fn(&str) -> Option<&'a Value>,
{
    build_entity_attrs_with_shape(lookup, entities, attrs_shape)
}

/// Uses the schema as guide in building the attributes
fn build_entity_attrs_with_shape<'a, F>(
    lookup: F,
    entities: &BuiltEntities,
    attrs_shape: &HashMap<SmolStr, AttrsShape>,
) -> Result<HashMap<String, RestrictedExpression>, BuildAttrsErrorVec>
where
    F: Fn(&str) -> Option<&'a Value>,
{
    let mut errs = Vec::new();
    let mut attrs = HashMap::new();
    let mut required_missing_claims: Vec<SmolStr> = Vec::new();
    let mut missing_entity_ref_sets: Vec<SmolStr> = Vec::new();

    for (attr_name, attr_shape) in attrs_shape {
        match attr_shape.src() {
            AttrSrc::JwtClaim(claim_src) => {
                // skip if the source couldn't be found and was not required
                let Some(src) = lookup(attr_name.as_str()) else {
                    if attr_shape.is_required() {
                        required_missing_claims.push(attr_name.to_smolstr());
                    }

                    continue;
                };

                match claim_src.build_expr(src) {
                    Ok(Some(expr)) => {
                        attrs.insert(attr_name.to_string(), expr);
                    },
                    Err(e) if attr_shape.is_required() => {
                        errs.push(BuildAttrsError::from(e));
                    },
                    _ => {},
                }
            },
            AttrSrc::EntityRef(entity_ref_src) => {
                let eid_opt = match entities.get_single(entity_ref_src) {
                    Some(eid) => Some(eid),
                    None => match entities.get_multiple(entity_ref_src) {
                        Some(eids) => {
                            let src_val_opt =
                                lookup(attr_name.as_str()).and_then(|val| val.as_str());

                            match src_val_opt {
                                Some(src_val) if eids.contains(&src_val.to_smolstr()) => {
                                    Some(src_val)
                                },
                                _ => None,
                            }
                        },
                        None => None,
                    },
                };

                let Some(eid) = eid_opt else {
                    if attr_shape.is_required() {
                        errs.push(BuildAttrsError::MissingEntityRefs(vec![
                            (**entity_ref_src).clone(),
                        ]));
                    }
                    continue;
                };

                match entity_ref_src.build_expr(eid) {
                    Ok(src) => {
                        attrs.insert(attr_name.to_string(), src);
                    },
                    Err(e) => {
                        errs.push(BuildAttrsError::from(e));
                    },
                }
            },
            AttrSrc::EntityRefSet(entity_ref_set_src) => {
                let Some(eids) = entities.get_multiple(entity_ref_set_src) else {
                    if attr_shape.is_required() {
                        missing_entity_ref_sets.push((*entity_ref_set_src).clone());
                    }
                    continue;
                };

                match entity_ref_set_src.build_expr(eids) {
                    Ok(src) => {
                        attrs.insert(attr_name.to_string(), src);
                    },
                    Err(e) => {
                        errs.push(BuildAttrsError::from(e));
                    },
                }
            },
        }
    }

    if !required_missing_claims.is_empty() {
        errs.push(BuildAttrsError::MissingClaims(required_missing_claims));
    }
    if !missing_entity_ref_sets.is_empty() {
        errs.push(BuildAttrsError::MissingEntityRefs(missing_entity_ref_sets));
    }

    if !errs.is_empty() {
        return Err(BuildAttrsErrorVec(errs));
    }
    Ok(attrs)
}

/// Will do it's best to create the entity without a schema
fn build_entity_attrs_without_schema(
    attrs_src: &HashMap<String, Value>,
) -> Result<HashMap<String, RestrictedExpression>, BuildAttrsErrorVec> {
    let mut errs = Vec::new();
    let mut attrs = HashMap::new();

    for (name, src) in attrs_src {
        // without claim mapping
        match value_to_expr(src) {
            Ok(Some(expr)) => {
                attrs.insert(name.clone(), expr);
            },
            Err(e) => {
                errs.push(e);
            },
            _ => {},
        }
    }

    if !errs.is_empty() {
        return Err(BuildAttrsErrorVec::from(
            errs.into_iter().flatten().collect::<Vec<_>>(),
        ));
    }

    Ok(attrs)
}

#[cfg(test)]
mod test {
    use super::super::test::assert_entity_eq;
    use super::*;
    use crate::entity_builder::schema::MappingSchema;
    use cedar_policy::{Entity, Schema};
    use cedar_policy_core::validator::ValidatorSchema;
    use serde_json::json;
    use std::{collections::HashSet, str::FromStr};

    #[test]
    fn can_build_entity_with_schema() {
        let schema_src = r"
        namespace SomeNamespace {
            entity AnotherEntity;
            entity SomeEntity {
                bool_attr: Bool,
                str_attr: String,
                long_attr: Long,
                set_bool_attr: Set<Bool>,
                set_set_bool_attr: Set<Set<Bool>>,
                record_attr: {
                    inner_record_attr: Bool,
                },
                record_record_attr: {
                    inner_record_attr: {
                        inner_inner_record_attr: Bool,
                    },
                },
                entity_ref_attr: AnotherEntity,
                decimal_attr: decimal,
                ip_attr: ipaddr,
                optional_attr?: Bool,
            };
        }
        ";
        let cedar_schema = Schema::from_str(schema_src).expect("builds cedar Schema");
        let mapping_schema: MappingSchema = (&ValidatorSchema::from_str(schema_src)
            .expect("builds ValidatorSchema"))
            .try_into()
            .expect("builds MappingSchema");
        let entity_name = "SomeNamespace::SomeEntity";
        let attrs_src = HashMap::from([
            ("bool_attr".into(), json!(true)),
            ("str_attr".into(), json!("some_str")),
            ("long_attr".into(), json!(1234)),
            ("set_bool_attr".into(), json!([true, true])),
            ("set_set_bool_attr".into(), json!([[true], [true]])),
            ("record_attr".into(), json!({"inner_record_attr": true})),
            (
                "record_record_attr".into(),
                json!({"inner_record_attr": {"inner_inner_record_attr": true}}),
            ),
            ("decimal_attr".into(), json!("0.0")),
            ("ip_attr".into(), json!("0.0.0.0")),
        ]);
        let mut built_entities = BuiltEntities::default();
        built_entities.insert(
            &"SomeNamespace::AnotherEntity::\"another_id\""
                .parse()
                .expect("a valid entity uid"),
        );

        let attrs_shape = mapping_schema
            .get_entity_shape(entity_name)
            .expect("get entity requirements");
        let attrs = build_entity_attrs(&attrs_src, &built_entities, Some(attrs_shape))
            .expect("builds entity attrs");

        let dummy_entity = Entity::new(
            "SomeNamespace::SomeEntity::\"some_id\""
                .parse()
                .expect("a valid entity uid"),
            attrs,
            HashSet::new(),
        )
        .expect("builds dummy entity");

        assert_entity_eq(
            &dummy_entity,
            &json!({
                "uid": {"type": "SomeNamespace::SomeEntity", "id": "some_id"},
                "attrs": {
                    "bool_attr": true,
                    "str_attr": "some_str",
                    "long_attr": 1234,
                    "set_bool_attr": [true],
                    "set_set_bool_attr": [[true]],
                    "record_attr": {
                        "inner_record_attr": true,
                    },
                    "record_record_attr": {
                        "inner_record_attr": {
                            "inner_inner_record_attr": true,
                        },
                    },
                    "entity_ref_attr": {"__entity": {
                        "type": "SomeNamespace::AnotherEntity", "id": "another_id",
                    }},
                    "decimal_attr": {"__extn": {
                        "fn": "decimal",
                        "arg": "0.0",
                    }},
                    "ip_attr": {"__extn": {
                        "fn": "ip",
                        "arg": "0.0.0.0",
                    }},
                },
                "parents": [],
            }),
            Some(&cedar_schema),
        );
    }
}