From 29e42bdc00aed8a7ff7813dae99303c689a546c9 Mon Sep 17 00:00:00 2001
From: Brad Baker <bbaker@atlassian.com>
Date: Wed, 7 Sep 2022 13:39:33 +1000
Subject: [PATCH] Stable port of Fix field visibility bug with enum with enum
 args (#2926)

Fixes a bug that happened when the field visibility transformation
was applied to a schema that has some type(s) reachable only via
directive definitions.
---
 .../FieldVisibilitySchemaTransformation.java  | 18 +++++-
 ...dVisibilitySchemaTransformationTest.groovy | 63 +++++++++++++++++++
 2 files changed, 78 insertions(+), 3 deletions(-)

diff --git a/src/main/java/graphql/schema/transform/FieldVisibilitySchemaTransformation.java b/src/main/java/graphql/schema/transform/FieldVisibilitySchemaTransformation.java
index a290258838..303b81b9c8 100644
--- a/src/main/java/graphql/schema/transform/FieldVisibilitySchemaTransformation.java
+++ b/src/main/java/graphql/schema/transform/FieldVisibilitySchemaTransformation.java
@@ -1,10 +1,12 @@
 package graphql.schema.transform;
 
+import com.google.common.collect.ImmutableList;
 import graphql.PublicApi;
 import graphql.schema.GraphQLEnumType;
 import graphql.schema.GraphQLFieldDefinition;
 import graphql.schema.GraphQLImplementingType;
 import graphql.schema.GraphQLInputObjectField;
+import graphql.schema.GraphQLInputObjectType;
 import graphql.schema.GraphQLInterfaceType;
 import graphql.schema.GraphQLNamedSchemaElement;
 import graphql.schema.GraphQLNamedType;
@@ -60,7 +62,7 @@ public final GraphQLSchema apply(GraphQLSchema schema) {
         Set<GraphQLType> markedForRemovalTypes = new HashSet<>();
 
         // query, mutation, and subscription types should not be removed
-        final Set<String> protectedTypeNames = getRootTypes(schema).stream()
+        final Set<String> protectedTypeNames = getOperationTypes(schema).stream()
                 .map(GraphQLObjectType::getName)
                 .collect(Collectors.toSet());
 
@@ -216,6 +218,7 @@ public TraversalControl visitGraphQLType(GraphQLSchemaElement node,
                     !observedAfterTransform.contains(node) &&
                     (node instanceof GraphQLObjectType ||
                             node instanceof GraphQLEnumType ||
+                            node instanceof GraphQLInputObjectType ||
                             node instanceof GraphQLInterfaceType ||
                             node instanceof GraphQLUnionType)) {
 
@@ -250,12 +253,21 @@ public TraversalControl visitGraphQLType(GraphQLSchemaElement node,
         }
     }
 
-    private List<GraphQLObjectType> getRootTypes(GraphQLSchema schema) {
+    private List<GraphQLSchemaElement> getRootTypes(GraphQLSchema schema) {
+        return ImmutableList.<GraphQLSchemaElement>builder()
+                .addAll(getOperationTypes(schema))
+                // Include directive definitions as roots, since they won't be removed in the filtering process.
+                // Some types (enums, input types, etc.) might be reachable only by directive definitions (and
+                // not by other types or fields).
+                .addAll(schema.getDirectives())
+                .build();
+    }
+
+    private List<GraphQLObjectType> getOperationTypes(GraphQLSchema schema) {
         return Stream.of(
                 schema.getQueryType(),
                 schema.getSubscriptionType(),
                 schema.getMutationType()
         ).filter(Objects::nonNull).collect(Collectors.toList());
     }
-
 }
diff --git a/src/test/groovy/graphql/schema/transform/FieldVisibilitySchemaTransformationTest.groovy b/src/test/groovy/graphql/schema/transform/FieldVisibilitySchemaTransformationTest.groovy
index c6f25fd367..a6f82018a2 100644
--- a/src/test/groovy/graphql/schema/transform/FieldVisibilitySchemaTransformationTest.groovy
+++ b/src/test/groovy/graphql/schema/transform/FieldVisibilitySchemaTransformationTest.groovy
@@ -1178,4 +1178,67 @@ class FieldVisibilitySchemaTransformationTest extends Specification {
         (restrictedSchema.getType("Account") as GraphQLObjectType).getFieldDefinition("billingStatus") == null
         restrictedSchema.getType("BillingStatus") != null
     }
+
+    def "can remove a field with a directive containing enum argument"() {
+        given:
+        GraphQLSchema schema = TestUtil.schema("""
+
+        directive @private(privateType: SecretType) on FIELD_DEFINITION
+        enum SecretType {
+            SUPER_SECRET
+            NOT_SO_SECRET
+        }
+
+        type Query {
+            account: Account
+        }
+        
+        type Account {
+            name: String
+            billingStatus: BillingStatus @private(privateType: NOT_SO_SECRET)
+        }
+        
+        type BillingStatus {
+            accountNumber: String
+        }
+        """)
+
+        when:
+        GraphQLSchema restrictedSchema = visibilitySchemaTransformation.apply(schema)
+
+        then:
+        (restrictedSchema.getType("Account") as GraphQLObjectType).getFieldDefinition("billingStatus") == null
+        restrictedSchema.getType("BillingStatus") == null
+    }
+
+    def "can remove a field with a directive containing type argument"() {
+        given:
+        GraphQLSchema schema = TestUtil.schema("""
+
+        directive @private(privateType: SecretType) on FIELD_DEFINITION
+        input SecretType {
+            description: String
+        }
+
+        type Query {
+            account: Account
+        }
+        
+        type Account {
+            name: String
+            billingStatus: BillingStatus @private(privateType: { description: "secret" })
+        }
+        
+        type BillingStatus {
+            accountNumber: String
+        }
+        """)
+
+        when:
+        GraphQLSchema restrictedSchema = visibilitySchemaTransformation.apply(schema)
+
+        then:
+        (restrictedSchema.getType("Account") as GraphQLObjectType).getFieldDefinition("billingStatus") == null
+        restrictedSchema.getType("BillingStatus") == null
+    }
 }