parser: limit maximum number of tokens

Replicates graphql/graphql-js@9df9079

Author: Cito

src/graphql/language/parser.py

@@ -84,6 +84,7 @@
 def parse(
     source: SourceType,
     no_location: bool = False,
+    max_tokens: Optional[int] = None,
     allow_legacy_fragment_variables: bool = False,
     experimental_client_controlled_nullability: bool = False,
 ) -> DocumentNode:
@@ -95,6 +96,12 @@ def parse(
     they correspond to. The ``no_location`` option disables that behavior for
     performance or testing.
 
+    Parser CPU and memory usage is linear to the number of tokens in a document,
+    however in extreme cases it becomes quadratic due to memory exhaustion.
+    Parsing happens before validation so even invalid queries can burn lots of
+    CPU time and memory.
+    To prevent this you can set a maximum number of tokens allowed within a document.
+
     Legacy feature (will be removed in v3.3):
 
     If ``allow_legacy_fragment_variables`` is set to ``True``, the parser will
@@ -131,6 +138,7 @@ def parse(
     parser = Parser(
         source,
         no_location=no_location,
+        max_tokens=max_tokens,
         allow_legacy_fragment_variables=allow_legacy_fragment_variables,
         experimental_client_controlled_nullability=experimental_client_controlled_nullability,  # noqa
     )
@@ -140,6 +148,7 @@ def parse(
 def parse_value(
     source: SourceType,
     no_location: bool = False,
+    max_tokens: Optional[int] = None,
     allow_legacy_fragment_variables: bool = False,
 ) -> ValueNode:
     """Parse the AST for a given string containing a GraphQL value.
@@ -155,6 +164,7 @@ def parse_value(
     parser = Parser(
         source,
         no_location=no_location,
+        max_tokens=max_tokens,
         allow_legacy_fragment_variables=allow_legacy_fragment_variables,
     )
     parser.expect_token(TokenKind.SOF)
@@ -166,6 +176,7 @@ def parse_value(
 def parse_const_value(
     source: SourceType,
     no_location: bool = False,
+    max_tokens: Optional[int] = None,
     allow_legacy_fragment_variables: bool = False,
 ) -> ConstValueNode:
     """Parse the AST for a given string containing a GraphQL constant value.
@@ -176,6 +187,7 @@ def parse_const_value(
     parser = Parser(
         source,
         no_location=no_location,
+        max_tokens=max_tokens,
         allow_legacy_fragment_variables=allow_legacy_fragment_variables,
     )
     parser.expect_token(TokenKind.SOF)
@@ -187,6 +199,7 @@ def parse_const_value(
 def parse_type(
     source: SourceType,
     no_location: bool = False,
+    max_tokens: Optional[int] = None,
     allow_legacy_fragment_variables: bool = False,
 ) -> TypeNode:
     """Parse the AST for a given string containing a GraphQL Type.
@@ -202,6 +215,7 @@ def parse_type(
     parser = Parser(
         source,
         no_location=no_location,
+        max_tokens=max_tokens,
         allow_legacy_fragment_variables=allow_legacy_fragment_variables,
     )
     parser.expect_token(TokenKind.SOF)
@@ -222,27 +236,32 @@ class Parser:
     library, please use the `__version_info__` variable for version detection.
     """
 
-    _lexer: Lexer
     _no_location: bool
+    _max_tokens: Optional[int]
     _allow_legacy_fragment_variables: bool
     _experimental_client_controlled_nullability: bool
+    _lexer: Lexer
+    _token_counter: int
 
     def __init__(
         self,
         source: SourceType,
         no_location: bool = False,
+        max_tokens: Optional[int] = None,
         allow_legacy_fragment_variables: bool = False,
         experimental_client_controlled_nullability: bool = False,
     ):
         if not is_source(source):
             source = Source(cast(str, source))
 
-        self._lexer = Lexer(source)
         self._no_location = no_location
+        self._max_tokens = max_tokens
         self._allow_legacy_fragment_variables = allow_legacy_fragment_variables
         self._experimental_client_controlled_nullability = (
             experimental_client_controlled_nullability
         )
+        self._lexer = Lexer(source)
+        self._token_counter = 0
 
     def parse_name(self) -> NameNode:
         """Convert a name lex token into a name parse node."""
@@ -546,7 +565,7 @@ def parse_value_literal(self, is_const: bool) -> ValueNode:
 
     def parse_string_literal(self, _is_const: bool = False) -> StringValueNode:
         token = self._lexer.token
-        self._lexer.advance()
+        self.advance_lexer()
         return StringValueNode(
             value=token.value,
             block=token.kind == TokenKind.BLOCK_STRING,
@@ -583,18 +602,18 @@ def parse_object(self, is_const: bool) -> ObjectValueNode:
 
     def parse_int(self, _is_const: bool = False) -> IntValueNode:
         token = self._lexer.token
-        self._lexer.advance()
+        self.advance_lexer()
         return IntValueNode(value=token.value, loc=self.loc(token))
 
     def parse_float(self, _is_const: bool = False) -> FloatValueNode:
         token = self._lexer.token
-        self._lexer.advance()
+        self.advance_lexer()
         return FloatValueNode(value=token.value, loc=self.loc(token))
 
     def parse_named_values(self, _is_const: bool = False) -> ValueNode:
         token = self._lexer.token
         value = token.value
-        self._lexer.advance()
+        self.advance_lexer()
         if value == "true":
             return BooleanValueNode(value=True, loc=self.loc(token))
         if value == "false":
@@ -1089,7 +1108,7 @@ def expect_token(self, kind: TokenKind) -> Token:
         """
         token = self._lexer.token
         if token.kind == kind:
-            self._lexer.advance()
+            self.advance_lexer()
             return token
 
         raise GraphQLSyntaxError(
@@ -1106,7 +1125,7 @@ def expect_optional_token(self, kind: TokenKind) -> bool:
         """
         token = self._lexer.token
         if token.kind == kind:
-            self._lexer.advance()
+            self.advance_lexer()
             return True
 
         return False
@@ -1119,7 +1138,7 @@ def expect_keyword(self, value: str) -> None:
         """
         token = self._lexer.token
         if token.kind == TokenKind.NAME and token.value == value:
-            self._lexer.advance()
+            self.advance_lexer()
         else:
             raise GraphQLSyntaxError(
                 self._lexer.source,
@@ -1135,7 +1154,7 @@ def expect_optional_keyword(self, value: str) -> bool:
         """
         token = self._lexer.token
         if token.kind == TokenKind.NAME and token.value == value:
-            self._lexer.advance()
+            self.advance_lexer()
             return True
 
         return False
@@ -1223,6 +1242,20 @@ def delimited_many(
                 break
         return nodes
 
+    def advance_lexer(self) -> None:
+        """Advance the lexer."""
+        token = self._lexer.advance()
+        max_tokens = self._max_tokens
+        if max_tokens is not None and token.kind is not TokenKind.EOF:
+            self._token_counter += 1
+            if self._token_counter > max_tokens:
+                raise GraphQLSyntaxError(
+                    self._lexer.source,
+                    token.start,
+                    f"Document contains more that {max_tokens} tokens."
+                    " Parsing aborted.",
+                )
+
 
 def get_token_desc(token: Token) -> str:
     """Describe a token as a string for debugging."""

tests/language/test_parser.py

@@ -115,6 +115,22 @@ def parse_provides_useful_error_when_using_source():
             """
         )
 
+    def limits_maximum_number_of_tokens():
+        parse("{ foo }", max_tokens=3)
+        with raises(
+            GraphQLSyntaxError,
+            match="Syntax Error:"
+            r" Document contains more that 2 tokens\. Parsing aborted\.",
+        ):
+            parse("{ foo }", max_tokens=2)
+        parse('{ foo(bar: "baz") }', max_tokens=8)
+        with raises(
+            GraphQLSyntaxError,
+            match="Syntax Error:"
+            r" Document contains more that 7 tokens\. Parsing aborted\.",
+        ):
+            parse('{ foo(bar: "baz") }', max_tokens=7)
+
     def parses_variable_inline_values():
         parse("{ field(complex: { a: { b: [ $var ] } }) }")