graphql-java · andimarek · Mar 18, 2024 · Mar 14, 2024 · Mar 14, 2024 · Mar 18, 2024
diff --git a/src/main/java/graphql/execution/Execution.java b/src/main/java/graphql/execution/Execution.java
@@ -107,6 +107,7 @@ public CompletableFuture<ExecutionResult> execute(Document document, GraphQLSche
                 .executionInput(executionInput)
                 .build();
 
+        executionContext.getGraphQLContext().put(ResultNodesInfo.RESULT_NODES_INFO, executionContext.getResultNodesInfo());
 
         InstrumentationExecutionParameters parameters = new InstrumentationExecutionParameters(
                 executionInput, graphQLSchema, instrumentationState

diff --git a/src/main/java/graphql/execution/ExecutionContext.java b/src/main/java/graphql/execution/ExecutionContext.java
@@ -63,6 +63,8 @@ public class ExecutionContext {
     // this is modified after creation so it needs to be volatile to ensure visibility across Threads
     private volatile DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = DataLoaderDispatchStrategy.NO_OP;
 
+    private final ResultNodesInfo resultNodesInfo = new ResultNodesInfo();
+
     ExecutionContext(ExecutionContextBuilder builder) {
         this.graphQLSchema = builder.graphQLSchema;
         this.executionId = builder.executionId;
@@ -304,4 +306,8 @@ public ExecutionContext transform(Consumer<ExecutionContextBuilder> builderConsu
         builderConsumer.accept(builder);
         return builder.build();
     }
+
+    public ResultNodesInfo getResultNodesInfo() {
+        return resultNodesInfo;
+    }
 }
diff --git a/src/main/java/graphql/execution/ExecutionStrategy.java b/src/main/java/graphql/execution/ExecutionStrategy.java
@@ -67,6 +67,7 @@
 import static graphql.execution.FieldValueInfo.CompleteValueType.NULL;
 import static graphql.execution.FieldValueInfo.CompleteValueType.OBJECT;
 import static graphql.execution.FieldValueInfo.CompleteValueType.SCALAR;
+import static graphql.execution.ResultNodesInfo.MAX_RESULT_NODES;
 import static graphql.execution.instrumentation.SimpleInstrumentationContext.nonNullCtx;
 import static graphql.schema.DataFetchingEnvironmentImpl.newDataFetchingEnvironment;
 import static graphql.schema.GraphQLTypeUtil.isEnum;
@@ -381,6 +382,17 @@ protected CompletableFuture<FetchedValue> fetchField(ExecutionContext executionC
     }
 
     private CompletableFuture<FetchedValue> fetchField(GraphQLFieldDefinition fieldDef, ExecutionContext executionContext, ExecutionStrategyParameters parameters) {
+
+        int resultNodesCount = executionContext.getResultNodesInfo().incrementAndGetResultNodesCount();
+
+        Integer maxNodes;
+        if ((maxNodes = executionContext.getGraphQLContext().get(MAX_RESULT_NODES)) != null) {
+            if (resultNodesCount > maxNodes) {
+                executionContext.getResultNodesInfo().maxResultNodesExceeded();
+                return CompletableFuture.completedFuture(new FetchedValue(null, Collections.emptyList(), null));
+            }
+        }
+
         MergedField field = parameters.getField();
         GraphQLObjectType parentType = (GraphQLObjectType) parameters.getExecutionStepInfo().getUnwrappedNonNullType();
 
@@ -712,6 +724,15 @@ protected FieldValueInfo completeValueForList(ExecutionContext executionContext,
         List<FieldValueInfo> fieldValueInfos = new ArrayList<>(size.orElse(1));
         int index = 0;
         for (Object item : iterableValues) {
+            int resultNodesCount = executionContext.getResultNodesInfo().incrementAndGetResultNodesCount();
+            Integer maxNodes;
+            if ((maxNodes = executionContext.getGraphQLContext().get(MAX_RESULT_NODES)) != null) {
+                if (resultNodesCount > maxNodes) {
+                    executionContext.getResultNodesInfo().maxResultNodesExceeded();
+                    return new FieldValueInfo(NULL, completedFuture(null), fieldValueInfos);
+                }
+            }
+
             ResultPath indexedPath = parameters.getPath().segment(index);
 
             ExecutionStepInfo stepInfoForListElement = executionStepInfoFactory.newExecutionStepInfoForListElement(executionStepInfo, indexedPath);

diff --git a/src/main/java/graphql/execution/ResultNodesInfo.java b/src/main/java/graphql/execution/ResultNodesInfo.java
@@ -0,0 +1,55 @@
+package graphql.execution;
+
+import graphql.Internal;
+import graphql.PublicApi;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * This class is used to track the number of result nodes that have been created during execution.
+ * After each execution the GraphQLContext contains a ResultNodeInfo object under the key {@link ResultNodesInfo#RESULT_NODES_INFO}
+ * <p>
+ * The number of result can be limited (and should be for security reasons) by setting the maximum number of result nodes
+ * in the GraphQLContext under the key {@link ResultNodesInfo#MAX_RESULT_NODES} to an Integer
+ * </p>
+ */
+@PublicApi
+public class ResultNodesInfo {
+
+    public static final String MAX_RESULT_NODES = "__MAX_RESULT_NODES";
+    public static final String RESULT_NODES_INFO = "__RESULT_NODES_INFO";
+
+    private volatile boolean maxResultNodesExceeded = false;
+    private final AtomicInteger resultNodesCount = new AtomicInteger(0);
+
+    @Internal
+    public int incrementAndGetResultNodesCount() {
+        return resultNodesCount.incrementAndGet();
+    }
+
+    @Internal
+    public void maxResultNodesExceeded() {
+        this.maxResultNodesExceeded = true;
+    }
+
+    /**
+     * The number of result nodes created.
+     * Note: this can be higher than max result nodes because
+     * a each node that exceeds the number of max nodes is set to null,
+     * but still is a result node (with value null)
+     *
+     * @return number of result nodes created
+     */
+    public int getResultNodesCount() {
+        return resultNodesCount.get();
+    }
+
+    /**
+     * If the number of result nodes has exceeded the maximum allowed numbers.
+     *
+     * @return true if the number of result nodes has exceeded the maximum allowed numbers
+     */
+    public boolean isMaxResultNodesExceeded() {
+        return maxResultNodesExceeded;
+    }
+}
diff --git a/src/test/groovy/graphql/GraphQLTest.groovy b/src/test/groovy/graphql/GraphQLTest.groovy
@@ -13,6 +13,7 @@ import graphql.execution.ExecutionId
 import graphql.execution.ExecutionIdProvider
 import graphql.execution.ExecutionStrategyParameters
 import graphql.execution.MissingRootTypeException
+import graphql.execution.ResultNodesInfo
 import graphql.execution.SubscriptionExecutionStrategy
 import graphql.execution.ValueUnboxer
 import graphql.execution.instrumentation.Instrumentation
@@ -49,6 +50,7 @@ import static graphql.ExecutionInput.Builder
 import static graphql.ExecutionInput.newExecutionInput
 import static graphql.Scalars.GraphQLInt
 import static graphql.Scalars.GraphQLString
+import static graphql.execution.ResultNodesInfo.MAX_RESULT_NODES
 import static graphql.schema.GraphQLArgument.newArgument
 import static graphql.schema.GraphQLFieldDefinition.newFieldDefinition
 import static graphql.schema.GraphQLInputObjectField.newInputObjectField
@@ -1427,4 +1429,143 @@ many lines''']
         then:
         !er.errors.isEmpty()
     }
+
+    def "max result nodes not breached"() {
+        given:
+        def sdl = '''
+
+        type Query {
+          hello: String
+        }
+        '''
+        def df = { env -> "world" } as DataFetcher
+        def fetchers = ["Query": ["hello": df]]
+        def schema = TestUtil.schema(sdl, fetchers)
+        def graphQL = GraphQL.newGraphQL(schema).build()
+
+        def query = "{ hello h1: hello h2: hello h3: hello } "
+        def ei = newExecutionInput(query).build()
+        ei.getGraphQLContext().put(MAX_RESULT_NODES, 4);
+
+        when:
+        def er = graphQL.execute(ei)
+        def rni = ei.getGraphQLContext().get(ResultNodesInfo.RESULT_NODES_INFO) as ResultNodesInfo
+        then:
+        !rni.maxResultNodesExceeded
+        rni.resultNodesCount == 4
+        er.data == [hello: "world", h1: "world", h2: "world", h3: "world"]
+    }
+
+    def "max result nodes breached"() {
+        given:
+        def sdl = '''
+
+        type Query {
+          hello: String
+        }
+        '''
+        def df = { env -> "world" } as DataFetcher
+        def fetchers = ["Query": ["hello": df]]
+        def schema = TestUtil.schema(sdl, fetchers)
+        def graphQL = GraphQL.newGraphQL(schema).build()
+
+        def query = "{ hello h1: hello h2: hello h3: hello } "
+        def ei = newExecutionInput(query).build()
+        ei.getGraphQLContext().put(MAX_RESULT_NODES, 3);
+
+        when:
+        def er = graphQL.execute(ei)
+        def rni = ei.getGraphQLContext().get(ResultNodesInfo.RESULT_NODES_INFO) as ResultNodesInfo
+        then:
+        rni.maxResultNodesExceeded
+        rni.resultNodesCount == 4
+        er.data == [hello: "world", h1: "world", h2: "world", h3: null]
+    }
+
+    def "max result nodes breached with list"() {
+        given:
+        def sdl = '''
+
+        type Query {
+          hello: [String]
+        }
+        '''
+        def df = { env -> ["w1", "w2", "w3"] } as DataFetcher
+        def fetchers = ["Query": ["hello": df]]
+        def schema = TestUtil.schema(sdl, fetchers)
+        def graphQL = GraphQL.newGraphQL(schema).build()
+
+        def query = "{ hello}"
+        def ei = newExecutionInput(query).build()
+        ei.getGraphQLContext().put(MAX_RESULT_NODES, 3);
+
+        when:
+        def er = graphQL.execute(ei)
+        def rni = ei.getGraphQLContext().get(ResultNodesInfo.RESULT_NODES_INFO) as ResultNodesInfo
+        then:
+        rni.maxResultNodesExceeded
+        rni.resultNodesCount == 4
+        er.data == [hello: null]
+    }
+
+    def "max result nodes breached with list 2"() {
+        given:
+        def sdl = '''
+
+        type Query {
+          hello: [Foo]
+        }
+        type Foo {
+            name: String
+        }
+        '''
+        def df = { env -> [[name: "w1"], [name: "w2"], [name: "w3"]] } as DataFetcher
+        def fetchers = ["Query": ["hello": df]]
+        def schema = TestUtil.schema(sdl, fetchers)
+        def graphQL = GraphQL.newGraphQL(schema).build()
+
+        def query = "{ hello {name}}"
+        def ei = newExecutionInput(query).build()
+        // we have 7 result nodes overall
+        ei.getGraphQLContext().put(MAX_RESULT_NODES, 6);
+
+        when:
+        def er = graphQL.execute(ei)
+        def rni = ei.getGraphQLContext().get(ResultNodesInfo.RESULT_NODES_INFO) as ResultNodesInfo
+        then:
+        rni.resultNodesCount == 7
+        rni.maxResultNodesExceeded
+        er.data == [hello: [[name: "w1"], [name: "w2"], [name: null]]]
+    }
+
+    def "max result nodes not breached with list"() {
+        given:
+        def sdl = '''
+
+        type Query {
+          hello: [Foo]
+        }
+        type Foo {
+            name: String
+        }
+        '''
+        def df = { env -> [[name: "w1"], [name: "w2"], [name: "w3"]] } as DataFetcher
+        def fetchers = ["Query": ["hello": df]]
+        def schema = TestUtil.schema(sdl, fetchers)
+        def graphQL = GraphQL.newGraphQL(schema).build()
+
+        def query = "{ hello {name}}"
+        def ei = newExecutionInput(query).build()
+        // we have 7 result nodes overall
+        ei.getGraphQLContext().put(MAX_RESULT_NODES, 7);
+
+        when:
+        def er = graphQL.execute(ei)
+        def rni = ei.getGraphQLContext().get(ResultNodesInfo.RESULT_NODES_INFO) as ResultNodesInfo
+        then:
+        !rni.maxResultNodesExceeded
+        rni.resultNodesCount == 7
+        er.data == [hello: [[name: "w1"], [name: "w2"], [name: "w3"]]]
+    }
+
 }