From af53798d8fdd7d244da344585602f4f24c09806b Mon Sep 17 00:00:00 2001 From: Dillon O'Donovan <39768972+dillon-odonovan@users.noreply.github.com> Date: Fri, 18 Jun 2021 13:47:52 -0400 Subject: [PATCH] fix(stepfunctions-tasks): checking for task token in EcsRunTask containerOverrides 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* --- .../lib/ecs/run-ecs-task-base.ts | 4 +- .../lib/ecs/run-task.ts | 5 +- .../test/ecs/ecs-tasks.test.ts | 71 ++++++++++++++++++ .../test/ecs/run-tasks.test.ts | 73 +++++++++++++++++++ 4 files changed, 149 insertions(+), 4 deletions(-) diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/ecs/run-ecs-task-base.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/ecs/run-ecs-task-base.ts index bc56f1a57d16f..381c6ab324c73 100644 --- a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/ecs/run-ecs-task-base.ts +++ b/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 || []) { diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/ecs/run-task.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/ecs/run-task.ts index 60f0ede44be83..b53b7cbab9dcf 100644 --- a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/ecs/run-task.ts +++ b/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) { diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs/ecs-tasks.test.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs/ecs-tasks.test.ts index e5654a5e752fe..74344f8139758 100644 --- a/packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs/ecs-tasks.test.ts +++ b/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(); +}); diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs/run-tasks.test.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs/run-tasks.test.ts index 626d1000fde01..947b5cb0c8986 100644 --- a/packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs/run-tasks.test.ts +++ b/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(); +});