Add limit when recursively walking selection set

goto-bus-stop · goto-bus-stop · commit c7e90f085ca9 · 2023-11-17T16:09:31.000+01:00
diff --git a/crates/apollo-compiler/src/diagnostics.rs b/crates/apollo-compiler/src/diagnostics.rs
@@ -244,6 +244,8 @@ pub(crate) enum DiagnosticData {
         /// Name of the argument where variable is used
         arg_name: String,
     },
+    #[error("too much recursion")]
+    RecursionError {},
 }
 
 impl DiagnosticData {
diff --git a/crates/apollo-compiler/src/validation/fragment.rs b/crates/apollo-compiler/src/validation/fragment.rs
@@ -316,7 +316,6 @@ pub(crate) fn validate_fragment_cycles(
                     }
 
                     if let Some(fragment) = named_fragments.get(&spread.fragment_name) {
-                        println!("going into {}", spread.fragment_name);
                         detect_fragment_cycles(
                             named_fragments,
                             &fragment.selection_set,
diff --git a/crates/apollo-compiler/src/validation/variable.rs b/crates/apollo-compiler/src/validation/variable.rs
@@ -1,11 +1,15 @@
 use crate::{
     ast,
     diagnostics::{ApolloDiagnostic, DiagnosticData, Label},
-    schema, FileId, Node, NodeLocation, ValidationDatabase,
+    schema,
+    validation::{RecursionGuard, RecursionStack},
+    FileId, Node, NodeLocation, ValidationDatabase,
 };
 use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet};
 
+use super::RecursionLimitError;
+
 pub(crate) fn validate_variable_definitions2(
     db: &dyn ValidationDatabase,
     variables: &[Node<ast::VariableDefinition>],
@@ -107,16 +111,11 @@ pub(crate) fn validate_variable_definitions2(
     diagnostics
 }
 
-enum SelectionPathElement<'ast> {
-    Field(&'ast ast::Field),
-    Fragment(&'ast ast::FragmentSpread),
-    InlineFragment(&'ast ast::InlineFragment),
-}
 fn walk_selections(
     document: &ast::Document,
     selections: &[ast::Selection],
     mut f: impl FnMut(&ast::Selection),
-) {
+) -> Result<(), RecursionLimitError> {
     type NamedFragments = HashMap<ast::Name, Node<ast::FragmentDefinition>>;
     let named_fragments: NamedFragments = document
         .definitions
@@ -129,55 +128,48 @@ fn walk_selections(
         })
         .collect();
 
-    fn walk_selections_inner<'ast: 'path, 'path>(
+    fn walk_selections_inner<'ast, 'guard>(
         named_fragments: &'ast NamedFragments,
         selections: &'ast [ast::Selection],
-        path: &mut Vec<SelectionPathElement<'path>>,
+        guard: &mut RecursionGuard<'guard>,
         f: &mut dyn FnMut(&ast::Selection),
-    ) {
+    ) -> Result<(), RecursionLimitError> {
         for selection in selections {
             f(selection);
             match selection {
                 ast::Selection::Field(field) => {
-                    path.push(SelectionPathElement::Field(field));
-                    walk_selections_inner(named_fragments, &field.selection_set, path, f);
-                    path.pop();
+                    walk_selections_inner(named_fragments, &field.selection_set, guard, f)?;
                 }
                 ast::Selection::FragmentSpread(fragment) => {
-                    // prevent recursion
-                    if path.iter().any(|element| {
-                        matches!(element, SelectionPathElement::Fragment(walked) if walked.fragment_name == fragment.fragment_name)
-                    }) {
+                    // Prevent chasing a cyclical reference.
+                    // Note we do not report `CycleError::Recursed` here, as that is already caught
+                    // by the cyclical fragment validation--we just need to ensure that we don't
+                    // overflow the stack.
+                    if guard.contains(&fragment.fragment_name) {
                         continue;
                     }
 
-                    path.push(SelectionPathElement::Fragment(fragment));
                     if let Some(fragment_definition) = named_fragments.get(&fragment.fragment_name)
                     {
                         walk_selections_inner(
                             named_fragments,
                             &fragment_definition.selection_set,
-                            path,
+                            &mut guard.push(&fragment.fragment_name)?,
                             f,
-                        );
+                        )?;
                     }
-                    path.pop();
                 }
                 ast::Selection::InlineFragment(fragment) => {
-                    path.push(SelectionPathElement::InlineFragment(fragment));
-                    walk_selections_inner(named_fragments, &fragment.selection_set, path, f);
-                    path.pop();
+                    walk_selections_inner(named_fragments, &fragment.selection_set, guard, f)?;
                 }
             }
         }
+        Ok(())
     }
 
-    walk_selections_inner(
-        &named_fragments,
-        selections,
-        &mut Default::default(),
-        &mut f,
-    );
+    let mut stack = RecursionStack::new(500);
+    let result = walk_selections_inner(&named_fragments, selections, &mut stack.guard(), &mut f);
+    result
 }
 
 fn variables_in_value(value: &ast::Value) -> impl Iterator<Item = ast::Name> + '_ {
@@ -220,6 +212,8 @@ pub(crate) fn validate_unused_variables(
     file_id: FileId,
     operation: Node<ast::OperationDefinition>,
 ) -> Vec<ApolloDiagnostic> {
+    let mut diagnostics = Vec::new();
+
     let defined_vars: HashSet<_> = operation
         .variables
         .iter()
@@ -236,24 +230,31 @@ pub(crate) fn validate_unused_variables(
         })
         .collect();
     let mut used_vars = HashSet::<ast::Name>::new();
-    walk_selections(
-        &db.ast(file_id),
-        &operation.selection_set,
-        |selection| match selection {
-            ast::Selection::Field(field) => {
-                used_vars.extend(variables_in_directives(&field.directives));
-                used_vars.extend(variables_in_arguments(&field.arguments));
-            }
-            ast::Selection::FragmentSpread(fragment) => {
-                used_vars.extend(variables_in_directives(&fragment.directives));
-            }
-            ast::Selection::InlineFragment(fragment) => {
-                used_vars.extend(variables_in_directives(&fragment.directives));
-            }
-        },
-    );
-
-    let mut diagnostics = Vec::new();
+    let walked =
+        walk_selections(
+            &db.ast(file_id),
+            &operation.selection_set,
+            |selection| match selection {
+                ast::Selection::Field(field) => {
+                    used_vars.extend(variables_in_directives(&field.directives));
+                    used_vars.extend(variables_in_arguments(&field.arguments));
+                }
+                ast::Selection::FragmentSpread(fragment) => {
+                    used_vars.extend(variables_in_directives(&fragment.directives));
+                }
+                ast::Selection::InlineFragment(fragment) => {
+                    used_vars.extend(variables_in_directives(&fragment.directives));
+                }
+            },
+        );
+    if walked.is_err() {
+        diagnostics.push(ApolloDiagnostic::new(
+            db,
+            None,
+            DiagnosticData::RecursionError {},
+        ));
+        return diagnostics;
+    }
 
     let unused_vars = defined_vars.difference(&used_vars);
 
diff --git a/crates/apollo-compiler/test_data/diagnostics/0070_self_referential_directive_definition.txt b/crates/apollo-compiler/test_data/diagnostics/0070_self_referential_directive_definition.txt
@@ -5,14 +5,14 @@ Error: `invalidExample` directive definition cannot reference itself
    │ ────────────┬────────────             ───────┬───────  
    │             ╰────────────────────────────────────────── recursive directive definition
    │                                              │         
-   │                                              ╰───────── refers to itself here
+   │                                              ╰───────── `invalidExample` circularly references `invalidExample` here
 ───╯
 Error: `deprecatedType` directive definition cannot reference itself
     ╭─[0070_self_referential_directive_definition.graphql:11:1]
     │
  10 │ extend scalar String @deprecatedType(reason: "use OurCustomString instead")
     │                      ───────────────────────────┬──────────────────────────  
-    │                                                 ╰──────────────────────────── refers to itself here
+    │                                                 ╰──────────────────────────── `deprecatedType` circularly references `deprecatedType` here
  11 │ directive @deprecatedType(reason: String!) on OBJECT | INTERFACE | ENUM | SCALAR | UNION
     │ ────────────┬────────────  
     │             ╰────────────── recursive directive definition
@@ -21,39 +21,53 @@ Error: `loopA` directive definition cannot reference itself
     ╭─[0070_self_referential_directive_definition.graphql:13:1]
     │
  13 │ directive @loopA(arg: Boolean @loopB) on ARGUMENT_DEFINITION
-    │ ────────┬───────  
-    │         ╰───────── recursive directive definition
-    │ 
+    │ ────────┬───────              ───┬──  
+    │         ╰───────────────────────────── recursive directive definition
+    │                                  │    
+    │                                  ╰──── `loopA` references `loopB` here...
+ 14 │ directive @loopB(arg: Boolean @loopC) on ARGUMENT_DEFINITION
+    │                               ───┬──  
+    │                                  ╰──── `loopB` references `loopC` here...
  15 │ directive @loopC(arg: Boolean @loopA) on ARGUMENT_DEFINITION
     │                               ───┬──  
-    │                                  ╰──── refers to itself here
+    │                                  ╰──── `loopC` circularly references `loopA` here
 ────╯
 Error: `loopB` directive definition cannot reference itself
     ╭─[0070_self_referential_directive_definition.graphql:14:1]
     │
  13 │ directive @loopA(arg: Boolean @loopB) on ARGUMENT_DEFINITION
     │                               ───┬──  
-    │                                  ╰──── refers to itself here
+    │                                  ╰──── `loopA` circularly references `loopB` here
  14 │ directive @loopB(arg: Boolean @loopC) on ARGUMENT_DEFINITION
-    │ ────────┬───────  
-    │         ╰───────── recursive directive definition
+    │ ────────┬───────              ───┬──  
+    │         ╰───────────────────────────── recursive directive definition
+    │                                  │    
+    │                                  ╰──── `loopB` references `loopC` here...
+ 15 │ directive @loopC(arg: Boolean @loopA) on ARGUMENT_DEFINITION
+    │                               ───┬──  
+    │                                  ╰──── `loopC` references `loopA` here...
 ────╯
 Error: `loopC` directive definition cannot reference itself
     ╭─[0070_self_referential_directive_definition.graphql:15:1]
     │
+ 13 │ directive @loopA(arg: Boolean @loopB) on ARGUMENT_DEFINITION
+    │                               ───┬──  
+    │                                  ╰──── `loopA` references `loopB` here...
  14 │ directive @loopB(arg: Boolean @loopC) on ARGUMENT_DEFINITION
     │                               ───┬──  
-    │                                  ╰──── refers to itself here
+    │                                  ╰──── `loopB` circularly references `loopC` here
  15 │ directive @loopC(arg: Boolean @loopA) on ARGUMENT_DEFINITION
-    │ ────────┬───────  
-    │         ╰───────── recursive directive definition
+    │ ────────┬───────              ───┬──  
+    │         ╰───────────────────────────── recursive directive definition
+    │                                  │    
+    │                                  ╰──── `loopC` references `loopA` here...
 ────╯
 Error: `wrong` directive definition cannot reference itself
     ╭─[0070_self_referential_directive_definition.graphql:21:1]
     │
  18 │ extend enum Enum { RECURSIVE @wrong }
     │                              ───┬──  
-    │                                 ╰──── refers to itself here
+    │                                 ╰──── `wrong` circularly references `wrong` here
     │ 
  21 │ directive @wrong(input: InputObject) on INPUT_FIELD_DEFINITION | ENUM_VALUE
     │ ────────┬───────  
diff --git a/crates/apollo-compiler/test_data/diagnostics/0084_circular_non_nullable_input_objects.txt b/crates/apollo-compiler/test_data/diagnostics/0084_circular_non_nullable_input_objects.txt
@@ -2,50 +2,93 @@ Error: `First` input object cannot reference itself
     ╭─[0084_circular_non_nullable_input_objects.graphql:6:1]
     │
   6 │ ╭─▶ input First {
+  7 │ │     second: Second!
+    │ │     ───────┬───────  
+    │ │            ╰───────── `First` references `second` here...
     ┆ ┆   
   9 │ ├─▶ }
     │ │       
     │ ╰─────── cyclical input object definition
     │ 
+ 12 │       third: Third!
+    │       ──────┬──────  
+    │             ╰──────── `second` references `third` here...
+    │ 
+ 17 │       fourth: Fourth!
+    │       ───────┬───────  
+    │              ╰───────── `third` references `fourth` here...
+    │ 
  22 │       first: First!
     │       ──────┬──────  
-    │             ╰──────── refers to itself here
+    │             ╰──────── `fourth` circularly references `first` here
 ────╯
 Error: `Second` input object cannot reference itself
     ╭─[0084_circular_non_nullable_input_objects.graphql:11:1]
     │
   7 │       second: Second!
     │       ───────┬───────  
-    │              ╰───────── refers to itself here
+    │              ╰───────── `first` circularly references `second` here
     │ 
  11 │ ╭─▶ input Second {
+ 12 │ │     third: Third!
+    │ │     ──────┬──────  
+    │ │           ╰──────── `Second` references `third` here...
     ┆ ┆   
  14 │ ├─▶ }
     │ │       
     │ ╰─────── cyclical input object definition
+    │ 
+ 17 │       fourth: Fourth!
+    │       ───────┬───────  
+    │              ╰───────── `third` references `fourth` here...
+    │ 
+ 22 │       first: First!
+    │       ──────┬──────  
+    │             ╰──────── `fourth` references `first` here...
 ────╯
 Error: `Third` input object cannot reference itself
     ╭─[0084_circular_non_nullable_input_objects.graphql:16:1]
     │
+  7 │       second: Second!
+    │       ───────┬───────  
+    │              ╰───────── `first` references `second` here...
+    │ 
  12 │       third: Third!
     │       ──────┬──────  
-    │             ╰──────── refers to itself here
+    │             ╰──────── `second` circularly references `third` here
     │ 
  16 │ ╭─▶ input Third {
+ 17 │ │     fourth: Fourth!
+    │ │     ───────┬───────  
+    │ │            ╰───────── `Third` references `fourth` here...
     ┆ ┆   
  19 │ ├─▶ }
     │ │       
     │ ╰─────── cyclical input object definition
+    │ 
+ 22 │       first: First!
+    │       ──────┬──────  
+    │             ╰──────── `fourth` references `first` here...
 ────╯
 Error: `Fourth` input object cannot reference itself
     ╭─[0084_circular_non_nullable_input_objects.graphql:21:1]
     │
+  7 │       second: Second!
+    │       ───────┬───────  
+    │              ╰───────── `first` references `second` here...
+    │ 
+ 12 │       third: Third!
+    │       ──────┬──────  
+    │             ╰──────── `second` references `third` here...
+    │ 
  17 │       fourth: Fourth!
     │       ───────┬───────  
-    │              ╰───────── refers to itself here
+    │              ╰───────── `third` circularly references `fourth` here
     │ 
  21 │ ╭─▶ input Fourth {
-    ┆ ┆   
+ 22 │ │     first: First!
+    │ │     ──────┬──────  
+    │ │           ╰──────── `Fourth` references `first` here...
  23 │ ├─▶ }
     │ │      
     │ ╰────── cyclical input object definition
diff --git a/crates/apollo-compiler/tests/validation/fragments.rs b/crates/apollo-compiler/tests/validation/fragments.rs
@@ -13,7 +13,7 @@ fn long_fragment_chains_do_not_overflow_stack() {
     "#
     .to_string();
 
-    let fragments: usize = 10_000;
+    let fragments: usize = 1_000;
     for i in 1..fragments {
         query.push_str(&format!(
             "
@@ -41,5 +41,20 @@ fn long_fragment_chains_do_not_overflow_stack() {
         ),
         "overflow.graphql",
     );
-    executable.validate(&schema).unwrap();
+
+    let errors = executable
+        .validate(&schema)
+        .expect_err("must have recursion errors");
+
+    let expected = expect_test::expect![[r#"
+        Error: too much recursion
+        Error: `typeFragment1` fragment cannot reference itself
+            ╭─[overflow.graphql:11:11]
+            │
+         11 │           fragment typeFragment1 on __Type {
+            │           ───────────┬──────────  
+            │                      ╰──────────── recursive fragment definition
+        ────╯
+    "#]];
+    expected.assert_eq(&errors.to_string_no_color());
 }

Original file line number	Diff line number	Diff line change
`@@ -244,6 +244,8 @@ pub(crate) enum DiagnosticData {`
`244`	`244`	`/// Name of the argument where variable is used`
`245`	`245`	`arg_name: String,`
`246`	`246`	`},`
	`247`	`+ #[error("too much recursion")]`
	`248`	`+ RecursionError {},`
`247`	`249`	`}`
`248`	`250`
`249`	`251`	`impl DiagnosticData {`
Original file line number	Diff line number	Diff line change
`@@ -316,7 +316,6 @@ pub(crate) fn validate_fragment_cycles(`
`316`	`316`	`}`
`317`	`317`
`318`	`318`	`if let Some(fragment) = named_fragments.get(&spread.fragment_name) {`
`319`		`- println!("going into {}", spread.fragment_name);`
`320`	`319`	`detect_fragment_cycles(`
`321`	`320`	`named_fragments,`
`322`	`321`	`&fragment.selection_set,`