fix(stepfunctions-tasks): checking for task token in EcsRunTask conta…

…inerOverrides causes memory explosion (#15187)

EcsRunTask containerOverrides had references to the entire construct tree causing a large amount of recursion
when trying to validate that the task token was provided in the waitForTaskToken integration pattern. The only place for the task token to appear is in the environment of a containerOverride, so we should only check for it there.

I also updated the error to mention JsonPath instead of Context since Context is deprecated.

In addition to the unit tests, I packaged the library to dotnet and succesfully synthed the minimal reproduction.

thanks to @BenChaimberg for providing the suggested fix

fixes #15124

----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*

packages/@aws-cdk/aws-stepfunctions-tasks/lib/ecs/run-ecs-task-base.ts

@@ -84,8 +84,8 @@ export class EcsRunTaskBase implements ec2.IConnectable, sfn.IStepFunctionsTask
  }
 
  if (this.integrationPattern === sfn.ServiceIntegrationPattern.WAIT_FOR_TASK_TOKEN
- && !sfn.FieldUtils.containsTaskToken(props.containerOverrides)) {
- throw new Error('Task Token is missing in containerOverrides (pass JsonPath.taskToken somewhere in containerOverrides)');
+ && !sfn.FieldUtils.containsTaskToken(props.containerOverrides?.map(override => override.environment))) {
+ throw new Error('Task Token is required in at least one `containerOverrides.environment` for callback. Use JsonPath.taskToken to set the token.');
  }
 
  for (const override of this.props.containerOverrides || []) {

packages/@aws-cdk/aws-stepfunctions-tasks/lib/ecs/run-task.ts

@@ -242,8 +242,9 @@ export class EcsRunTask extends sfn.TaskStateBase implements ec2.IConnectable {
 
  validatePatternSupported(this.integrationPattern, EcsRunTask.SUPPORTED_INTEGRATION_PATTERNS);
 
- if (this.integrationPattern === sfn.IntegrationPattern.WAIT_FOR_TASK_TOKEN && !sfn.FieldUtils.containsTaskToken(props.containerOverrides)) {
- throw new Error('Task Token is required in `containerOverrides` for callback. Use Context.taskToken to set the token.');
+ if (this.integrationPattern === sfn.IntegrationPattern.WAIT_FOR_TASK_TOKEN
+ && !sfn.FieldUtils.containsTaskToken(props.containerOverrides?.map(override => override.environment))) {
+ throw new Error('Task Token is required in at least one `containerOverrides.environment` for callback. Use JsonPath.taskToken to set the token.');
  }
 
  if (!this.props.taskDefinition.defaultContainer) {

packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs/ecs-tasks.test.ts

@@ -373,3 +373,74 @@ test('Running an EC2 Task with overridden number values', () => {
  Type: 'Task',
  });
 });
+
+test('Cannot create a task with WAIT_FOR_TASK_TOKEN if no TaskToken provided', () => {
+ const taskDefinition = new ecs.TaskDefinition(stack, 'TaskDefinition', {
+ compatibility: ecs.Compatibility.EC2,
+ });
+
+ const containerDefinition = taskDefinition.addContainer('ContainerDefinition', {
+ image: ecs.ContainerImage.fromRegistry('foo/bar'),
+ });
+
+ expect(() =>
+ new tasks.RunEcsEc2Task({
+ cluster,
+ containerOverrides: [
+ {
+ containerDefinition,
+ environment: [
+ {
+ name: 'Foo',
+ value: 'Bar',
+ },
+ ],
+ },
+ ],
+ integrationPattern: sfn.ServiceIntegrationPattern.WAIT_FOR_TASK_TOKEN,
+ taskDefinition,
+ }),
+ ).toThrowError(/Task Token is required in at least one `containerOverrides.environment`/);
+});
+
+test('Running a task with WAIT_FOR_TASK_TOKEN and task token in environment', () => {
+ const taskDefinition = new ecs.TaskDefinition(stack, 'TaskDefinition', {
+ compatibility: ecs.Compatibility.EC2,
+ });
+
+ const primaryContainerDef = taskDefinition.addContainer('PrimaryContainerDef', {
+ image: ecs.ContainerImage.fromRegistry('foo/primary'),
+ essential: true,
+ });
+
+ const sidecarContainerDef = taskDefinition.addContainer('SideCarContainerDef', {
+ image: ecs.ContainerImage.fromRegistry('foo/sidecar'),
+ essential: false,
+ });
+
+ expect(() => new tasks.RunEcsEc2Task({
+ cluster,
+ containerOverrides: [
+ {
+ containerDefinition: primaryContainerDef,
+ environment: [
+ {
+ name: 'Foo',
+ value: 'Bar',
+ },
+ ],
+ },
+ {
+ containerDefinition: sidecarContainerDef,
+ environment: [
+ {
+ name: 'TaskToken.$',
+ value: sfn.JsonPath.taskToken,
+ },
+ ],
+ },
+ ],
+ integrationPattern: sfn.ServiceIntegrationPattern.WAIT_FOR_TASK_TOKEN,
+ taskDefinition,
+ })).not.toThrow();
+});

packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs/run-tasks.test.ts

@@ -411,3 +411,76 @@ test('Running an EC2 Task with overridden number values', () => {
  Type: 'Task',
  });
 });
