Harden stack overflow protection (#966)

goto-bus-stop · web-flow · commit ebe09fed4713 · 2025-04-18T10:28:56.000+02:00
* Add failing tests for stack overflow

* Add overflow protection to variable validation

* You need a lot to blow the stack here

* Lint

* Add overflow protection in subscription validation

* Respect recursion limit in Fragment Is Used validation

Also speed it up significantly by only walking the entire AST once,
instead of doing it again for every fragment definition.

* lint

* comment tweaks

* Update snapshots
diff --git a/crates/apollo-compiler/src/executable/validation.rs b/crates/apollo-compiler/src/executable/validation.rs
@@ -1,5 +1,5 @@
 use super::FieldSet;
-use crate::validation::fragment::validate_fragment_used;
+use crate::validation::fragment::validate_fragments_used;
 use crate::validation::operation::validate_operation_definitions;
 use crate::validation::selection::FieldsInSetCanMerge;
 use crate::validation::DiagnosticList;
@@ -44,9 +44,7 @@ pub(crate) fn validate_with_or_without_schema(
 ) {
     let context = ExecutableValidationContext::new(schema);
     validate_operation_definitions(errors, document, &context);
-    for def in document.fragments.values() {
-        validate_fragment_used(errors, document, def);
-    }
+    validate_fragments_used(errors, document);
 }
 
 pub(crate) fn validate_field_set(
diff --git a/crates/apollo-compiler/src/validation/fragment.rs b/crates/apollo-compiler/src/validation/fragment.rs
@@ -7,10 +7,12 @@ use crate::executable;
 use crate::schema;
 use crate::schema::Implementers;
 use crate::validation::diagnostics::DiagnosticData;
+use crate::validation::variable::walk_selections_with_deduped_fragments;
 use crate::validation::CycleError;
 use crate::validation::DiagnosticList;
 use crate::validation::OperationValidationContext;
 use crate::validation::RecursionGuard;
+use crate::validation::RecursionLimitError;
 use crate::validation::RecursionStack;
 use crate::validation::SourceSpan;
 use crate::ExecutableDocument;
@@ -376,49 +378,40 @@ pub(crate) fn validate_fragment_type_condition(
     }
 }
 
-pub(crate) fn validate_fragment_used(
-    diagnostics: &mut DiagnosticList,
+fn collect_used_fragments(
     document: &ExecutableDocument,
-    fragment: &Node<executable::Fragment>,
-) {
-    let fragment_name = &fragment.name;
-
-    let mut all_selections = document
-        .operations
-        .iter()
-        .map(|operation| &operation.selection_set)
-        .chain(
-            document
-                .fragments
-                .values()
-                .map(|fragment| &fragment.selection_set),
-        )
-        .flat_map(|set| &set.selections);
-
-    let is_used = all_selections.any(|sel| selection_uses_fragment(sel, fragment_name));
-
-    // Fragments must be used within the schema
-    //
-    // Returns Unused Fragment error.
-    if !is_used {
-        diagnostics.push(
-            fragment.location(),
-            DiagnosticData::UnusedFragment {
-                name: fragment_name.clone(),
-            },
-        )
+) -> Result<HashSet<&Name>, RecursionLimitError> {
+    let mut names = HashSet::default();
+    for operation in document.operations.iter() {
+        walk_selections_with_deduped_fragments(document, &operation.selection_set, |selection| {
+            if let executable::Selection::FragmentSpread(spread) = selection {
+                names.insert(&spread.fragment_name);
+            }
+        })?;
     }
+    Ok(names)
 }
 
-fn selection_uses_fragment(sel: &executable::Selection, name: &str) -> bool {
-    let sub_selections = match sel {
-        executable::Selection::FragmentSpread(fragment) => return fragment.fragment_name == name,
-        executable::Selection::Field(field) => &field.selection_set,
-        executable::Selection::InlineFragment(inline) => &inline.selection_set,
+pub(crate) fn validate_fragments_used(
+    diagnostics: &mut DiagnosticList,
+    document: &ExecutableDocument,
+) {
+    let Ok(used_fragments) = collect_used_fragments(document) else {
+        diagnostics.push(None, super::Details::RecursionLimitError);
+        return;
     };
 
-    sub_selections
-        .selections
-        .iter()
-        .any(|sel| selection_uses_fragment(sel, name))
+    for fragment in document.fragments.values() {
+        // Fragments must be used within the schema
+        //
+        // Returns Unused Fragment error.
+        if !used_fragments.contains(&fragment.name) {
+            diagnostics.push(
+                fragment.location(),
+                DiagnosticData::UnusedFragment {
+                    name: fragment.name.clone(),
+                },
+            )
+        }
+    }
 }
diff --git a/crates/apollo-compiler/src/validation/mod.rs b/crates/apollo-compiler/src/validation/mod.rs
@@ -1089,6 +1089,60 @@ const DEFAULT_RECURSION_LIMIT: usize = 32;
 #[non_exhaustive]
 struct RecursionLimitError {}
 
+/// Track recursion depth to prevent stack overflow.
+#[derive(Debug)]
+struct DepthCounter {
+    value: usize,
+    high: usize,
+    limit: usize,
+}
+
+impl DepthCounter {
+    fn new() -> Self {
+        Self {
+            value: 0,
+            high: 0,
+            limit: DEFAULT_RECURSION_LIMIT,
+        }
+    }
+
+    fn with_limit(mut self, limit: usize) -> Self {
+        self.limit = limit;
+        self
+    }
+
+    /// Return the actual API for tracking recursive uses.
+    pub(crate) fn guard(&mut self) -> DepthGuard<'_> {
+        DepthGuard(self)
+    }
+}
+
+/// Track call depth in a recursive function.
+///
+/// Pass the result of `guard.increment()` to recursive calls. When a guard is dropped,
+/// its value is decremented.
+struct DepthGuard<'a>(&'a mut DepthCounter);
+
+impl DepthGuard<'_> {
+    /// Mark that we are recursing. If we reached the limit, return an error.
+    fn increment(&mut self) -> Result<DepthGuard<'_>, RecursionLimitError> {
+        self.0.value += 1;
+        self.0.high = self.0.high.max(self.0.value);
+        if self.0.value > self.0.limit {
+            Err(RecursionLimitError {})
+        } else {
+            Ok(DepthGuard(self.0))
+        }
+    }
+}
+
+impl Drop for DepthGuard<'_> {
+    fn drop(&mut self) {
+        // This may already be 0 if it's the original `counter.guard()` result, but that's fine
+        self.0.value = self.0.value.saturating_sub(1);
+    }
+}
+
 /// Track used names in a recursive function.
 #[derive(Debug)]
 struct RecursionStack {
diff --git a/crates/apollo-compiler/src/validation/operation.rs b/crates/apollo-compiler/src/validation/operation.rs
@@ -1,6 +1,8 @@
 use crate::collections::HashSet;
 use crate::executable;
 use crate::validation::diagnostics::DiagnosticData;
+use crate::validation::DepthCounter;
+use crate::validation::DepthGuard;
 use crate::validation::DiagnosticList;
 use crate::validation::ExecutableValidationContext;
 use crate::validation::RecursionLimitError;
@@ -23,6 +25,7 @@ fn walk_selections<'doc>(
         document: &'doc ExecutableDocument,
         selection_set: &'doc executable::SelectionSet,
         seen: &mut HashSet<&'doc Name>,
+        mut guard: DepthGuard<'_>,
         f: &mut dyn FnMut(&'doc executable::Selection),
     ) -> Result<(), RecursionLimitError> {
         for selection in &selection_set.selections {
@@ -45,19 +48,38 @@ fn walk_selections<'doc>(
                             document,
                             &fragment_definition.selection_set,
                             seen,
+                            guard.increment()?,
                             f,
                         )?;
                     }
                 }
                 executable::Selection::InlineFragment(fragment) => {
-                    walk_selections_inner(document, &fragment.selection_set, seen, f)?;
+                    walk_selections_inner(
+                        document,
+                        &fragment.selection_set,
+                        seen,
+                        guard.increment()?,
+                        f,
+                    )?;
                 }
             }
         }
         Ok(())
     }
 
