From d31d658a7d006a81060708dac8938f0792916fb9 Mon Sep 17 00:00:00 2001
From: Torbjorn van Heeswijck <torbjorn.vanheeswijck@aligent.com.au>
Date: Fri, 5 May 2023 13:36:08 +0930
Subject: [PATCH] feat: add support dynamoDB batch tasks using AWS SDK service
 integration

---
 lib/deploy/stepFunctions/compileIamRole.js    |  27 ++++
 .../stepFunctions/compileIamRole.test.js      | 116 ++++++++++++++++++
 2 files changed, 143 insertions(+)

diff --git a/lib/deploy/stepFunctions/compileIamRole.js b/lib/deploy/stepFunctions/compileIamRole.js
index 58fd7812..9710db87 100644
--- a/lib/deploy/stepFunctions/compileIamRole.js
+++ b/lib/deploy/stepFunctions/compileIamRole.js
@@ -219,6 +219,28 @@ function getDynamoDBPermissions(action, state) {
     resource,
   }];
 }
+
+function getBatchDynamoDBPermissions(action, state) {
+  if (state.Parameters['RequestItems.$']) {
+    // When the RequestItems object is only known at runtime,
+    // we have to provide * permissions during deployment.
+    return [{
+      action,
+      resource: '*',
+    }];
+  }
+  // If RequestItems is specified it must contain the target
+  // table names as keys. We can use these to generate roles
+  // whether the array of requests for that table is known
+  // at deploy time or not
+  const tableNames = Object.keys(state.Parameters.RequestItems);
+
+  return tableNames.map(tableName => ({
+    action,
+    resource: getDynamoDBArn(tableName.replace('.$', '')),
+  }));
+}
+
 function getRedshiftDataPermissions(action, state) {
   if (['redshift-data:ExecuteStatement', 'redshift-data:BatchExecuteStatement'].includes(action)) {
     const clusterName = _.has(state, 'Parameters.ClusterIdentifier') ? state.Parameters.ClusterIdentifier : '*';
@@ -515,6 +537,11 @@ function getIamPermissions(taskStates) {
       case 'arn:aws:states:::aws-sdk:dynamodb:query':
         return getDynamoDBPermissions('dynamodb:Query', state);
 
+      case 'arn:aws:states:::aws-sdk:dynamodb:batchGetItem':
+        return getBatchDynamoDBPermissions('dynamodb:BatchGetItem', state);
+      case 'arn:aws:states:::aws-sdk:dynamodb:batchWriteItem':
+        return getBatchDynamoDBPermissions('dynamodb:BatchWriteItem', state);
+
       case 'arn:aws:states:::aws-sdk:redshiftdata:executeStatement':
         return getRedshiftDataPermissions('redshift-data:ExecuteStatement', state);
       case 'arn:aws:states:::aws-sdk:redshiftdata:batchExecuteStatement':
diff --git a/lib/deploy/stepFunctions/compileIamRole.test.js b/lib/deploy/stepFunctions/compileIamRole.test.js
index 303a8902..3d1735c6 100644
--- a/lib/deploy/stepFunctions/compileIamRole.test.js
+++ b/lib/deploy/stepFunctions/compileIamRole.test.js
@@ -908,6 +908,122 @@ describe('#compileIamRole', () => {
     expect(policy.PolicyDocument.Statement[0].Resource[0]).to.equal('*');
   });
 
+  it('should give batch dynamodb permission for only tables referenced by state machine', () => {
+    const helloTable = 'hello';
+    const helloTableArn = {
+      'Fn::Join': [
+        ':', ['arn', { Ref: 'AWS::Partition' }, 'dynamodb', { Ref: 'AWS::Region' }, { Ref: 'AWS::AccountId' }, 'table/hello'],
+      ],
+    };
+    const worldTable = 'world';
+    const worldTableArn = {
+      'Fn::Join': [
+        ':', ['arn', { Ref: 'AWS::Partition' }, 'dynamodb', { Ref: 'AWS::Region' }, { Ref: 'AWS::AccountId' }, 'table/world'],
+      ],
+    };
+
+    const genStateMachine = (id, tableName) => ({
+      id,
+      definition: {
+        StartAt: 'A',
+        States: {
+          A: {
+            Type: 'Task',
+            Resource: 'arn:aws:states:::aws-sdk:dynamodb:batchWriteItem',
+            Parameters: {
+              RequestItems: {
+                [tableName]: [],
+              },
+            },
+            Next: 'B',
+          },
+          B: {
+            Type: 'Task',
+            Resource: 'arn:aws:states:::aws-sdk:dynamodb:batchGetItem',
+            Parameters: {
+              RequestItems: {
+                [tableName]: {},
+              },
+            },
+            End: true,
+          },
+        },
+      },
+    });
+    serverless.service.stepFunctions = {
+      stateMachines: {
+        myStateMachine1: genStateMachine('StateMachine1', helloTable),
+        myStateMachine2: genStateMachine('StateMachine2', worldTable),
+      },
+    };
+
+    serverlessStepFunctions.compileIamRole();
+    const resources = serverlessStepFunctions.serverless.service
+      .provider.compiledCloudFormationTemplate.Resources;
+    const policy1 = resources.StateMachine1Role.Properties.Policies[0];
+    const policy2 = resources.StateMachine2Role.Properties.Policies[0];
+
+    [policy1, policy2].forEach((policy) => {
+      expect(policy.PolicyDocument.Statement[0].Action)
+        .to.be.deep.equal([
+          'dynamodb:BatchWriteItem',
+          'dynamodb:BatchGetItem',
+        ]);
+    });
+
+    expect(policy1.PolicyDocument.Statement[0].Resource)
+      .to.be.deep.equal([helloTableArn]);
+    expect(policy2.PolicyDocument.Statement[0].Resource)
+      .to.be.deep.equal([worldTableArn]);
+  });
+
+  it('should give batch dynamodb permission to * whenever RequestItems.$ is seen', () => {
+    const genStateMachine = id => ({
+      id,
+      definition: {
+        StartAt: 'A',
+        States: {
+          A: {
+            Type: 'Task',
+            Resource: 'arn:aws:states:::aws-sdk:dynamodb:batchWriteItem',
+            Parameters: {
+              RequestItems: {
+                tableName: [],
+              },
+            },
+            Next: 'B',
+          },
+          B: {
+            Type: 'Task',
+            Resource: 'arn:aws:states:::aws-sdk:dynamodb:batchWriteItem',
+            Parameters: {
+              'RequestItems.$': '$.requestItems',
+            },
+            End: true,
+          },
+        },
+      },
+    });
+
+    serverless.service.stepFunctions = {
+      stateMachines: {
+        myStateMachine1: genStateMachine('StateMachine1'),
+      },
+    };
+
+    serverlessStepFunctions.compileIamRole();
+    const policy = serverlessStepFunctions.serverless.service
+      .provider.compiledCloudFormationTemplate.Resources.StateMachine1Role
+      .Properties.Policies[0];
+    expect(policy.PolicyDocument.Statement[0].Action)
+      .to.be.deep.equal(['dynamodb:BatchWriteItem']);
+
+    // even though some tasks target specific tables, because RequestItems.$ is used we
+    // have to give broad permissions to allow execution to talk to whatever table
+    // the input specifies
+    expect(policy.PolicyDocument.Statement[0].Resource).to.equal('*');
+  });
+
   it('should give Redshift Data permissions to * for safe actions', () => {
     serverless.service.stepFunctions = {
       stateMachines: {