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/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/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/credentials/interface.go

@@ -20,4 +20,6 @@ type Manager interface {
 	SetTaskCredentials(*TaskIAMRoleCredentials) error
 	GetTaskCredentials(string) (TaskIAMRoleCredentials, bool)
 	RemoveCredentials(string)
+	IsCredentialsPending(string) bool
+	AddKnownCredentialsID(string)
 }

agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/credentials/manager.go

@@ -18,6 +18,7 @@ import (
 	"sync"
 
 	"github.com/aws/amazon-ecs-agent/ecs-agent/acs/model/ecsacs"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/logger"
 	"github.com/aws/aws-sdk-go-v2/aws"
 )
 
@@ -86,7 +87,10 @@ func (roleCredentials *IAMRoleCredentials) GenerateCredentialsEndpointRelativeUR
 // the credentials endpoint
 type credentialsManager struct {
 	// idToTaskCredentials maps credentials id to its corresponding TaskIAMRoleCredentials object
+	// this map contains credentials id for which agent received the actual credentials from ACS already
 	idToTaskCredentials map[string]TaskIAMRoleCredentials
+	// knownCredentialsIDs tracks all credentials IDs we know about
+	knownCredentialsIDs map[string]bool
 	taskCredentialsLock sync.RWMutex
 }
 
@@ -108,6 +112,7 @@ func IAMRoleCredentialsFromACS(roleCredentials *ecsacs.IAMRoleCredentials, roleT
 func NewManager() Manager {
 	return &credentialsManager{
 		idToTaskCredentials: make(map[string]TaskIAMRoleCredentials),
+		knownCredentialsIDs: make(map[string]bool),
 	}
 }
 
@@ -132,6 +137,8 @@ func (manager *credentialsManager) SetTaskCredentials(taskCredentials *TaskIAMRo
 		IAMRoleCredentials: taskCredentials.GetIAMRoleCredentials(),
 	}
 
+	manager.knownCredentialsIDs[credentials.CredentialsID] = true
+
 	return nil
 }
 
@@ -157,4 +164,37 @@ func (manager *credentialsManager) RemoveCredentials(id string) {
 	defer manager.taskCredentialsLock.Unlock()
 
 	delete(manager.idToTaskCredentials, id)
+	delete(manager.knownCredentialsIDs, id)
+}
+
+// IsCredentialsPending returns true if credentials ID is known but has not yet arrived from ACS.
+func (manager *credentialsManager) IsCredentialsPending(id string) bool {
+	manager.taskCredentialsLock.RLock()
+	defer manager.taskCredentialsLock.RUnlock()
+
+	_, isKnown := manager.knownCredentialsIDs[id]
+	_, hasCredentials := manager.idToTaskCredentials[id]
+
+	result := isKnown && !hasCredentials
+	logger.Info("IsCredentialsPending", logger.Fields{
+		"id":             id,
+		"isKnown":        isKnown,
+		"hasCredentials": hasCredentials,
+		"result":         result,
+	})
+
+	return result
+}
+
+// AddKnownCredentialsID adds a credentials ID to the known set.
+// This is useful when agent needs to track known credentials IDs
+// for which the credentials themselves have not arrived from ACS.
+func (manager *credentialsManager) AddKnownCredentialsID(id string) {
+	manager.taskCredentialsLock.Lock()
+	defer manager.taskCredentialsLock.Unlock()
+
+	logger.Info("AddKnownCredentialsID: Adding credentials ID", logger.Fields{
+		"credentialsID": id,
+	})
+	manager.knownCredentialsIDs[id] = true
 }

agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/credentials/mocks/credentials_mocks.go

@@ -172,3 +172,54 @@ func TestRemoveExistingCredentials(t *testing.T) {
 		t.Error("Expected GetTaskCredentials to return false for removed credentials")
 	}
 }
+
+// TestAddKnownCredentialsID tests that AddKnownCredentialsID properly tracks credentials IDs
+func TestAddKnownCredentialsID(t *testing.T) {
+	manager := NewManager()
+	credentialsID := "test-creds-id"
+
+	// Initially, credentials should not be pending
+	assert.False(t, manager.IsCredentialsPending(credentialsID))
+
+	// Add known credentials ID
+	manager.AddKnownCredentialsID(credentialsID)
+
+	// Now it should be pending (known but no actual credentials)
+	assert.True(t, manager.IsCredentialsPending(credentialsID))
+
+	// Verify no actual credentials exist
+	_, ok := manager.GetTaskCredentials(credentialsID)
+	assert.False(t, ok)
+}
+
+// TestIsCredentialsPending tests the IsCredentialsPending method behavior
+func TestIsCredentialsPending(t *testing.T) {
+	manager := NewManager()
+	credentialsID := "test-creds-id"
+
+	// Case 1: Unknown credentials ID - should return false
+	assert.False(t, manager.IsCredentialsPending(credentialsID))
+
+	// Case 2: Known but no actual credentials - should return true
+	manager.AddKnownCredentialsID(credentialsID)
+	assert.True(t, manager.IsCredentialsPending(credentialsID))
+
+	// Case 3: Known and has actual credentials - should return false
+	credentials := TaskIAMRoleCredentials{
+		ARN: "arn:aws:ecs:us-east-1:123456789012:task/test-task",
+		IAMRoleCredentials: IAMRoleCredentials{
+			AccessKeyID:     "akid1",
+			SecretAccessKey: "skid1",
+			SessionToken:    "stkn",
+			Expiration:      "ts",
+			CredentialsID:   credentialsID,
+		},
+	}
+	err := manager.SetTaskCredentials(&credentials)
+	assert.NoError(t, err)
+	assert.False(t, manager.IsCredentialsPending(credentialsID))
+
+	// Case 4: After removal - should return false
+	manager.RemoveCredentials(credentialsID)
+	assert.False(t, manager.IsCredentialsPending(credentialsID))
+}