Add repetition detection to the GraphQL package

Signed-off-by: peterdeme <demepeter93@gmail.com>

Author: peterdeme

ddos_vulnerability_test.go

@@ -175,16 +175,15 @@ func TestDDoSVulnerability_WithFix(t *testing.T) {
 		t.Fatal("Expected validation errors, but got none")
 	}
 
-	// Check that the error message mentions the max selection set size
 	found := false
 	for _, err := range errors {
-		if strings.Contains(err.Message, "exceeds the maximum allowed size") {
+		if strings.Contains(err.Message, "invalid query") {
 			found = true
 			break
 		}
 	}
 	if !found {
-		t.Errorf("Expected error about exceeding max selection set size, but got: %v", errors)
+		t.Errorf("Expected error about exceeding max selection set size or repeated tokens, but got: %v", errors)
 	}
 
 	// Validation should be fast (< 100ms)
@@ -215,3 +214,99 @@ func TestDDoSVulnerability_FixWithReasonableQuery(t *testing.T) {
 		}
 	}
 }
+
+// TestDDoSVulnerability_RepeatedTokenDetection tests early detection of repeated field patterns
+func TestDDoSVulnerability_RepeatedTokenDetection(t *testing.T) {
+	// Create a query with 150 repeated "a" fields - should be rejected during parsing
+	numFields := 150
+	fields := make([]string, numFields)
+	for i := 0; i < numFields; i++ {
+		fields[i] = "a"
+	}
+
+	maliciousQuery := "query { " + strings.Join(fields, " ") + " }"
+
+	schema := graphql.MustParseSchema(simpleSchema, &simpleResolver{})
+
+	start := time.Now()
+	result := schema.Exec(context.Background(), maliciousQuery, "", nil)
+	duration := time.Since(start)
+
+	t.Logf("Query with %d repeated tokens took %v", numFields, duration)
+	t.Logf("Errors: %v", result.Errors)
+
+	// Should have errors
+	if len(result.Errors) == 0 {
+		t.Fatal("Expected errors for repeated token attack, but got none")
+	}
+
+	// Should be rejected quickly (parsing phase)
+	if duration > 100*time.Millisecond {
+		t.Errorf("Query took too long (%v) - should be rejected during parsing", duration)
+	}
+
+	// Check error message mentions repetition or DDoS
+	foundRelevantError := false
+	for _, err := range result.Errors {
+		if strings.Contains(err.Message, "syntax error: invalid query") {
+			foundRelevantError = true
+			t.Logf("Found expected error: %s", err.Message)
+			break
+		}
+	}
+
+	if !foundRelevantError {
+		t.Errorf("Expected error about repeated tokens, got: %v", result.Errors)
+	}
+}
+
+// TestDDoSVulnerability_RepeatedTokenAtThreshold tests behavior at the threshold
+func TestDDoSVulnerability_RepeatedTokenAtThreshold(t *testing.T) {
+	// Create a query with exactly 20 repeated "a" fields - should be allowed
+	numFields := 20
+	fields := make([]string, numFields)
+	for i := 0; i < numFields; i++ {
+		fields[i] = "a"
+	}
+
+	queryAtThreshold := "query { " + strings.Join(fields, " ") + " }"
+
+	schema := graphql.MustParseSchema(simpleSchema, &simpleResolver{})
+
+	start := time.Now()
+	result := schema.Exec(context.Background(), queryAtThreshold, "", nil)
+	duration := time.Since(start)
+
+	t.Logf("Query with %d repeated tokens (at threshold) took %v", numFields, duration)
+	t.Logf("Errors: %v", result.Errors)
+
+	// Should be fast
+	if duration > 100*time.Millisecond {
+		t.Errorf("Query took too long (%v)", duration)
+	}
+
+	if len(result.Errors) > 0 {
+		t.Errorf("Expected no errors for query at threshold, but got: %v", result.Errors)
+	}
+}
+
+// TestDDoSVulnerability_NonRepeatedFields tests that different fields don't trigger detection
+func TestDDoSVulnerability_NonRepeatedFields(t *testing.T) {
+	// Create a query with many different fields - should not be blocked by repetition detection
+	// Note: This will still be caught by MaxSelectionSetSize if configured
+	queryWithDifferentFields := "query { a b c d e f g h i j k l m n o p q r s t u v w x y z }"
+
+	schema := graphql.MustParseSchema(simpleSchema, &simpleResolver{})
+
+	start := time.Now()
+	result := schema.Exec(context.Background(), queryWithDifferentFields, "", nil)
+	duration := time.Since(start)
+
+	t.Logf("Query with different fields took %v", duration)
+	t.Logf("Errors: %v", result.Errors)
+
+	err := result.Errors[0]
+	if !strings.Contains(err.Message, "Cannot query field") {
+		t.Fatalf("Expected error about unknown fields, got: %s", err.Message)
+	}
+}

internal/common/lexer.go

@@ -57,6 +57,11 @@ func (l *Lexer) Peek() rune {
 	return l.next
 }
 
+// TokenText returns the text of the current token without consuming it.
+func (l *Lexer) TokenText() string {
+	return l.sc.TokenText()
+}
+
 // ConsumeWhitespace consumes whitespace and tokens equivalent to whitespace (e.g. commas and comments).
 //
 // Consumed comment characters will build the description for the next type or field encountered.

internal/query/query.go

@@ -96,9 +96,33 @@ func parseFragment(l *common.Lexer) *ast.FragmentDefinition {
 }
 
 func parseSelectionSet(l *common.Lexer) []ast.Selection {
+	const maxRepeatedTokens = 20 // Prevent DDoS attacks with repeated fields like "a a a a..."
+
 	var sels []ast.Selection
+	var prevFieldName string
+	repeatCount := 0
+
 	l.ConsumeToken('{')
 	for l.Peek() != '}' {
+		// Detect repeated identical field names to prevent DDoS attacks
+		// Check the token text before consuming to detect patterns like "a a a a..."
+		if l.Peek() == scanner.Ident {
+			currentFieldName := l.TokenText()
+			if currentFieldName == prevFieldName {
+				repeatCount++
+				if repeatCount >= maxRepeatedTokens {
+					l.SyntaxError("invalid query")
+				}
+			} else {
+				repeatCount = 0
+				prevFieldName = currentFieldName
+			}
+		} else {
+			// Reset for non-field tokens (fragments, etc.)
+			repeatCount = 0
+			prevFieldName = ""
+		}
+
 		sels = append(sels, parseSelection(l))
 	}
 	l.ConsumeToken('}')