-    walk_selections_inner(document, selections, &mut HashSet::default(), &mut f)
+    // This has a much higher limit than comparable recursive walks, like the one in
+    // `validate_fragment_cycles`, despite doing similar work. This is because this limit
+    // was introduced later and should not break (reasonable) existing queries that are
+    // under that pre-existing limit. Luckily the existing limit was very conservative.
+    let mut depth = DepthCounter::new().with_limit(500);
+
+    walk_selections_inner(
+        document,
+        selections,
+        &mut HashSet::default(),
+        depth.guard(),
+        &mut f,
+    )
 }
 
 pub(crate) fn validate_subscription(
diff --git a/crates/apollo-compiler/src/validation/variable.rs b/crates/apollo-compiler/src/validation/variable.rs
@@ -4,6 +4,8 @@ use crate::collections::HashSet;
 use crate::executable;
 use crate::validation::diagnostics::DiagnosticData;
 use crate::validation::value::value_of_correct_type;
+use crate::validation::DepthCounter;
+use crate::validation::DepthGuard;
 use crate::validation::DiagnosticList;
 use crate::validation::RecursionLimitError;
 use crate::validation::SourceSpan;
@@ -88,7 +90,7 @@ pub(crate) fn validate_variable_definitions(
 ///
 /// Named fragments are "deduplicated": only visited once even if spread multiple times *in
 /// different locations*. This is only appropriate for certain kinds of validations, so reuser beware.
-fn walk_selections_with_deduped_fragments<'doc>(
+pub(super) fn walk_selections_with_deduped_fragments<'doc>(
     document: &'doc ExecutableDocument,
     selections: &'doc executable::SelectionSet,
     mut f: impl FnMut(&'doc executable::Selection),
@@ -97,13 +99,20 @@ fn walk_selections_with_deduped_fragments<'doc>(
         document: &'doc ExecutableDocument,
         selection_set: &'doc executable::SelectionSet,
         seen: &mut HashSet<&'doc Name>,
+        mut guard: DepthGuard<'_>,
         f: &mut dyn FnMut(&'doc executable::Selection),
     ) -> Result<(), RecursionLimitError> {
         for selection in &selection_set.selections {
             f(selection);
             match selection {
                 executable::Selection::Field(field) => {
-                    walk_selections_inner(document, &field.selection_set, seen, f)?;
+                    walk_selections_inner(
+                        document,
+                        &field.selection_set,
+                        seen,
+                        guard.increment()?,
+                        f,
+                    )?;
                 }
                 executable::Selection::FragmentSpread(fragment) => {
                     let new = seen.insert(&fragment.fragment_name);
@@ -118,19 +127,37 @@ fn walk_selections_with_deduped_fragments<'doc>(
                             document,
                             &fragment_definition.selection_set,
                             seen,
+                            guard.increment()?,
                             f,
                         )?;
                     }
                 }
                 executable::Selection::InlineFragment(fragment) => {
-                    walk_selections_inner(document, &fragment.selection_set, seen, f)?;
+                    walk_selections_inner(
+                        document,
+                        &fragment.selection_set,
+                        seen,
+                        guard.increment()?,
+                        f,
+                    )?;
                 }
             }
         }
         Ok(())
     }
 
-    walk_selections_inner(document, selections, &mut HashSet::default(), &mut f)
+    // This has a much higher limit than comparable recursive walks, like the one in
+    // `validate_fragment_cycles`, despite doing similar work. This is because this limit
+    // was introduced later and should not break (reasonable) existing queries that are
+    // under that pre-existing limit. Luckily the existing limit was very conservative.
+    let mut depth = DepthCounter::new().with_limit(500);
+    walk_selections_inner(
+        document,
+        selections,
+        &mut HashSet::default(),
+        depth.guard(),
+        &mut f,
+    )
 }
 
 fn variables_in_value(value: &ast::Value) -> impl Iterator<Item = &Name> + '_ {
diff --git a/crates/apollo-compiler/tests/validation/recursion.rs b/crates/apollo-compiler/tests/validation/recursion.rs
