fix(stepfunctions-tasks): in-flight stepfunction executions fail on update of versioned lambdas

packages/@aws-cdk/aws-stepfunctions-tasks/lib/lambda/invoke.ts

@@ -1,5 +1,6 @@
 import * as iam from '@aws-cdk/aws-iam';
 import * as lambda from '@aws-cdk/aws-lambda';
+import { IAlias, IFunction, IVersion } from '@aws-cdk/aws-lambda';
 import * as sfn from '@aws-cdk/aws-stepfunctions';
 import * as cdk from '@aws-cdk/core';
 import { Construct } from 'constructs';
@@ -120,7 +121,7 @@ export class LambdaInvoke extends sfn.TaskStateBase {
 
     this.taskPolicies = [
       new iam.PolicyStatement({
-        resources: this.props.lambdaFunction.resourceArnsForGrantInvoke,
+        resources: this.determineResourceArnsForGrantInvoke(props.lambdaFunction),
         actions: ['lambda:InvokeFunction'],
       }),
     ];
@@ -161,6 +162,35 @@ export class LambdaInvoke extends sfn.TaskStateBase {
       };
     }
   }
+
+  /**
+   * Determine the ARN(s) to put into the resource field of the generated
+   * IAM policy based on the type of the provided lambda function.
+   *
+   * When invoking Lambda Versions, we need to grant permissions to all
+   * qualifiers. Otherwise in-flight StepFunction executions will fail with
+   * due to missing permissions. This is because a change of the referenced
+   * Version will cause the Policy to remove permissions for the previous
+   * Version - which is currently in-flight.
+   *
+   * @see https://github.com/aws/aws-cdk/issues/17515
+   */
+  private determineResourceArnsForGrantInvoke(lambdaFunction: IFunction): string[] {
+    if (isVersion(lambdaFunction)) {
+      return lambdaFunction.lambda.resourceArnsForGrantInvoke;
+    }
+
+    return lambdaFunction.resourceArnsForGrantInvoke;
+  }
+}
+
+/**
+ * Type guard to determine if a given `IFunction` implements IVersion
+ */
+function isVersion(lambdaFunction: IFunction | IAlias | IVersion): lambdaFunction is IVersion {
+  return !(lambdaFunction as IAlias).aliasName
+      && (lambdaFunction as IVersion).lambda
+      && Boolean((lambdaFunction as IVersion).version);
 }
 
 /**

packages/@aws-cdk/aws-stepfunctions-tasks/test/lambda/integ.invoke.qualifiers.ts

@@ -0,0 +1,80 @@
+import { Code, Function, Runtime } from '@aws-cdk/aws-lambda';
+import * as sfn from '@aws-cdk/aws-stepfunctions';
+import * as cdk from '@aws-cdk/core';
+import * as integ from '@aws-cdk/integ-tests';
+import { LambdaInvoke } from '../../lib';
+
+
+/*
+ * Creates a state machine with a task state to invoke a Lambda function
+ * via Alias and Version qualifiers.
+ *
+ * The state machine creates a couple of Lambdas that pass results forward
+ * and into a Choice state that validates the output.
+ */
+const app = new cdk.App();
+const stack = new cdk.Stack(app, 'aws-stepfunctions-tasks-lambda-invoke-integ');
+
+const submitJobLambda = new Function(stack, 'submitJobLambda', {
+  code: Code.fromInline(`exports.handler = async () => {
+        return {
+          statusCode: '200',
+          body: 'hello, world!'
+        };
+      };`),
+  runtime: Runtime.NODEJS_14_X,
+  handler: 'index.handler',
+});
+
+const submitJob = new LambdaInvoke(stack, 'Invoke Handler', {
+  lambdaFunction: submitJobLambda.currentVersion,
+  outputPath: '$.Payload',
+});
+
+const checkJobStateLambda = new Function(stack, 'checkJobStateLambda', {
+  code: Code.fromInline(`exports.handler = async function(event, context) {
+        return {
+          status: event.statusCode === '200' ? 'SUCCEEDED' : 'FAILED'
+        };
+  };`),
+  runtime: Runtime.NODEJS_14_X,
+  handler: 'index.handler',
+});
+const checkJobStateLambdaAlias = checkJobStateLambda.addAlias('IntegTest');
+
+const checkJobState = new LambdaInvoke(stack, 'Check the job state', {
+  lambdaFunction: checkJobStateLambdaAlias,
+  resultSelector: {
+    status: sfn.JsonPath.stringAt('$.Payload.status'),
+  },
+});
+
+const isComplete = new sfn.Choice(stack, 'Job Complete?');
+const jobFailed = new sfn.Fail(stack, 'Job Failed', {
+  cause: 'Job Failed',
+  error: 'Received a status that was not 200',
+});
+const finalStatus = new sfn.Pass(stack, 'Final step');
+
+const chain = sfn.Chain.start(submitJob)
+  .next(checkJobState)
+  .next(
+    isComplete
+      .when(sfn.Condition.stringEquals('$.status', 'FAILED'), jobFailed)
+      .when(sfn.Condition.stringEquals('$.status', 'SUCCEEDED'), finalStatus),
+  );
+
+const sm = new sfn.StateMachine(stack, 'StateMachine', {
+  definition: chain,
+  timeout: cdk.Duration.seconds(30),
+});
+
+new cdk.CfnOutput(stack, 'stateMachineArn', {
+  value: sm.stateMachineArn,
+});
+
+new integ.IntegTest(app, 'StepFunctionsInvokeLambdaQualifiersTest', {
+  testCases: [stack],
+});
+
+app.synth();

packages/@aws-cdk/aws-stepfunctions-tasks/test/lambda/invoke.qualifiers.integ.snapshot/StepFunctionsInvokeLambdaQualifiersTestDefaultTestDeployAssert57F7A2BF.template.json

@@ -0,0 +1 @@
+{}