Token limit fix CVE-2023-49559 (#291)

* Add directive limit to prevent overloading * Added token limit to parser for query and schema parsing, removed my previous directive limit * Added token limit to parser for query and schema parsing, removed my previous directive limit * Update parser/parser.go * Update parser/parser.go * Update parser/query.go * Fix lint Signed-off-by: Steve Coffman <[email protected]> --------- Signed-off-by: Steve Coffman <[email protected]> Co-authored-by: Yuval Moravchick <[email protected]> Co-authored-by: Steve Coffman <[email protected]> Co-authored-by: Steve Coffman <[email protected]>
vektah · Jun 11, 2024 · d457fc0 · d457fc0
1 parent 6db1bd3
commit d457fc0
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 4 deletions.
diff --git a/parser/parser.go b/parser/parser.go
@@ -1,6 +1,7 @@
 package parser
 
 import (
+	"fmt"
 	"strconv"
 
 	"github.com/vektah/gqlparser/v2/ast"
@@ -20,6 +21,13 @@ type parser struct {
 
 	comment          *ast.CommentGroup
 	commentConsuming bool
+
+	tokenCount    int
+	maxTokenLimit int
+}
+
+func (p *parser) SetMaxTokenLimit(maxToken int) {
+	p.maxTokenLimit = maxToken
 }
 
 func (p *parser) consumeComment() (*ast.Comment, bool) {
@@ -95,6 +103,12 @@ func (p *parser) next() lexer.Token {
 	if p.err != nil {
 		return p.prev
 	}
+	// Increment the token count before reading the next token
+	p.tokenCount++
+	if p.maxTokenLimit != 0 && p.tokenCount > p.maxTokenLimit {
+		p.err = fmt.Errorf("exceeded token limit of %d", p.maxTokenLimit)
+		return p.prev
+	}
 	if p.peeked {
 		p.peeked = false
 		p.comment = nil

diff --git a/parser/parser_test.go b/parser/parser_test.go
@@ -166,5 +166,8 @@ func TestParserUtils(t *testing.T) {
 }
 
 func newParser(input string) parser {
-	return parser{lexer: lexer.New(&ast.Source{Input: input, Name: "input.graphql"})}
+	return parser{
+		lexer:         lexer.New(&ast.Source{Input: input, Name: "input.graphql"}),
+		maxTokenLimit: 15000, // 15000 is the default value
+	}
 }
diff --git a/parser/query.go b/parser/query.go
@@ -2,14 +2,14 @@ package parser
 
 import (
 	"github.com/vektah/gqlparser/v2/lexer"
-
 	//nolint:revive
 	. "github.com/vektah/gqlparser/v2/ast"
 )
 
 func ParseQuery(source *Source) (*QueryDocument, error) {
 	p := parser{
-		lexer: lexer.New(source),
+		lexer:         lexer.New(source),
+		maxTokenLimit: 0, // 0 is the default value
 	}
 	return p.parseQueryDocument(), p.err
 }

diff --git a/parser/schema.go b/parser/schema.go
@@ -20,7 +20,8 @@ func ParseSchemas(inputs ...*Source) (*SchemaDocument, error) {
 
 func ParseSchema(source *Source) (*SchemaDocument, error) {
 	p := parser{
-		lexer: lexer.New(source),
+		lexer:         lexer.New(source),
+		maxTokenLimit: 15000, // default value
 	}
 	sd, err := p.parseSchemaDocument(), p.err
 	if err != nil {