Add depth guard for schema type parsing

cedar-policy-core/Cargo.toml

@@ -83,6 +83,7 @@ similar-asserts = "2.0.0"
 cool_asserts = "2.0"
 miette = { version = "7.6.0", features = ["fancy"] }
 insta = { version = "1.47.2", features = ["glob"] }
+rstest = "0.26.1"
 
 [lints]
 workspace = true

cedar-policy-core/src/validator.rs

@@ -52,12 +52,14 @@ pub use schema::err::*;
 pub use schema::*;
 mod deprecated_schema_compat;
 pub mod json_schema;
+pub(crate) mod schema_type_depth;
 mod str_checks;
 pub use str_checks::confusable_string_checks;
 pub mod cedar_schema;
 pub mod typecheck;
 use typecheck::Typechecker;
 mod partition_nonempty;
+mod topo_sort;
 pub mod types;
 
 /// Used to select how a policy will be validated.

cedar-policy-core/src/validator/cedar_schema.rs

@@ -18,6 +18,8 @@
 
 mod ast;
 pub use ast::Path;
+/// Depth computation for schema types
+pub mod depth;
 mod err;
 pub mod fmt;
 pub mod parser;

cedar-policy-core/src/validator/cedar_schema/ast.rs

@@ -307,6 +307,35 @@ pub enum Type {
     Record(Vec<Node<Annotated<AttrDecl>>>),
 }
 
