From fbb941f848fc363c93ae79d899532ed4522f31b1 Mon Sep 17 00:00:00 2001
From: Jonathan Goldwasser <jogold@users.noreply.github.com>
Date: Wed, 21 Sep 2022 16:39:13 +0200
Subject: [PATCH] feat(stepfunctions-tasks): additional IAM statements for AWS
 SDK service integration (#22070)

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


----

### All Submissions:

* [x] Have you followed the guidelines in our [Contributing guide?](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md)

### Adding new Unconventional Dependencies:

* [ ] This PR adds new unconventional dependencies following the process described [here](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md/#adding-new-unconventional-dependencies)

### New Features

* [ ] Have you added the new feature to an [integration test](https://github.com/aws/aws-cdk/blob/main/INTEGRATION_TESTS.md)?
	* [ ] Did you use `yarn integ` to deploy the infrastructure and generate the snapshot (i.e. `yarn integ` without `--dry-run`)?

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
---
 .../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',
+    },
+  });
+});