fix: transform auth on fields of supported extended types

packages/amplify-graphql-auth-transformer/src/__tests__/extended-types.test.ts

@@ -0,0 +1,191 @@
+import { ModelTransformer } from '@aws-amplify/graphql-model-transformer';
+import { testTransform } from '@aws-amplify/graphql-transformer-test-utils';
+import { AppSyncAuthConfiguration } from '@aws-amplify/graphql-transformer-interfaces';
+import { DocumentNode, FieldDefinitionNode, Kind, ObjectTypeDefinitionNode, parse } from 'graphql';
+import { AuthTransformer } from '../graphql-auth-transformer';
+
+describe('@auth directive on extended types', () => {
+  const authConfig: AppSyncAuthConfiguration = {
+    defaultAuthentication: {
+      authenticationType: 'AMAZON_COGNITO_USER_POOLS',
+    },
+    additionalAuthenticationProviders: [{ authenticationType: 'AWS_IAM' }],
+  };
+
+  const getObjectType = (doc: DocumentNode, type: string): ObjectTypeDefinitionNode | undefined => {
+    return doc.definitions.find((def) => def.kind === Kind.OBJECT_TYPE_DEFINITION && def.name.value === type) as
+      | ObjectTypeDefinitionNode
+      | undefined;
+  };
+
+  const expectNoDirectives = (fieldOrType: ObjectTypeDefinitionNode | FieldDefinitionNode | undefined): void => {
+    expect(fieldOrType?.directives?.length).toEqual(0);
+  };
+
+  const expectOneDirective = (fieldOrType: ObjectTypeDefinitionNode | FieldDefinitionNode | undefined, directiveName: string): void => {
+    expect(fieldOrType?.directives?.length).toEqual(1);
+    expect(fieldOrType?.directives?.find((d) => d.name.value === directiveName)).toBeDefined();
+  };
+
+  const expectDirectiveWithName = (
+    fieldOrType: ObjectTypeDefinitionNode | FieldDefinitionNode | undefined,
+    directiveName: string,
+  ): void => {
+    expect(fieldOrType?.directives?.find((d) => d.name.value === directiveName)).toBeDefined();
+  };
+
+  const getField = (type: ObjectTypeDefinitionNode | undefined, name: string): FieldDefinitionNode | undefined =>
+    type?.fields?.find((f) => f.name.value === name);
+
+  test('supports @auth directive on a type that has been extended', () => {
+    const schema = /* GraphQL */ `
+      type Todo @model @auth(rules: [{ allow: public, provider: iam }]) {
+        id: ID!
+        description: String
+      }
+
+      extend type Todo {
+        extendedField: String!
+      }
+    `;
+
+    const testTransformParams = {
+      schema: schema,
+      authConfig,
+      transformers: [new ModelTransformer(), new AuthTransformer()],
+    };
+
+    const out = testTransform(testTransformParams);
+    expect(out).toBeDefined();
+
+    const transformedSchema = out.schema;
+    const doc = parse(transformedSchema);
+
+    const todoType = getObjectType(doc, 'Todo');
+    expect(todoType).toBeDefined();
+    expectDirectiveWithName(todoType, 'aws_iam');
+  });
+
+  test('does not support @auth directive on the model type extension itself', () => {
+    const schema = /* GraphQL */ `
+      type Todo @model {
+        id: ID!
+        description: String
+      }
+
+      extend type Todo @auth(rules: [{ allow: public, provider: iam }]) {
+        extendedField: String!
+      }
+    `;
+
+    const testTransformParams = {
+      schema: schema,
+      authConfig,
+      transformers: [new ModelTransformer(), new AuthTransformer()],
+    };
+
+    expect(() => testTransform(testTransformParams)).toThrow(
+      "Directives are not supported on object or interface extensions. See the '@auth' directive on 'Todo'",
+    );
+  });
+
+  test('does not support @auth directive on fields of model type extensions', () => {
+    const schema = /* GraphQL */ `
+      type Todo @model {
+        id: ID!
+        description: String
+      }
+
+      extend type Todo @auth(rules: [{ allow: public, provider: iam }]) {
+        extendedField: String!
+      }
+    `;
+
+    const testTransformParams = {
+      schema: schema,
+      authConfig,
+      transformers: [new ModelTransformer(), new AuthTransformer()],
+    };
+
+    expect(() => testTransform(testTransformParams)).toThrow(
+      "Directives are not supported on object or interface extensions. See the '@auth' directive on 'Todo'",
+    );
+  });
+
+  test.each(['Query', 'Mutation', 'Subscription'])('supports @auth directive on fields of %s type extensions', (builtInType: string) => {
+    const schema = /* GraphQL */ `
+      type ${builtInType} {
+        customOperation1: String! @auth(rules: [{ allow: public, provider: iam }])
+      }
+
+      extend type ${builtInType} {
+        customOperation2: String! @auth(rules: [{ allow: public, provider: iam }])
+      }
+    `;
+
+    const testTransformParams = {
+      schema: schema,
+      authConfig,
+      transformers: [new AuthTransformer()],
+    };
+
+    const out = testTransform(testTransformParams);
+    expect(out).toBeDefined();
+
+    const transformedSchema = out.schema;
+    const doc = parse(transformedSchema);
+
+    const operationType = getObjectType(doc, builtInType);
+    expect(operationType).toBeDefined();
+    expectNoDirectives(operationType!);
+    expectOneDirective(getField(operationType, 'customOperation1'), 'aws_iam');
+    expectOneDirective(getField(operationType, 'customOperation2'), 'aws_iam');
+  });
+
+  test.each(['Query', 'Mutation', 'Subscription'])(
+    'does not support @auth directive on %s object extension itself',
+    (builtInType: string) => {
+      const schema = /* GraphQL */ `
+        type ${builtInType} {
+          customOperation1: String! @auth(rules: [{ allow: public, provider: iam }])
+        }
+
+        extend type ${builtInType} @auth(rules: [{ allow: public, provider: iam }]) {
+          customOperation2: String! 
+        }
+      `;
+
+      const testTransformParams = {
+        schema: schema,
+        authConfig,
+        transformers: [new AuthTransformer()],
+      };
+
+      expect(() => testTransform(testTransformParams)).toThrow(
+        `Directives are not supported on object or interface extensions. See the '@auth' directive on '${builtInType}'`,
+      );
+    },
+  );
+
+  test('does not support @auth directive on fields of non-model type extensions', () => {
+    const schema = /* GraphQL */ `
+      type Foo {
+        customField1: String! @auth(rules: [{ allow: public, provider: iam }])
+      }
+
+      extend type Foo {
+        customField2: String! @auth(rules: [{ allow: public, provider: iam }])
+      }
+    `;
+
+    const testTransformParams = {
+      schema: schema,
+      authConfig,
+      transformers: [new AuthTransformer()],
+    };
+
+    expect(() => testTransform(testTransformParams)).toThrow(
+      "The '@auth' directive cannot be used on fields of type extensions other than 'Query', 'Mutation', and 'Subscription'. See Foo.customField2",
+    );
+  });
+});

