From 07c1a2e684d01872ecbc6259bb6bdc19650e78e5 Mon Sep 17 00:00:00 2001
From: Sunli <scott_s829@163.com>
Date: Thu, 21 Jul 2022 10:14:40 +0800
Subject: [PATCH] Limit recursive depth to 256 by default

---
 src/schema.rs | 96 +++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 89 insertions(+), 7 deletions(-)

diff --git a/src/schema.rs b/src/schema.rs
index 1af2e037..33727652 100644
--- a/src/schema.rs
+++ b/src/schema.rs
@@ -15,7 +15,10 @@ use crate::{
     model::__DirectiveLocation,
     parser::{
         parse_query,
-        types::{Directive, DocumentOperations, OperationType, Selection, SelectionSet},
+        types::{
+            Directive, DocumentOperations, ExecutableDocument, OperationType, Selection,
+            SelectionSet,
+        },
         Positioned,
     },
     registry::{MetaDirective, MetaInputValue, Registry, SDLExportOptions},
@@ -24,8 +27,8 @@ use crate::{
     types::QueryRoot,
     validation::{check_rules, ValidationMode},
     BatchRequest, BatchResponse, CacheControl, ContextBase, EmptyMutation, EmptySubscription,
-    InputType, ObjectType, OutputType, QueryEnv, Request, Response, ServerError, SubscriptionType,
-    Variables, ID,
+    InputType, ObjectType, OutputType, QueryEnv, Request, Response, ServerError, ServerResult,
+    SubscriptionType, Variables, ID,
 };
 
 /// Introspection mode
@@ -55,6 +58,7 @@ pub struct SchemaBuilder<Query, Mutation, Subscription> {
     data: Data,
     complexity: Option<usize>,
     depth: Option<usize>,
+    recursive_depth: usize,
     extensions: Vec<Box<dyn ExtensionFactory>>,
     custom_directives: HashMap<&'static str, Box<dyn CustomDirectiveFactory>>,
 }
@@ -110,6 +114,16 @@ impl<Query, Mutation, Subscription> SchemaBuilder<Query, Mutation, Subscription>
         self
     }
 
+    /// Set the maximum recursive depth a query can have. (default: 256)
+    ///
+    /// If the value is too large, stack overflow may occur, usually `256` is
+    /// enough.
+    #[must_use]
+    pub fn limit_recursive_depth(mut self, depth: usize) -> Self {
+        self.recursive_depth = depth;
+        self
+    }
+
     /// Add an extension to the schema.
     ///
     /// # Examples
@@ -219,6 +233,7 @@ impl<Query, Mutation, Subscription> SchemaBuilder<Query, Mutation, Subscription>
             subscription: self.subscription,
             complexity: self.complexity,
             depth: self.depth,
+            recursive_depth: self.recursive_depth,
             extensions: self.extensions,
             env: SchemaEnv(Arc::new(SchemaEnvInner {
                 registry: self.registry,
@@ -256,6 +271,7 @@ pub struct SchemaInner<Query, Mutation, Subscription> {
     pub(crate) subscription: Subscription,
     pub(crate) complexity: Option<usize>,
     pub(crate) depth: Option<usize>,
+    pub(crate) recursive_depth: usize,
     pub(crate) extensions: Vec<Box<dyn ExtensionFactory>>,
     pub(crate) env: SchemaEnv,
 }
@@ -339,6 +355,7 @@ where
             data: Default::default(),
             complexity: None,
             depth: None,
+            recursive_depth: 256,
             extensions: Default::default(),
             custom_directives: Default::default(),
         }
@@ -490,11 +507,14 @@ where
         let mut document = {
             let query = &request.query;
             let parsed_doc = request.parsed_query.take();
+            let recursive_depth = self.recursive_depth;
             let fut_parse = async move {
-                match parsed_doc {
-                    Some(parsed_doc) => Ok(parsed_doc),
-                    None => parse_query(query).map_err(Into::into),
-                }
+                let doc = match parsed_doc {
+                    Some(parsed_doc) => parsed_doc,
+                    None => parse_query(query)?,
+                };
+                check_recursive_depth(&doc, recursive_depth)?;
+                Ok(doc)
             };
             futures_util::pin_mut!(fut_parse);
             extensions
@@ -771,3 +791,65 @@ fn remove_skipped_selection(selection_set: &mut SelectionSet, variables: &Variab
         }
     }
 }
+
+fn check_recursive_depth(doc: &ExecutableDocument, max_depth: usize) -> ServerResult<()> {
+    fn check_selection_set(
+        doc: &ExecutableDocument,
+        selection_set: &Positioned<SelectionSet>,
+        current_depth: usize,
+        max_depth: usize,
+    ) -> ServerResult<()> {
+        if current_depth > max_depth {
+            return Err(ServerError::new(
+                format!(
+                    "The recursion depth of the query cannot be greater than `{}`",
+                    max_depth
+                ),
+                Some(selection_set.pos),
+            ));
+        }
+
+        for selection in &selection_set.node.items {
+            match &selection.node {
+                Selection::Field(field) => {
+                    if !field.node.selection_set.node.items.is_empty() {
+                        check_selection_set(
+                            doc,
+                            &field.node.selection_set,
+                            current_depth + 1,
+                            max_depth,
+                        )?;
+                    }
+                }
+                Selection::FragmentSpread(fragment_spread) => {
+                    if let Some(fragment) =
+                        doc.fragments.get(&fragment_spread.node.fragment_name.node)
+                    {
+                        check_selection_set(
+                            doc,
+                            &fragment.node.selection_set,
+                            current_depth + 1,
+                            max_depth,
+                        )?;
+                    }
+                }
+                Selection::InlineFragment(inline_fragment) => {
+                    check_selection_set(
+                        doc,
+                        &inline_fragment.node.selection_set,
+                        current_depth + 1,
+                        max_depth,
+                    )?;
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    for (_, operation) in doc.operations.iter() {
+        check_selection_set(doc, &operation.node.selection_set, 0, max_depth)?;
+    }
+
+    Ok(())
+}