From f52f39179b4eecc0b00d12a2c01d9c6c142f64fd Mon Sep 17 00:00:00 2001
From: Heming Han <hanhm@amazon.com>
Date: Fri, 27 Aug 2021 11:46:54 -0700
Subject: [PATCH] Fix a issue that agent does not clean task execution
 credentials from credential manager

---
 agent/engine/task_manager.go      | 12 +++++-
 agent/engine/task_manager_test.go | 70 +++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/agent/engine/task_manager.go b/agent/engine/task_manager.go
index 0c44f5d7fb..bba3b40aca 100644
--- a/agent/engine/task_manager.go
+++ b/agent/engine/task_manager.go
@@ -350,7 +350,8 @@ func (mtask *managedTask) steadyState() bool {
 	}
 }
 
-// cleanupCredentials removes credentials for a stopped task
+// cleanupCredentials removes credentials for a stopped task (execution credentials are removed in cleanupTask
+// due to its potential usage in the later phase of the task cleanup such as sending logs)
 func (mtask *managedTask) cleanupCredentials() {
 	taskCredentialsID := mtask.GetCredentialsID()
 	if taskCredentialsID != "" {
@@ -1424,6 +1425,7 @@ func (mtask *managedTask) time() ttime.Time {
 }
 
 func (mtask *managedTask) cleanupTask(taskStoppedDuration time.Duration) {
+	taskExecutionCredentialsID := mtask.GetExecutionCredentialsID()
 	cleanupTimeDuration := mtask.GetKnownStatusTime().Add(taskStoppedDuration).Sub(ttime.Now())
 	cleanupTime := make(<-chan time.Time)
 	if cleanupTimeDuration < 0 {
@@ -1464,6 +1466,14 @@ func (mtask *managedTask) cleanupTask(taskStoppedDuration time.Duration) {
 	mtask.engine.sweepTask(mtask.Task)
 	mtask.engine.deleteTask(mtask.Task)
 
+	// Remove TaskExecutionCredentials from credentialsManager
+	if taskExecutionCredentialsID != "" {
+		logger.Info("Cleaning up task's execution credentials", logger.Fields{
+			field.TaskARN: mtask.Arn,
+		})
+		mtask.credentialsManager.RemoveCredentials(taskExecutionCredentialsID)
+	}
+
 	// The last thing to do here is to cancel the context, which should cancel
 	// all outstanding go routines associated with this managed task.
 	mtask.cancel()
diff --git a/agent/engine/task_manager_test.go b/agent/engine/task_manager_test.go
index b22958fcf9..6bb494e41b 100644
--- a/agent/engine/task_manager_test.go
+++ b/agent/engine/task_manager_test.go
@@ -38,6 +38,7 @@ import (
 	apitaskstatus "github.com/aws/amazon-ecs-agent/agent/api/task/status"
 	"github.com/aws/amazon-ecs-agent/agent/config"
 	"github.com/aws/amazon-ecs-agent/agent/credentials"
+	mock_credentials "github.com/aws/amazon-ecs-agent/agent/credentials/mocks"
 	"github.com/aws/amazon-ecs-agent/agent/data"
 	mock_dockerapi "github.com/aws/amazon-ecs-agent/agent/dockerclient/dockerapi/mocks"
 	"github.com/aws/amazon-ecs-agent/agent/engine/dependencygraph"
@@ -1335,6 +1336,75 @@ func TestTaskWaitForExecutionCredentials(t *testing.T) {
 	}
 }
 
+func TestCleanupTaskWithExecutionCredentials(t *testing.T) {
+	cfg := getTestConfig()
+	ctrl := gomock.NewController(t)
+	mockTime := mock_ttime.NewMockTime(ctrl)
+	mockState := mock_dockerstate.NewMockTaskEngineState(ctrl)
+	mockClient := mock_dockerapi.NewMockDockerClient(ctrl)
+	mockImageManager := mock_engine.NewMockImageManager(ctrl)
+	mockCredentialsManager := mock_credentials.NewMockManager(ctrl)
+	mockResource := mock_taskresource.NewMockTaskResource(ctrl)
+	defer ctrl.Finish()
+
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
+
+	taskEngine := &DockerTaskEngine{
+		ctx:                ctx,
+		cfg:                &cfg,
+		dataClient:         data.NewNoopClient(),
+		state:              mockState,
+		client:             mockClient,
+		imageManager:       mockImageManager,
+		credentialsManager: mockCredentialsManager,
+	}
+	mTask := &managedTask{
+		ctx:                      ctx,
+		cancel:                   cancel,
+		Task:                     testdata.LoadTask("sleep5"),
+		credentialsManager:       mockCredentialsManager,
+		_time:                    mockTime,
+		engine:                   taskEngine,
+		acsMessages:              make(chan acsTransition),
+		dockerMessages:           make(chan dockerContainerChange),
+		resourceStateChangeEvent: make(chan resourceStateChange),
+		cfg:                      taskEngine.cfg,
+	}
+
+	mTask.Task.ResourcesMapUnsafe = make(map[string][]taskresource.TaskResource)
+	mTask.Task.SetExecutionRoleCredentialsID("executionRoleCredentialsId")
+	mTask.AddResource("mockResource", mockResource)
+	mTask.SetKnownStatus(apitaskstatus.TaskStopped)
+	mTask.SetSentStatus(apitaskstatus.TaskStopped)
+	container := mTask.Containers[0]
+	dockerContainer := &apicontainer.DockerContainer{
+		DockerName: "dockerContainer",
+	}
+
+	// Expectations for triggering cleanup
+	now := mTask.GetKnownStatusTime()
+	taskStoppedDuration := 1 * time.Minute
+	mockTime.EXPECT().Now().Return(now).AnyTimes()
+	cleanupTimeTrigger := make(chan time.Time)
+	mockTime.EXPECT().After(gomock.Any()).Return(cleanupTimeTrigger)
+	go func() {
+		cleanupTimeTrigger <- now
+	}()
+
+	// Expectations to verify the execution credentials get removed
+	mockCredentialsManager.EXPECT().RemoveCredentials("executionRoleCredentialsId")
+
+	// Expectations to verify that the task gets removed
+	mockState.EXPECT().ContainerMapByArn(mTask.Arn).Return(map[string]*apicontainer.DockerContainer{container.Name: dockerContainer}, true)
+	mockClient.EXPECT().RemoveContainer(gomock.Any(), dockerContainer.DockerName, gomock.Any()).Return(nil)
+	mockImageManager.EXPECT().RemoveContainerReferenceFromImageState(container).Return(nil)
+	mockState.EXPECT().RemoveTask(mTask.Task)
+	mockResource.EXPECT().Cleanup()
+	mockResource.EXPECT().GetName()
+	mTask.cleanupTask(taskStoppedDuration)
+}
+
 func TestCleanupTaskWithInvalidInterval(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	mockTime := mock_ttime.NewMockTime(ctrl)