packages/amplify-graphql-auth-transformer/src/graphql-auth-transformer.ts

@@ -13,20 +13,20 @@ import {
   getFilterInputName,
   getConditionInputName,
   getSubscriptionFilterInputName,
-  getConnectionName,
-  InputFieldWrapper,
   isDynamoDbModel,
+  isBuiltInGraphqlNode,
 } from '@aws-amplify/graphql-transformer-core';
 import {
   DataSourceProvider,
   MutationFieldType,
   QueryFieldType,
-  TransformerTransformSchemaStepContextProvider,
+  TransformerAuthProvider,
+  TransformerBeforeStepContextProvider,
   TransformerContextProvider,
   TransformerResolverProvider,
   TransformerSchemaVisitStepContextProvider,
-  TransformerAuthProvider,
-  TransformerBeforeStepContextProvider,
+  TransformerTransformSchemaStepContextProvider,
+  TransformerValidationStepContextProvider,
 } from '@aws-amplify/graphql-transformer-interfaces';
 import {
   DirectiveNode,
@@ -37,6 +37,7 @@ import {
   TypeDefinitionNode,
   ListValueNode,
   StringValueNode,
+  ObjectTypeExtensionNode,
 } from 'graphql';
 import { merge } from 'lodash';
 import {
@@ -268,9 +269,27 @@ export class AuthTransformer extends TransformerAuthBase implements TransformerA
   ): void => {
     if (parent.kind === Kind.INTERFACE_TYPE_DEFINITION) {
       throw new InvalidDirectiveError(
-        `The @auth directive cannot be placed on an interface's field. See ${parent.name.value}${field.name.value}`,
+        `The @auth directive cannot be placed on an interface's field. See ${parent.name.value}.${field.name.value}`,
       );
     }
+    this.transformField(parent, field, directive, context);
+  };
+
+  fieldOfExtendedType = (
+    parent: ObjectTypeExtensionNode,
+    field: FieldDefinitionNode,
+    directive: DirectiveNode,
+    context: TransformerSchemaVisitStepContextProvider,
+  ): void => {
+    this.transformField(parent, field, directive, context);
+  };
+
+  private transformField = (
+    parent: ObjectTypeDefinitionNode | InterfaceTypeDefinitionNode | ObjectTypeExtensionNode,
+    field: FieldDefinitionNode,
+    directive: DirectiveNode,
+    context: TransformerSchemaVisitStepContextProvider,
+  ): void => {
     const isParentTypeBuiltinType =
       parent.name.value === context.output.getQueryTypeName() ||
       parent.name.value === context.output.getMutationTypeName() ||
@@ -720,11 +739,12 @@ export class AuthTransformer extends TransformerAuthBase implements TransformerA
     typeName: string,
     field: FieldDefinitionNode,
     fieldRoles: Array<string>,
-    needsFieldResolver: boolean = false,
+    needsFieldResolver = false,
   ): void => {
     let fieldAuthExpression: string;
     let relatedAuthExpression: string;
-    // Relational field redaction is default to `needsFieldResolver`, which stays consistent with current behavior of always redacting relational field when field resolver is needed
+    // Relational field redaction is default to `needsFieldResolver`, which stays consistent with current behavior of always redacting
+    // relational field when field resolver is needed
     let redactRelationalField: boolean = needsFieldResolver;
     const fieldIsRequired = field.type.kind === Kind.NON_NULL_TYPE;
     if (fieldIsRequired) {
@@ -779,7 +799,7 @@ export class AuthTransformer extends TransformerAuthBase implements TransformerA
          * Once there is one role detected to have access on both side, the auth role definitions will be compared to determine whether
          * to redac the field or not
          */
-        for (let fieldRole of fieldReadRoleDefinitions) {
+        for (const fieldRole of fieldReadRoleDefinitions) {
           // When two role definitions have an overlap
           if (isFieldRoleHavingAccessToBothSide(fieldRole, filteredRelatedModelReadRoleDefinitions)) {
             // Check if two role definitions are identical without dynamic auth role or custom auth role