fix tmds creds response

Author: singholt

agent/api/task/task.go

@@ -245,10 +245,10 @@ type Task struct {
 	// perform some action at the task level, such as pulling image from ECR
 	ExecutionCredentialsID string `json:"executionCredentialsID"`
 
-	// credentialsID is used to set the CredentialsId field for the
+	// CredentialsID is used to set the CredentialsId field for the
 	// IAMRoleCredentials object associated with the task. This id can be
 	// used to look up the credentials for task in the credentials manager
-	credentialsID                string
+	CredentialsID                string `json:"credentialsID"`
 	credentialsRelativeURIUnsafe string
 
 	// ENIs is the list of Elastic Network Interfaces assigned to this task. The
@@ -2805,15 +2805,15 @@ func (task *Task) SetCredentialsID(id string) {
 	task.lock.Lock()
 	defer task.lock.Unlock()
 
-	task.credentialsID = id
+	task.CredentialsID = id
 }
 
 // GetCredentialsID gets the credentials ID for the task
 func (task *Task) GetCredentialsID() string {
 	task.lock.RLock()
 	defer task.lock.RUnlock()
 
-	return task.credentialsID
+	return task.CredentialsID
 }
 
 // SetCredentialsRelativeURI sets the credentials relative uri for the task

agent/api/task/task_test.go

@@ -1004,7 +1004,7 @@ func TestGetCredentialsEndpointWhenCredentialsAreSet(t *testing.T) {
 				Name:        "c2",
 				Environment: make(map[string]string),
 			}},
