Fix 1176

cedar-policy-core/src/entities/conformance.rs

@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+use std::collections::BTreeMap;
+
 use super::{
     json::err::TypeMismatchError, schematype_of_restricted_expr, EntityTypeDescription,
     GetSchemaTypeError, HeterogeneousSetError, Schema, SchemaType,
@@ -24,6 +26,7 @@ use crate::ast::{
 use crate::extensions::{ExtensionFunctionLookupError, Extensions};
 use either::Either;
 use miette::Diagnostic;
+use smol_str::SmolStr;
 use thiserror::Error;
 pub mod err;
 
@@ -176,6 +179,68 @@ pub fn typecheck_value_against_schematype(
     }
 }
 
+/// Check whether the given `RestrictedExpr` is a valid instance of `SchemaType`
+pub fn does_restricted_expr_implement_schematype(
+    expr: BorrowedRestrictedExpr<'_>,
+    expr_ty: &SchemaType,
+    expected_ty: &SchemaType,
+    extensions: &Extensions<'_>,
+) -> bool {
+    use SchemaType::*;
+    if expr_ty == expected_ty {
+        return true;
+    }
+
+    match (expr_ty, expected_ty) {
+        (Set { .. }, EmptySet) => true,
+        (EmptySet, Set { .. }) => true,
+        (
+            Set {
+                element_ty: expr_elm_ty,
+            },
+            Set { element_ty: elty },
+        ) => match expr.as_set_elements() {
+            Some(mut els) => els.all(|e| {
+                does_restricted_expr_implement_schematype(e, expr_elm_ty, elty, extensions)
+            }),
+            None => false,
+        },
+        (
+            Record {
+                attrs: expr_attrs, ..
+            },
+            Record { attrs, open_attrs },
+        ) => match expr.as_record_pairs() {
+            Some(pairs) => {
+                let pairs_map: BTreeMap<&SmolStr, BorrowedRestrictedExpr<'_>> = pairs.collect();
+                let all_req_schema_attrs_in_record =
+                    attrs.iter().all(|(k, v)| match pairs_map.get(k) {
+                        Some(inner_e) => does_restricted_expr_implement_schematype(
+                            *inner_e,
+                            &expr_attrs.get(k).unwrap().attr_type,
+                            &v.attr_type,
+                            extensions,
+                        ),
+                        None => !v.required,
+                    });
+                let all_rec_attrs_in_schema =
+                    pairs_map.iter().all(|(k, inner_e)| match attrs.get(*k) {
+                        Some(sch_ty) => does_restricted_expr_implement_schematype(
+                            *inner_e,
+                            &expr_attrs.get(*k).unwrap().attr_type,
+                            &sch_ty.attr_type,
+                            extensions,
+                        ),
+                        None => false,
+                    });
+                (*open_attrs || all_rec_attrs_in_schema) && all_req_schema_attrs_in_record
+            }
+            None => false,
+        },
+        _ => false,
+    }
+}
+
 /// Check whether the given `RestrictedExpr` typechecks with the given `SchemaType`.
 /// If the typecheck passes, return `Ok(())`.
 /// If the typecheck fails, return an appropriate `Err`.
@@ -190,7 +255,8 @@ pub fn typecheck_restricted_expr_against_schematype(
     // directly?
     match schematype_of_restricted_expr(expr, extensions) {
         Ok(actual_ty) => {
-            if actual_ty.is_consistent_with(expected_ty) {
+            if does_restricted_expr_implement_schematype(expr, &actual_ty, expected_ty, extensions)
+            {
                 // typecheck passes
                 Ok(())
             } else {

cedar-policy-core/src/entities/json/schema_types.rs

@@ -69,9 +69,9 @@ pub enum SchemaType {
 #[derive(Debug, Hash, PartialEq, Eq, Clone)]
 pub struct AttributeType {
     /// Type of the attribute
-    attr_type: SchemaType,
+    pub(crate) attr_type: SchemaType,
     /// Is the attribute required
-    required: bool,
+    pub(crate) required: bool,
 }
 
 impl SchemaType {
@@ -354,11 +354,8 @@ impl Diagnostic for NontrivialResidualError {
 /// required or optional.
 /// This function, when given a record that has keys A, B, and C, will return a
 /// `SchemaType` where A, B, and C are all marked as optional attributes, but no
-/// other attributes are possible.
-/// That is, this assumes that all existing attributes are optional, but that no
-/// other optional attributes are possible.
-/// Compared to marking A, B, and C as required, this allows the returned
-/// `SchemaType` to `is_consistent_with()` more types.
+/// other attributes are possible. This maximized flexibility while avoiding
+/// heterogeneous sets.
 ///
 /// This function may return `GetSchemaTypeError`, but should never return
 /// `NontrivialResidual`, because `RestrictedExpr`s can't contain nontrivial
@@ -382,9 +379,8 @@ pub fn schematype_of_restricted_expr(
                         BorrowedRestrictedExpr::new_unchecked(v), // assuming the invariant holds for the record as a whole, it will also hold for each attribute value
                         extensions,
                     )?;
-                    // we can't know if the attribute is required or optional,
-                    // but marking it optional is more flexible -- allows the
-                    // attribute type to `is_consistent_with()` more types
+                    // We can't know if the attribute is required or optional.
+                    // Keep as optional to minimize heterogeneous sets.
                     Ok((k.clone(), AttributeType::optional(attr_type)))
                 }).collect::<Result<BTreeMap<_,_>, GetSchemaTypeError>>()?,
                 open_attrs: false,

cedar-policy/src/tests.rs

@@ -1382,6 +1382,7 @@ mod ancestors_tests {
 /// schema-based parsing.
 mod entity_validate_tests {
     use super::*;
+    use cool_asserts::assert_matches;
     use entities::err::EntitiesError;
     use serde_json::json;
 
@@ -1709,6 +1710,158 @@ mod entity_validate_tests {
             }
         }
     }
+
+    /// Record inside entity doesn't conform to schema
+    #[test]
+    fn issue_1176_should_fail() {
+        let (schema, _) = Schema::from_cedarschema_str(
+            "
+            entity E {
+              rec: {
+                foo: Long
+              }
+            };
+            action Act appliesTo {
+              principal: [E],
+              resource: [E],
+            };
+        ",
+        )
+        .unwrap();
+        let entity = Entity::new(
+            EntityUid::from_str(r#"E::"abc""#).unwrap(),
+            HashMap::from_iter([(
+                "rec".into(),
+                RestrictedExpression::new_record([
+                    ("foo".into(), RestrictedExpression::new_long(4567)),
+                    (
+                        "extra".into(),
+                        RestrictedExpression::new_string("bad".into()),
+                    ),
+                ])
+                .unwrap(),
+            )]),
+            HashSet::new(),
+        )
+        .unwrap();
+        assert_matches!(
+            Entities::from_entities([entity], Some(&schema)),
+            Err(EntitiesError::InvalidEntity(_))
+        );
+    }
+
+    #[test]
+    fn issue_1176_should_pass_1() {
+        let (schema, _) = Schema::from_cedarschema_str(
+            r###"
+entity A = {"foo": Set < Set < {"bar": __cedar::Bool, "baz"?: __cedar::Bool} > >};
+action "g" appliesTo {
+  principal: [A],
+  resource: [A],
+};
+        "###,
+        )
+        .unwrap();
+        let entity_str = r###"
+        {
+            "uid": {
+              "type": "A",
+              "id": "alice"
+            },
+            "attrs": {
+              "foo": [
+                [],
+                [
+                  {
+                    "bar": false
+                  },
+                  {
+                    "bar": true
+                  },
+                  {
+                    "bar": true,
+                    "baz": true
+                  }
+                ],
+                [
+                  {
+                    "bar": false,
+                    "baz": false
+                  },
+                  {
+                    "bar": true
+                  }
+                ],
+                [
+                  {
+                    "bar": true
+                  },
+                  {
+                    "bar": true,
+                    "baz": false
+                  }
+                ]
+              ]
+            },
+            "parents": []
+          }
+        "###;
+
+        assert_matches!(Entity::from_json_str(entity_str, Some(&schema)), Ok(_));
+    }
+
+    #[test]
+    fn issue_1176_should_pass_2() {
+        let (schema, _) = Schema::from_cedarschema_str(
+            r###"
+            entity User {
+              allowedTagsForRole: {
+                "Role-A"?: {
+                    production_status?: Set<String>,
+                    country?: Set<String>,
+                    stage?: Set<String>,
+                },
+                "Role-B"?: {
+                    production_status?: Set<String>,
+                    country?: Set<String>,
+                    stage?: Set<String>,
+                },
+              },
+            };
+
+            action UpdateWorkspace appliesTo {
+              principal: User,
+              resource: User,
+            };
+        "###,
+        )
+        .unwrap();
+        let entity_str = r###"
+        {
+            "uid": {
+                "type": "User",
+                "id": "Alice"
+            },
+            "attrs": {
+                "allowedTagsForRole": {
+                    "Role-B": {
+                        "production_status": [
+                            "production"
+                        ],
+                        "country": [
+                            "ALL"
+                        ],
+                        "stage": [
+                            "valuation"
+                        ]
+                    }
+                }
+            },
+            "parents": []
+        }
+        "###;
+        assert_matches!(Entity::from_json_str(entity_str, Some(&schema)), Ok(_));
+    }
 }
 
 /// The main unit tests for schema-based parsing live here, as they require both
@@ -1743,7 +1896,7 @@ mod schema_based_parsing_tests {
     #[test]
     #[allow(clippy::too_many_lines)]
     #[allow(clippy::cognitive_complexity)]
-    fn signle_attr_types() {
+    fn single_attr_types() {
         let schema = Schema::from_json_value(json!(
         {"": {
             "entityTypes": {