+impl Drop for Type {
+    fn drop(&mut self) {
+        let mut work: Vec<Node<Type>> = Vec::new();
+        collect_child_types(self, &mut work);
+        while let Some(mut type_node) = work.pop() {
+            collect_child_types(&mut type_node.node, &mut work);
+        }
+    }
+}
+
+/// Extract all `Node<Type>` from `ty`, moving them out of `ty` and appending
+/// them to `out`.  After this call, `ty` holds no recursive `Node<Type>`, so it
+/// can be dropped without risking stack overflow due to recursion.
+fn collect_child_types(ty: &mut Type, out: &mut Vec<Node<Type>>) {
+    match ty {
+        Type::Ident(_) => {}
+        Type::Set(inner) => {
+            let dummy = Node::with_maybe_source_loc(Type::Record(vec![]), None);
+            out.push(std::mem::replace(inner.as_mut(), dummy));
+        }
+        Type::Record(attrs) => {
+            for attr_node in attrs.iter_mut() {
+                let dummy = Node::with_maybe_source_loc(Type::Record(vec![]), None);
+                out.push(std::mem::replace(&mut attr_node.node.data.ty, dummy));
+            }
+        }
+    }
+}
+
 /// Primitive Type Definitions
 #[derive(Debug, Clone)]
 pub enum PrimitiveType {

cedar-policy-core/src/validator/cedar_schema/depth.rs

@@ -0,0 +1,151 @@
+/*
+ * 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.
+ */
+
+//! Iterative depth computation for Cedar schema type AST.
+
+use itertools::Either;
+
+use super::ast::{ActionDecl, AppDecl, Declaration, EntityDecl, Namespace, Schema, Type};
+use crate::parser::Node;
+
+/// Iteratively compute the maximum nesting depth across all types in a schema.
+pub(crate) fn schema_type_depth(schema: &Schema) -> usize {
+    schema
+        .iter()
+        .map(|ns| namespace_type_depth(&ns.data))
+        .max()
+        .unwrap_or(0)
+}
+
+fn namespace_type_depth(ns: &Namespace) -> usize {
+    ns.decls
+        .iter()
+        .map(|d| declaration_type_depth(&d.data.node))
+        .max()
+        .unwrap_or(0)
+}
+
+fn declaration_type_depth(decl: &Declaration) -> usize {
+    match decl {
+        Declaration::Entity(entity_decl) => entity_decl_type_depth(entity_decl),
+        Declaration::Action(action_decl) => action_decl_type_depth(action_decl),
+        Declaration::Type(type_decl) => type_depth(&type_decl.def),
+    }
+}
+
+fn entity_decl_type_depth(decl: &EntityDecl) -> usize {
+    match decl {
+        EntityDecl::Standard(standard) => {
+            let attrs_d = standard
+                .attrs
+                .node
+                .iter()
+                .map(|a| 1 + type_depth(&a.node.data.ty))
+                .max()
+                .unwrap_or(0);
+            let tags_d = standard
+                .tags
+                .iter()
+                .map(|a| type_depth(&a))
+                .max()
+                .unwrap_or(0);
+            attrs_d.max(tags_d)
+        }
+        EntityDecl::Enum(_) => 0,
+    }
+}
+
+fn action_decl_type_depth(decl: &ActionDecl) -> usize {
+    let Some(app_decls) = &decl.app_decls else {
+        return 0;
+    };
+    app_decls
+        .node
+        .iter()
+        .map(|decl| app_decl_type_depth(&decl.node))
+        .max()
+        .unwrap_or(0)
+}
+
+fn app_decl_type_depth(decl: &AppDecl) -> usize {
+    match decl {
+        AppDecl::Context(Either::Right(attrs)) => attrs
+            .node
+            .iter()
+            .map(|a| 1 + type_depth(&a.node.data.ty))
+            .max()
+            .unwrap_or(0),
+        AppDecl::PR(_) | AppDecl::Context(Either::Left(_)) => 0,
+    }
+}
+
+fn type_depth(root: &Node<Type>) -> usize {
+    let mut stack: Vec<(&Node<Type>, usize)> = vec![(root, 0)];
+    let mut max_depth: usize = 0;
+
+    while let Some((ty, depth)) = stack.pop() {
+        max_depth = max_depth.max(depth);
+
+        match &ty.node {
+            Type::Ident(_) => {}
+            Type::Set(inner) => {
+                stack.push((inner, depth + 1));
+            }
+            Type::Record(attrs) => {
+                for attr_node in attrs {
+                    stack.push((&attr_node.node.data.ty, depth + 1));
+                }
+            }
+        }
+    }
+
+    max_depth
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::validator::cedar_schema::parser::parse_schema;
+    use rstest::rstest;
+
+    fn schema_depth_of(src: &str) -> usize {
+        let schema = parse_schema(src).expect("parse failed");
+        schema_type_depth(&schema)
+    }
+
+    #[rstest]
+    #[case::primitive_entity("entity Foo;", 0)]
+    #[case::entity_with_attribute("entity Foo = { bar: Long };", 1)]
+    #[case::nested_record("entity Foo = { bar: { baz: Long } };", 2)]
+    #[case::set_of_long("entity Foo = { bar: Set<Long> };", 2)]
+    #[case::set_of_set("entity Foo = { bar: Set<Set<Long>> };", 3)]
+    #[case::deeply_nested_record("entity Foo = { a: { b: { c: Long } } };", 3)]
+    #[case::common_type("type T = Set<Set<Long>>;", 2)]
+    #[case::entity_tags("entity Foo tags Set<Long>;", 1)]
+    #[case::entity_tags_nested("entity Foo tags Set<Set<Long>>;", 2)]
+    #[case::enum_entity(r#"entity Flavor enum ["Vanilla", "Chocolate"];"#, 0)]
+    #[case::action_context(
+        "entity User; entity File; action Read appliesTo { principal: User, resource: File, context: { level: Set<Long> } };",
+        2
+    )]
+    #[case::multiple_namespaces(
+        "namespace A { entity Foo = { x: Long }; } namespace B { entity Bar = { x: { y: { z: Long } } }; }",
+        3
+    )]
+    fn schema_depth(#[case] src: &str, #[case] expected: usize) {
+        assert_eq!(schema_depth_of(src), expected);
+    }
+}

cedar-policy-core/src/validator/cedar_schema/parser.rs

@@ -97,6 +97,14 @@ pub enum CedarSchemaParseErrors {
     #[error(transparent)]
     #[diagnostic(transparent)]
     JsonError(#[from] ToJsonSchemaErrors),
+    /// Schema type nesting depth exceeds the configured limit
+    #[error("schema type depth {depth} exceeds the configured limit of {limit}")]
+    TypeTooDeep {
+        /// Actual depth of the deepest type
+        depth: usize,
+        /// Configured limit
+        limit: usize,
+    },
 }
 
 /// Parse a schema fragment, in the Cedar syntax, into a [`json_schema::Fragment`],
@@ -116,6 +124,51 @@ pub fn parse_cedar_schema_fragment<'a>(
     Ok(tuple)
 }
 