+
+test('Cannot create a task with WAIT_FOR_TASK_TOKEN if no TaskToken provided', () => {
+ const taskDefinition = new ecs.TaskDefinition(stack, 'TaskDefinition', {
+ compatibility: ecs.Compatibility.EC2,
+ });
+
+ const containerDefinition = taskDefinition.addContainer('ContainerDefinition', {
+ image: ecs.ContainerImage.fromRegistry('foo/bar'),
+ });
+
+ expect(() =>
+ new tasks.EcsRunTask(stack, 'RunTask', {
+ cluster,
+ containerOverrides: [
+ {
+ containerDefinition,
+ environment: [
+ {
+ name: 'Foo',
+ value: 'Bar',
+ },
+ ],
+ },
+ ],
+ integrationPattern: sfn.IntegrationPattern.WAIT_FOR_TASK_TOKEN,
+ launchTarget: new tasks.EcsEc2LaunchTarget(),
+ taskDefinition,
+ }),
+ ).toThrowError(/Task Token is required in at least one `containerOverrides.environment`/);
+});
+
+test('Running a task with WAIT_FOR_TASK_TOKEN and task token in environment', () => {
+ const taskDefinition = new ecs.TaskDefinition(stack, 'TaskDefinition', {
+ compatibility: ecs.Compatibility.EC2,
+ });
+
+ const primaryContainerDef = taskDefinition.addContainer('PrimaryContainerDef', {
+ image: ecs.ContainerImage.fromRegistry('foo/primary'),
+ essential: true,
+ });
+
+ const sidecarContainerDef = taskDefinition.addContainer('SideCarContainerDef', {
+ image: ecs.ContainerImage.fromRegistry('foo/sidecar'),
+ essential: false,
+ });
+
+ expect(() => new tasks.EcsRunTask(stack, 'RunTask', {
+ cluster,
+ containerOverrides: [
+ {
+ containerDefinition: primaryContainerDef,
+ environment: [
+ {
+ name: 'Foo',
+ value: 'Bar',
+ },
+ ],
+ },
+ {
+ containerDefinition: sidecarContainerDef,
+ environment: [
+ {
+ name: 'TaskToken.$',
+ value: sfn.JsonPath.taskToken,
+ },
+ ],
+ },
+ ],
+ integrationPattern: sfn.IntegrationPattern.WAIT_FOR_TASK_TOKEN,
+ launchTarget: new tasks.EcsEc2LaunchTarget(),
+ taskDefinition,
+ })).not.toThrow();
+});