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

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

diff --git a/agent/engine/task_manager.go b/agent/engine/task_manager.go
index 0c44f5d7fbe..c69273decba 100644
--- a/agent/engine/task_manager.go
+++ b/agent/engine/task_manager.go
@@ -350,7 +350,7 @@ 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)
 func (mtask *managedTask) cleanupCredentials() {
 	taskCredentialsID := mtask.GetCredentialsID()
 	if taskCredentialsID != "" {
@@ -1424,6 +1424,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 {
@@ -1453,6 +1454,12 @@ func (mtask *managedTask) cleanupTask(taskStoppedDuration time.Duration) {
 		return
 	}
 
+	// Remove TaskExecutionCredentials from credentialsManager
+	if taskExecutionCredentialsID != "" {
+		logger.Info("Cleaning up task's execution credentials")
+		mtask.credentialsManager.RemoveCredentials(taskExecutionCredentialsID)
+	}
+
 	logger.Info("Cleaning up task's containers and data", logger.Fields{
 		field.TaskARN: mtask.Arn,
 	})
diff --git a/agent/engine/task_manager_test.go b/agent/engine/task_manager_test.go
index b22958fcf9c..409bc8f4bab 100644
--- a/agent/engine/task_manager_test.go
+++ b/agent/engine/task_manager_test.go
@@ -43,6 +43,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/engine/dependencygraph"
 	mock_dockerstate "github.com/aws/amazon-ecs-agent/agent/engine/dockerstate/mocks"
 	mock_engine "github.com/aws/amazon-ecs-agent/agent/engine/mocks"
+	mock_credentials "github.com/aws/amazon-ecs-agent/agent/credentials/mocks"
 	"github.com/aws/amazon-ecs-agent/agent/engine/testdata"
 	"github.com/aws/amazon-ecs-agent/agent/eventstream"
 	"github.com/aws/amazon-ecs-agent/agent/sighandlers/exitcodes"
@@ -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)