+/// Like [`parse_cedar_schema_fragment`], but rejects the schema if any type's
+/// effective nesting depth exceeds `depth_limit`. Effective depth accounts for
+/// depth introduced through chains of common type (typedef) references.
+pub fn parse_cedar_schema_fragment_with_depth_limit<'a>(
+    src: &str,
+    extensions: &Extensions<'a>,
+    depth_limit: usize,
+) -> Result<
+    (
+        json_schema::Fragment<crate::validator::RawName>,
+        impl Iterator<Item = SchemaWarning> + 'a,
+    ),
+    CedarSchemaParseErrors,
+> {
+    let ast: Schema = parse_collect_errors(&*SCHEMA_PARSER, grammar::SchemaParser::parse, src)?;
+    let syntactic_depth = super::depth::schema_type_depth(&ast);
+    if syntactic_depth > depth_limit {
+        return Err(CedarSchemaParseErrors::TypeTooDeep {
+            depth: syntactic_depth,
+            limit: depth_limit,
+        });
+    }
+    // The syntactic check passed, so nothing directly exceeds the limit, but
+    // there might still be types composed from common types that will exceed it
+    // after inlining the type definitions.
+    let (fragment, warnings) = cedar_schema_to_json_schema(ast, extensions)?;
+    // Now check effective depth (after inlining common types) to ensure we
+    // avoid any downstream stackoverflow caused by recursion on the inlined
+    // types.  We first need to resolve `RawName` to `InternalName` to know
+    // exactly what each common type refers to when computing the depth.
+    if let Ok(resolved) = fragment.to_internal_name_fragment_with_resolved_types() {
+        if let Some(depth) =
+            crate::validator::schema_type_depth::fragment_effective_depth(&resolved)
+        {
+            if depth > depth_limit {
+                return Err(CedarSchemaParseErrors::TypeTooDeep {
+                    depth,
+                    limit: depth_limit,
+                });
+            }
+        }
+    }
+    Ok((fragment, warnings))
+}
+
 /// Parse schema from text
 pub fn parse_schema(text: &str) -> Result<Schema, err::ParseErrors> {
     parse_collect_errors(&*SCHEMA_PARSER, grammar::SchemaParser::parse, text)

cedar-policy-core/src/validator/cedar_schema/to_json_schema.rs

@@ -97,16 +97,28 @@ fn is_valid_ext_type(ty: &Id, extensions: &Extensions<'_>) -> bool {
 }
 
 /// Convert a `Type` into the JSON representation of the type.
-pub fn cedar_type_to_json_type(ty: Node<Type>) -> json_schema::Type<RawName> {
-    let variant = match ty.node {
+pub fn cedar_type_to_json_type(mut ty: Node<Type>) -> json_schema::Type<RawName> {
+    // Some anoying use of `std::mem::replace` here because we have `Type`
+    // implement `Drop` to avoid stack overflows when dropping deeply nested
+    // types and Rust won't let you move out of a type with `Drop`.
+    let variant = match &mut ty.node {
         Type::Set(t) => json_schema::TypeVariant::Set {
-            element: Box::new(cedar_type_to_json_type(*t)),
+            element: Box::new(cedar_type_to_json_type(std::mem::replace(
+                t.as_mut(),
+                Node::new(Type::Record(Vec::new())),
+            ))),
         },
         Type::Ident(p) => json_schema::TypeVariant::EntityOrCommon {
-            type_name: RawName::from(p),
+            type_name: RawName::from(std::mem::replace(
+                p,
+                Path::new(Id::new_unchecked_const(""), [], None),
+            )),
         },
         Type::Record(fields) => json_schema::TypeVariant::Record(json_schema::RecordType {
-            attributes: fields.into_iter().map(convert_attr_decl).collect(),
+            attributes: std::mem::take(fields)
+                .into_iter()
+                .map(convert_attr_decl)
+                .collect(),
             additional_attributes: false,
         }),
     };