-		credentialsID: credentialsIDInTask,
+		CredentialsID: credentialsIDInTask,
 	}
 
 	taskCredentials := credentials.TaskIAMRoleCredentials{

agent/api/task/taskvolume_test.go

@@ -121,6 +121,7 @@ func TestMarshalTaskVolumesEFS(t *testing.T) {
 		"ExecutionStoppedAt": "0001-01-01T00:00:00Z",
 		"SentStatus": "NONE",
 		"executionCredentialsID": "",
+		"credentialsID": "",
 		"ENI": null,
 		"AppMesh": null,
 		"PlatformFields": %s
@@ -168,6 +169,7 @@ func TestUnmarshalTaskVolumesEFS(t *testing.T) {
 		"ExecutionStoppedAt": "0001-01-01T00:00:00Z",
 		"SentStatus": "NONE",
 		"executionCredentialsID": "",
+		"credentialsID": "",
 		"ENI": null,
 		"AppMesh": null,
 		"PlatformFields": {}
@@ -241,6 +243,7 @@ func TestMarshalEBSVolumes(t *testing.T) {
 		"ExecutionStoppedAt": "0001-01-01T00:00:00Z",
 		"SentStatus": "NONE",
 		"executionCredentialsID": "",
+		"credentialsID": "",
 		"ENI": null,
 		"AppMesh": null,
 		"PlatformFields": %s
@@ -286,6 +289,7 @@ func TestUnmarshalEBSVolumes(t *testing.T) {
 		"ExecutionStoppedAt": "0001-01-01T00:00:00Z",
 		"SentStatus": "NONE",
 		"executionCredentialsID": "",
+		"credentialsID": "",
 		"ENI": null,
 		"AppMesh": null,
 		"PlatformFields": {}

agent/api/task/taskvolume_windows_test.go

@@ -78,6 +78,7 @@ func TestMarshalTaskVolumeFSxWindowsFileServer(t *testing.T) {
 		"ExecutionStoppedAt": "0001-01-01T00:00:00Z",
 		"SentStatus": "NONE",
 		"executionCredentialsID": "",
+		"credentialsID": "",
 		"ENI": null,
 		"AppMesh": null,
 		"PlatformFields": %s
@@ -117,6 +118,7 @@ func TestUnmarshalTaskVolumeFSxWindowsFileServer(t *testing.T) {
 		"ExecutionStoppedAt": "0001-01-01T00:00:00Z",
 		"SentStatus": "NONE",
 		"executionCredentialsID": "",
+		"credentialsID": "",
 		"ENI": null,
 		"AppMesh": null,
 		"PlatformFields": {}

agent/engine/data.go

@@ -56,6 +56,16 @@ func (engine *DockerTaskEngine) loadTasks() error {
 	for _, task := range tasks {
 		engine.state.AddTask(task)
 
+		// Register the task role's ID in credentials manager as soon as possible.
+		// This ensures TMDS can distinguish between invalid credentials requests (400) and
+		// known credentials that aren't available yet (503) after agent restart.
+		if credentialsID := task.GetCredentialsID(); credentialsID != "" {
+			seelog.Infof("loadTasks: Registering known credentials ID: %s for task: %s", credentialsID, task.Arn)
+			engine.credentialsManager.AddKnownCredentialsID(credentialsID)
+		} else {
+			seelog.Debugf("loadTasks: Task %s has no credentials ID", task.Arn)
+		}
+
 		// TODO: Will need to clean up all of the STOPPED managed daemon tasks
 		md, ok := task.IsManagedDaemonTask()
 		if ok {

agent/engine/data_test.go

@@ -27,8 +27,10 @@ import (
 	"github.com/aws/amazon-ecs-agent/ecs-agent/api/attachment"
 	apicontainerstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status"
 	apitaskstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/task/status"
+	mock_credentials "github.com/aws/amazon-ecs-agent/ecs-agent/credentials/mocks"
 	ni "github.com/aws/amazon-ecs-agent/ecs-agent/netlib/model/networkinterface"
 
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -376,3 +378,40 @@ func TestRemoveENIAttachmentData(t *testing.T) {
 	require.NoError(t, err)
 	assert.Len(t, res, 0)
 }
+
+func TestLoadTasksRegistersCredentialsID(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	dataClient := newTestDataClient(t)
+	mockCredentialsManager := mock_credentials.NewMockManager(ctrl)
+
+	// Create a task with credentials ID
+	testCredentialsID := "test-credentials-id"
+	taskWithCredentials := &apitask.Task{
+		Arn:        testTaskARN,
+		Containers: []*apicontainer.Container{testContainer},
+	}
+	// Set the credentials ID after task creation
+	taskWithCredentials.SetCredentialsID(testCredentialsID)
+
+	// Save task to data client
+	require.NoError(t, dataClient.SaveTask(taskWithCredentials))
+
+	engine := &DockerTaskEngine{
+		state:              dockerstate.NewTaskEngineState(),
+		dataClient:         dataClient,
+		credentialsManager: mockCredentialsManager,
+	}
+
+	// Expect credentials ID to be added to known credentials during loadTasks
+	mockCredentialsManager.EXPECT().AddKnownCredentialsID(testCredentialsID)
+
+	// Call loadTasks and verify credentials ID is registered
+	require.NoError(t, engine.loadTasks())
+
+	// Verify task was loaded into state
+	task, ok := engine.state.TaskByArn(testTaskARN)
+	assert.True(t, ok)
+	assert.Equal(t, testCredentialsID, task.GetCredentialsID())
+}

agent/handlers/task_server_setup_test.go

@@ -1106,7 +1106,17 @@ func getResponseForCredentialsRequest(t *testing.T, expectedStatus int,
 	recorder := httptest.NewRecorder()
 
 	creds, ok := getCredentials()
-	credentialsManager.EXPECT().GetTaskCredentials(gomock.Any()).Return(creds, ok)
+
+	// For credentials uninitialized tests, IsCredentialsPending should return true
+	if expectedStatus == http.StatusServiceUnavailable && ok {
+		credentialsManager.EXPECT().IsCredentialsPending(gomock.Any()).Return(true)
+		// GetTaskCredentials should not be called since IsCredentialsPending returns true
+	} else {
+		// For all other cases, IsCredentialsPending should return false
+		credentialsManager.EXPECT().IsCredentialsPending(gomock.Any()).Return(false)
+		credentialsManager.EXPECT().GetTaskCredentials(gomock.Any()).Return(creds, ok)
+	}
+
 	auditLog.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any())
 
 	params := make(url.Values)