fix: add missing aws_iam directive on custom types (#3270)

* fix: add missing aws_iam directive on custom types

* chore: update api extract

* chore: remove todos

* test: add iam auth test specifically for EventInvocationResponse type

---------

Co-authored-by: Jon Wire <wirej@amazon.com>

Author: svidgen

packages/amplify-graphql-auth-transformer/API.md

@@ -117,6 +117,8 @@ export class AuthTransformer extends TransformerAuthBase implements TransformerA
     // (undocumented)
     addFieldsToObject: (ctx: TransformerTransformSchemaStepContextProvider, modelName: string, ownerFields: Array<string>) => void;
     // (undocumented)
+    addIamAuthToCustomTypes: (ctx: TransformerTransformSchemaStepContextProvider) => void;
+    // (undocumented)
     after: (context: TransformerContextProvider) => void;
     // (undocumented)
     before: (context: TransformerBeforeStepContextProvider) => void;

packages/amplify-graphql-auth-transformer/src/__tests__/iam-custom-operations.test.ts

@@ -1,3 +1,4 @@
+import { FunctionTransformer } from '@aws-amplify/graphql-function-transformer';
 import { ModelTransformer } from '@aws-amplify/graphql-model-transformer';
 import { mockSqlDataSourceStrategy, testTransform } from '@aws-amplify/graphql-transformer-test-utils';
 import { PrimaryKeyTransformer } from '@aws-amplify/graphql-index-transformer';
@@ -40,6 +41,7 @@ const makeTransformers = (): TransformerPluginProvider[] => [
   new AuthTransformer(),
   new PrimaryKeyTransformer(),
   new SqlTransformer(),
+  new FunctionTransformer(),
 ];
 
 const makeSqlDirectiveDataSourceStrategies = (schema: string, strategy: ModelDataSourceStrategy): SqlDirectiveDataSourceStrategy[] =>
@@ -367,8 +369,7 @@ describe('Custom operations have @aws_iam directives when enableIamAuthorization
       expect(out.schema).not.toMatch(/onUpdateFooCustom: String.*@aws_iam/);
     });
 
-    // TODO: Enable this test once we fix https://github.com/aws-amplify/amplify-category-api/issues/2929
-    test.skip('Adds @aws_iam to non-model custom types when there is no model', () => {
+    test('Adds @aws_iam to non-model custom types when there is no model', () => {
       const strategy = makeStrategy(strategyType);
       const schema = /* GraphQL */ `
         type Foo {
@@ -403,8 +404,7 @@ describe('Custom operations have @aws_iam directives when enableIamAuthorization
       expect(out.schema).toMatch(/type Foo.*@aws_iam/);
     });
 
-    // TODO: Enable this test once we fix https://github.com/aws-amplify/amplify-category-api/issues/2929
-    test.skip('Adds @aws_iam to non-model custom types when there is a model', () => {
+    test('Adds @aws_iam to non-model custom types when there is a model', () => {
       const strategy = makeStrategy(strategyType);
       const schema = /* GraphQL */ `
         type Todo @model {
@@ -478,6 +478,42 @@ describe('Custom operations have @aws_iam directives when enableIamAuthorization
       expect(out.schema).toMatch(/description: String.*@aws_iam/);
     });
 
+    test('Adds @aws_iam to async function EventInvocationResponse type', () => {
+      const strategy = makeStrategy(strategyType);
+      const schema = /* GraphQL */ `
+        type Foo {
+          description: String
+        }
+        type EventInvocationResponse @aws_api_key {
+          success: Boolean!
+        }
+        type Query {
+          getFooCustom: Foo
+        }
+        type Mutation {
+          updateFooCustom: Foo
+          doSomethingAsync(body: String!): EventInvocationResponse
+            @function(name: "FnDoSomethingAsync", invocationType: Event)
+            @auth(rules: [{ allow: public, provider: apiKey }])
+        }
+        type Subscription {
+          onUpdateFooCustom: Foo @aws_subscribe(mutations: ["updateFooCustom"])
+        }
+      `;
+
+      const out = testTransform({
+        schema,
+        dataSourceStrategies: constructDataSourceStrategies(schema, strategy),
+        authConfig: makeAuthConfig(),
+        synthParameters: makeSynthParameters(),
+        transformers: makeTransformers(),
+        sqlDirectiveDataSourceStrategies: makeSqlDirectiveDataSourceStrategies(schema, strategy),
+      });
+
+      // Also expect the custom type referenced by the custom operation to be authorized
+      expect(out.schema).toMatch(/type EventInvocationResponse.*@aws_iam/);
+    });
+
     test('Does not add duplicate @aws_iam directive to custom type if already present', () => {
       const strategy = makeStrategy(strategyType);
       const schema = /* GraphQL */ `

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

@@ -10,6 +10,7 @@ import {
   isBuiltInGraphqlNode,
   isDynamoDbModel,
   isModelType,
+  isObjectTypeDefinitionNode,
   isSqlModel,
   MappingTemplate,
   TransformerAuthBase,
@@ -377,8 +378,32 @@ export class AuthTransformer extends TransformerAuthBase implements TransformerA
     });
   };
 
+  /**
+   * Adds the `aws_iam` directive to all custom/non-model types when IAM access is enabled.
+   */
+  addIamAuthToCustomTypes = (ctx: TransformerTransformSchemaStepContextProvider): void => {
+    if (!ctx.transformParameters.sandboxModeEnabled && !ctx.synthParameters.enableIamAccess) {
+      return;
+    }
+
+    const needsAwsIamDirective = (type: TypeDefinitionNode): boolean => {
+      return !type.directives?.some((dir) => dir.name.value === 'aws_iam');
+    };
+
+    const isNonModelType = (type: TypeDefinitionNode): boolean => {
+      return !type.directives?.some((dir) => dir.name.value === 'model');
+    };
+
+    ctx.inputDocument.definitions
+      .filter(isObjectTypeDefinitionNode)
+      .filter(isNonModelType)
+      .filter(needsAwsIamDirective)
+      .forEach((def) => extendTypeWithDirectives(ctx, def.name.value, [makeDirective('aws_iam', [])]));
+  };
+
   transformSchema = (context: TransformerTransformSchemaStepContextProvider): void => {
     this.addCustomOperationFieldsToAuthNonModelConfig(context);
+    this.addIamAuthToCustomTypes(context);
 
     const searchableAggregateServiceDirectives = new Set<AuthProvider>();