stepfunctions: ItemReader and ResultWriter do not grant KMS permissions

[BadMagic100]

Describe the bug

Generally, states added to a state machine will grant all necessary policy permissions to the execution role of the state machine. When using the DistributedMap construct, however, the necessary permissions for working with encrypted buckets are left out.

Regression Issue

• Select this option if this issue appears to be a regression.

Last Known Working CDK Version

No response

Expected Behavior

The step function should grant necessary KMS permissions on its returned policies.

Current Behavior

The step function does not grant necessary KMS permissions on its returned policies. Example error:

Failed to write a test manifest into the specified output bucket. Access denied. Ensure the correct permissions are added to your state machine's execution role. | Message from S3: User: is not authorized to perform: kms:GenerateDataKey on resource: because no identity-based policy allows the kms:GenerateDataKey action

Reproduction Steps

const inputBucket = new Bucket('DistributedMapInputBucket', {
    encryption: BucketEncryption.KMS,
    enfoceSSL: true
});
const outputBucket = new Bucket('DistributedMapOutputBucket', {
    encryption: BucketEncryption.KMS,
    enfoceSSL: true
});
new StateMachine(this, 'Test', {
    definitionBody: DefinitionBody.fromChainable(new DistributedMap(this, 'Map', {
        itemReader: new S3JsonItemReader({
            bucket: inputBucket,
            key: "some-object"
        }),
        resultWriter: new ResultWriter({
            bucket: outputBucket,
            prefix: "some-prefix"
       }),
    }).itemProcessor(new Pass(this, 'pass'))
});

Possible Solution

ResultWriter and ItemReader should add key read/write actions to their providePolicyStatements (both linked), similar to Bucket.grantRead and Bucket.grantWrite here

Additional Information/Context

No response

CDK CLI Version

2.164.1

Framework Version

No response

Node.js Version

18

OS

macOS

Language

TypeScript

Language Version

No response

Other information

No response

[khushail]

Hi @BadMagic100 , thanks for reaching out.

I tried to repro the scenario and it succeeded with synthesizing policies as well as deployment. Sharing the code as well as template snippet for your reference -

    const inputBucket = new Bucket(this,'DistributedMapInputBucket', {
      encryption: BucketEncryption.KMS,
      enforceSSL: true
    });
    const outputBucket = new Bucket(this, 'DistributedMapOutputBucket', {
        encryption: BucketEncryption.KMS,
        enforceSSL: true
    });

    const stateMachineObj = new StateMachine(this, 'Test', {
      definitionBody: DefinitionBody.fromChainable(new DistributedMap(this, 'Map', {
        itemReader: new S3JsonItemReader({
          bucket: inputBucket,
          key: "some-object"
        }),
        resultWriter: new ResultWriter({
          bucket: outputBucket,
          prefix: "some-prefix"
        }),
      }).itemProcessor(new Pass(this, 'pass')))
    });
    new CfnOutput(this, 'InputBucketName', { value: inputBucket.bucketName });
    new CfnOutput(this, 'OutputBucketName', { value: stateMachineObj.stateMachineArn });

Synthesized template for Policy -

"TestRoleDefaultPolicyAD214F97": {
      "Type": "AWS::IAM::Policy",
      "Properties": {
        "PolicyDocument": {
          "Statement": [
            {
              "Action": [
                "s3:AbortMultipartUpload",
                "s3:GetObject",
                "s3:ListMultipartUploadParts",
                "s3:PutObject"
              ],
              "Effect": "Allow",
              "Resource": {
                "Fn::Join": [
                  "",
                  [
                    "arn:",
                    {
                      "Ref": "AWS::Partition"
                    },
                    ":s3:::",
                    {
                      "Ref": "DistributedMapOutputBucket1136E660"
                    },
                    "/*"
                  ]
                ]
              }
            },
            {
              "Action": "s3:GetObject",
              "Effect": "Allow",
              "Resource": {
                "Fn::Join": [
                  "",
                  [
                    "arn:",
                    {
                      "Ref": "AWS::Partition"
                    },
                    ":s3:::",
                    {
                      "Ref": "DistributedMapInputBucketFDAB2D5D"
                    },
                    "/*"
                  ]
                ]
              }
            }
          ],
          "Version": "2012-10-17"
        },
        "PolicyName": "TestRoleDefaultPolicyAD214F97",
        "Roles": [
          {
            "Ref": "TestRole17AB2208"
          }
        ]
      },
      "Metadata": {
        "aws:cdk:path": "StepfunctionIssueStack/Test/Role/DefaultPolicy/Resource"
      }
    },
    "Test7BFAF513": {
      "Type": "AWS::StepFunctions::StateMachine",
      "Properties": {
        "DefinitionString": {
          "Fn::Join": [
            "",
            [
              "{\"StartAt\":\"Map\",\"States\":{\"Map\":{\"Type\":\"Map\",\"End\":true,\"ItemProcessor\":{\"ProcessorConfig\":{\"Mode\":\"DISTRIBUTED\",\"ExecutionType\":\"STANDARD\"},\"StartAt\":\"pass\",\"States\":{\"pass\":{\"Type\":\"Pass\",\"End\":true}}},\"ItemReader\":{\"Resource\":\"arn:",
              {
                "Ref": "AWS::Partition"
              },
              ":states:::s3:getObject\",\"ReaderConfig\":{\"InputType\":\"JSON\"},\"Parameters\":{\"Bucket\":\"",
              {
                "Ref": "DistributedMapInputBucketFDAB2D5D"
              },
              "\",\"Key\":\"some-object\"}},\"ResultWriter\":{\"Resource\":\"arn:",
              {
                "Ref": "AWS::Partition"
              },
              ":states:::s3:putObject\",\"Parameters\":{\"Bucket\":\"",
              {
                "Ref": "DistributedMapOutputBucket1136E660"
              },
              "\",\"Prefix\":\"some-prefix\"}}}}}"
            ]
          ]
        },
        "RoleArn": {
          "Fn::GetAtt": [
            "TestRole17AB2208",
            "Arn"
          ]
        }
      },
      "DependsOn": [
        "TestRoleDefaultPolicyAD214F97",
        "TestRole17AB2208"
      ],
      "UpdateReplacePolicy": "Delete",
      "DeletionPolicy": "Delete",
      "Metadata": {
        "aws:cdk:path": "StepfunctionIssueStack/Test/Resource"
      }
    },

In statemachine above, the ItemReader and ResultWrite have the required permissions needed to access the bucket.

Deployment also succeeds without error -

I assume this is what you are looking for. Let me know if this does not help or your usecase and ask is different

[BadMagic100]

Yes, this repro is right and you are correct that it can synthesize and deploy. What you will find is if you try to execute the step function it will fail because the buckets are encrypted and the step function execution role cannot read/write from them. The permissions for kms:DescribeKey and kms:Decrypt should be granted for the input bucket's key (from the item reader), and kms:Encrypt, kms:ReEncrypt, kms:Decrypt, and kms:GenerateDataKey should be granted for the output bucket's key (from the result writer) so that the state machine can execute correctly out of the box.

For now I am using a workaround like below to get the correct policies.

inputBucket.encryptionKey?.grant(stateMachineObj, 'kms:DescribeKey');
inputBucket.encryptionKey?.grantEncrypt(stateMachineObj);
outputBucket.encryptionKey?.grantEncryptDecrypt(ststeMachineObj);