From 2f13fec0e2660f97da2db7028b4949e995dbcaaf Mon Sep 17 00:00:00 2001
From: Jonathan Goldwasser <jonathan.goldwasser@gmail.com>
Date: Fri, 16 Sep 2022 10:16:50 +0200
Subject: [PATCH] feat(stepfunctions-tasks): additional IAM statements for AWS
 SDK service integration

Add a `additionalIamStatements` prop to pass additional IAM statements.
To be used when the call requires more than a single statement to be
executed.

Closes #22006
---
 .../aws-stepfunctions-tasks/README.md         | 18 +++++++++
 .../lib/aws-sdk/call-aws-service.ts           | 13 +++++++
 .../test/aws-sdk/call-aws-service.test.ts     | 39 +++++++++++++++++++
 3 files changed, 70 insertions(+)

diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/README.md b/packages/@aws-cdk/aws-stepfunctions-tasks/README.md
index ad1faf00cc029..41812352c2c14 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/README.md
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/README.md
@@ -223,6 +223,24 @@ const listBuckets = new tasks.CallAwsService(this, 'ListBuckets', {
 });
 ```
 
+Use the `additionalIamStatements` prop to pass additional IAM statements that will be added to the
+state machine role's policy. Use it in the case where the call requires more than a single statement
+to be executed:
+
+```ts
+const detectLabels = new tasks.CallAwsService(stack, 'DetectLabels', {
+  service: 'rekognition',
+  action: 'detectLabels',
+  iamResources: ['*'],
+  additionalIamStatements: [
+    new iam.PolicyStatement({
+      actions: ['s3:getObject'],
+      resources: ['arn:aws:s3:::my-bucket/*'],
+    }),
+  ],
+});
+```
+
 ## Athena
 
 Step Functions supports [Athena](https://docs.aws.amazon.com/step-functions/latest/dg/connect-athena.html) through the service integration pattern.
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/aws-sdk/call-aws-service.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/aws-sdk/call-aws-service.ts
index 34f1bd78fe3c2..7a2c33ed321cd 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/aws-sdk/call-aws-service.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/aws-sdk/call-aws-service.ts
@@ -52,6 +52,18 @@ export interface CallAwsServiceProps extends sfn.TaskStateBaseProps {
    * @default - service:action
    */
   readonly iamAction?: string;
+
+  /**
+   * Additional IAM statements that will be added to the state machine
+   * role's policy.
+   *
+   * Use in the case where the call requires more than a single statement to
+   * be executed, e.g. `rekognition:detectLabels` requires also S3 permissions
+   * to read the object on which it must act.
+   *
+   * @default - no additional statements are added
+   */
+  readonly additionalIamStatements?: iam.PolicyStatement[];
 }
 
 /**
@@ -75,6 +87,7 @@ export class CallAwsService extends sfn.TaskStateBase {
         // https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_elements_action.html
         actions: [props.iamAction ?? `${props.service}:${props.action}`],
       }),
+      ...props.additionalIamStatements ?? [],
     ];
   }
 
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/aws-sdk/call-aws-service.test.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/test/aws-sdk/call-aws-service.test.ts
index da26dc7d8450e..9f7321ac8acad 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/test/aws-sdk/call-aws-service.test.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/test/aws-sdk/call-aws-service.test.ts
@@ -1,4 +1,5 @@
 import { Template } from '@aws-cdk/assertions';
+import * as iam from '@aws-cdk/aws-iam';
 import * as sfn from '@aws-cdk/aws-stepfunctions';
 import * as cdk from '@aws-cdk/core';
 import * as tasks from '../../lib';
@@ -159,3 +160,41 @@ test('throws with invalid integration pattern', () => {
     iamResources: ['*'],
   })).toThrow(/The RUN_JOB integration pattern is not supported for CallAwsService/);
 });
+
+test('can pass additional IAM statements', () => {
+  // WHEN
+  const task = new tasks.CallAwsService(stack, 'DetectLabels', {
+    service: 'rekognition',
+    action: 'detectLabels',
+    iamResources: ['*'],
+    additionalIamStatements: [
+      new iam.PolicyStatement({
+        actions: ['s3:getObject'],
+        resources: ['arn:aws:s3:::my-bucket/*'],
+      }),
+    ],
+  });
+
+  new sfn.StateMachine(stack, 'StateMachine', {
+    definition: task,
+  });
+
+  // THEN
+  Template.fromStack(stack).hasResourceProperties('AWS::IAM::Policy', {
+    PolicyDocument: {
+      Statement: [
+        {
+          Action: 'rekognition:detectLabels',
+          Effect: 'Allow',
+          Resource: '*',
+        },
+        {
+          Action: 's3:getObject',
+          Effect: 'Allow',
+          Resource: 'arn:aws:s3:::my-bucket/*',
+        },
+      ],
+      Version: '2012-10-17',
+    },
+  });
+});