feat: support s3:putobject in sdk and resultwriter integrations

Author: tvhees

lib/deploy/stepFunctions/compileIamRole.js

@@ -32,6 +32,9 @@ function getTaskStates(states, stateMachineName) {
         if (state.ItemReader) {
           taskStates.push(state.ItemReader);
         }
+        if (state.ResultWriter) {
+          taskStates.push(state.ResultWriter);
+        }
         return taskStates;
       }
       default: {
@@ -429,22 +432,24 @@ function getEventBridgePermissions(state) {
   ];
 }
 
-function getS3GetObjectPermissions(state) {
-  const bucket = state.Parameters['Bucket.$'] ? state.Parameters['Bucket.$'] : state.Parameters.Bucket;
-  const key = state.Parameters['Key.$'] ? state.Parameters['Key.$'] : state.Parameters.Key;
+function getS3ObjectPermissions(action, state) {
+  const bucket = state.Parameters.Bucket || '*';
+  const key = state.Parameters.Key || '*';
+  const prefix = state.Parameters.Prefix;
+  let arn;
 
-  if (bucket.startsWith('$') && key.startsWith('$')) {
-    return [{
-      action: 's3:GetObject',
-      resource: [
-        '*',
-      ],
-    }];
+  if (prefix) {
+    arn = `arn:aws:s3:::${bucket}/${prefix}/${key}`;
+  } else if (bucket === '*' && key === '*') {
+    arn = '*';
+  } else {
+    arn = `arn:aws:s3:::${bucket}/${key}`;
   }
+
   return [{
-    action: 's3:GetObject',
+    action,
     resource: [
-      `arn:aws:s3:::${bucket}/${key}`,
+      arn,
     ],
   }];
 }
@@ -564,7 +569,10 @@ function getIamPermissions(taskStates) {
 
       case 'arn:aws:states:::s3:getObject':
       case 'arn:aws:states:::aws-sdk:s3:getObject':
-        return getS3GetObjectPermissions(state);
+        return getS3ObjectPermissions('s3:GetObject', state);
+      case 'arn:aws:states:::s3:putObject':
+      case 'arn:aws:states:::aws-sdk:s3:putObject':
+        return getS3ObjectPermissions('s3:PutObject', state);
 
       default:
         if (isIntrinsic(state.Resource) || !!state.Resource.match(/arn:aws(-[a-z]+)*:lambda/)) {

lib/deploy/stepFunctions/compileIamRole.test.js

@@ -1452,7 +1452,7 @@ describe('#compileIamRole', () => {
     expectDenyAllPolicy(policy);
   });
 
-  it('should give s3:GetObject permission for only objects referenced by state machine', () => {
+  it('should give s3 permissions for only objects referenced by state machine', () => {
     const hello = 'hello.txt';
     const world = 'world.txt';
     const testBucket = 'test-bucket';
@@ -1469,6 +1469,16 @@ describe('#compileIamRole', () => {
               Bucket: bucket,
               Key: key,
             },
+            Next: 'B',
+          },
+          B: {
+            Type: 'Task',
+            Resource: 'arn:aws:states:::aws-sdk:s3:putObject',
+            Parameters: {
+              Bucket: bucket,
+              Key: key,
+              Body: {},
+            },
             End: true,
           },
         },
@@ -1491,6 +1501,14 @@ describe('#compileIamRole', () => {
       .to.be.deep.equal([`arn:aws:s3:::${testBucket}/${hello}`]);
     expect(policy2.PolicyDocument.Statement[0].Resource)
       .to.be.deep.equal([`arn:aws:s3:::${testBucket}/${world}`]);
+
+    [policy1, policy2].forEach((policy) => {
+      expect(policy.PolicyDocument.Statement[0].Action)
+        .to.be.deep.equal([
+          's3:GetObject',
+          's3:PutObject',
+        ]);
+    });
   });
 
   it('should give s3:GetObject permission for only objects referenced by state machine with ItemReader', () => {
@@ -1612,9 +1630,135 @@ describe('#compileIamRole', () => {
       .provider.compiledCloudFormationTemplate.Resources;
     const policy1 = resources.StateMachine1Role.Properties.Policies[0];
 
-    // even though some tasks target specific topic ARNs, other states use Bucket.$
+    // even though some tasks target specific values, other states use Bucket.$
     // and Key.$  so we need to give broad permissions to be able to get any
-    // table and key the input specifies
+    // bucket and key the input specifies
+    expect(policy1.PolicyDocument.Statement[1].Resource)
+      .to.be.deep.equal('*');
+  });
+
+  it('should give s3:PutObject permission for only objects referenced by state machine with ResultWriter', () => {
+    const hello = 'hello';
+    const world = 'world';
+    const testBucket = 'test-bucket';
+
+    const genStateMachine = (id, lambdaArn, bucket, prefix) => ({
+      id,
+      definition: {
+        StartAt: 'A',
+        States: {
+          A: {
+            Type: 'Map',
+            ItemProcessor: {
+              StartAt: 'B',
+              States: {
+                B: {
+                  Type: 'Task',
+                  Resource: lambdaArn,
+                  End: true,
+                },
+              },
+            },
+            ResultWriter: {
+              Resource: 'arn:aws:states:::s3:putObject',
+              Parameters: {
+                Bucket: bucket,
+                Prefix: prefix,
+              },
+            },
+            End: true,
+          },
+        },
+      },
+    });
+
+    serverless.service.stepFunctions = {
+      stateMachines: {
+        myStateMachine1: genStateMachine('StateMachine1',
+          'arn:aws:lambda:us-west-2:1234567890:function:foo', testBucket, hello),
+        myStateMachine2: genStateMachine('StateMachine2',
+          'arn:aws:lambda:us-west-2:1234567890:function:foo', testBucket, world),
+      },
+    };
+
+    serverlessStepFunctions.compileIamRole();
+    const resources = serverlessStepFunctions.serverless.service
+      .provider.compiledCloudFormationTemplate.Resources;
+    const policy1 = resources.StateMachine1Role.Properties.Policies[0];
+    const policy2 = resources.StateMachine2Role.Properties.Policies[0];
+    expect(policy1.PolicyDocument.Statement[1].Resource)
+      .to.be.deep.equal([`arn:aws:s3:::${testBucket}/${hello}/*`]);
+    expect(policy2.PolicyDocument.Statement[1].Resource)
+      .to.be.deep.equal([`arn:aws:s3:::${testBucket}/${world}/*`]);
+  });
+
+  it('should give s3:PutObject permission to * when Bucket.$ and Prefix.$ are seen on ResultWriter', () => {
+    const genStateMachine = (id, lambdaArn) => ({
+      id,
+      definition: {
+        StartAt: 'A',
+        States: {
+          A: {
+            Type: 'Map',
+            ItemProcessor: {
+              StartAt: 'B',
+              States: {
+                B: {
+                  Type: 'Task',
+                  Resource: lambdaArn,
+                  End: true,
+                },
+              },
+            },
+            ResultWriter: {
+              Resource: 'arn:aws:states:::s3:putObject',
+              Parameters: {
+                Bucket: 'test-bucket',
+                Prefix: 'test-prefix',
+              },
+            },
+            Next: 'C',
+          },
+          C: {
+            Type: 'Map',
+            ItemProcessor: {
+              StartAt: 'D',
+              States: {
+                D: {
+                  Type: 'Task',
+                  Resource: lambdaArn,
+                  End: true,
+                },
+              },
+            },
+            ResultWriter: {
+              Resource: 'arn:aws:states:::s3:putObject',
+              Parameters: {
+                'Bucket.$': '$.testBucket',
+                'Prefix.$': '$.prefix',
+              },
+            },
+            End: true,
+          },
+        },
+      },
+    });
+
+    serverless.service.stepFunctions = {
+      stateMachines: {
+        myStateMachine1: genStateMachine('StateMachine1',
+          'arn:aws:lambda:us-west-2:1234567890:function:foo'),
+      },
+    };
+
+    serverlessStepFunctions.compileIamRole();
+    const resources = serverlessStepFunctions.serverless.service
+      .provider.compiledCloudFormationTemplate.Resources;
+    const policy1 = resources.StateMachine1Role.Properties.Policies[0];
+
+    // even though some tasks target specific values, other states use Bucket.$
+    // and Prefix.$ so we need to give broad permissions to be able to write to
+    // any bucket and prefix the input specifies
     expect(policy1.PolicyDocument.Statement[1].Resource)
       .to.be.deep.equal('*');
   });