Skip to content

Commit e284b92

Browse files
ShelbyZShelbyZ
committed
inject AWS_REGION and AWS_DEFAULT_REGION into task containers
Inject the instance region into normal container environments so the AWS SDK can resolve the correct region without explicit task definition configuration. Injection is skipped if either var is already set in the task definition, environment files, or container image. The precedence order is: task definition > environment files > container image > instance default. Environment file vars are merged into container.Environment before the region injection runs.
1 parent a15ee8d commit e284b92

12 files changed

Lines changed: 441 additions & 61 deletions

‎agent/api/task/task.go‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,10 @@ const (
8282
// credentials.
8383
awsSDKCredentialsRelativeURIPathEnvironmentVariableName = "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"
8484

85+
// awsRegionEnvVar and awsDefaultRegionEnvVar are the standard AWS SDK region env vars.
86+
awsRegionEnvVar = "AWS_REGION"
87+
awsDefaultRegionEnvVar = "AWS_DEFAULT_REGION"
88+
8589
NvidiaVisibleDevicesEnvVar = "NVIDIA_VISIBLE_DEVICES"
8690
GPUAssociationType = "gpu"
8791

@@ -1024,6 +1028,38 @@ func (task *Task) initializeCredentialsEndpoint(credentialsManager credentials.M
10241028
task.SetCredentialsRelativeURI(credentialsEndpointRelativeURI)
10251029
}
10261030

1031+
// ApplyRegionToContainer injects AWS_REGION and AWS_DEFAULT_REGION into the
1032+
// container environment. Injection is skipped if either var is already set in
1033+
// the task definition, environment files, or container image.
1034+
func (task *Task) ApplyRegionToContainer(container *apicontainer.Container, region string, imageManagedEnvKeys map[string]bool) {
1035+
if region == "" {
1036+
return
1037+
}
1038+
if container.IsInternal() {
1039+
// Internal containers (pause, SC relay, managed daemons) do not receive injected env vars.
1040+
return
1041+
}
1042+
1043+
// Skip if the customer set either var in the task definition or environment files.
1044+
// Environment file vars are merged into container.Environment.
1045+
_, hasRegion := container.Environment[awsRegionEnvVar]
1046+
_, hasDefaultRegion := container.Environment[awsDefaultRegionEnvVar]
1047+
if hasRegion || hasDefaultRegion {
1048+
return
1049+
}
1050+
1051+
// Skip if the image already has a region preference (e.g. Dockerfile ENV).
1052+
if imageManagedEnvKeys[awsRegionEnvVar] || imageManagedEnvKeys[awsDefaultRegionEnvVar] {
1053+
return
1054+
}
1055+
1056+
if container.Environment == nil {
1057+
container.Environment = make(map[string]string)
1058+
}
1059+
container.Environment[awsRegionEnvVar] = region
1060+
container.Environment[awsDefaultRegionEnvVar] = region
1061+
}
1062+
10271063
// initializeContainersV3MetadataEndpoint generates a v3 endpoint id for each container, constructs the
10281064
// v3 metadata endpoint, and injects it as an environment variable
10291065
func (task *Task) initializeContainersV3MetadataEndpoint(uuidProvider utils.UUIDProvider) {

‎agent/api/task/task_test.go‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1406,6 +1406,104 @@ func TestInitializeContainersV1AgentAPIEndpoint(t *testing.T) {
14061406
}
14071407
}
14081408

1409+
func TestApplyRegionToContainerPrecedence(t *testing.T) {
1410+
tests := []struct {
1411+
name string
1412+
containerType apicontainer.ContainerType
1413+
taskDefEnvs map[string]string
1414+
imageEnvKeys map[string]bool
1415+
instanceRegion string
1416+
wantRegion string // expected AWS_REGION ("" = not injected)
1417+
wantDefaultReg string // expected AWS_DEFAULT_REGION ("" = not injected)
1418+
}{
1419+
{
1420+
name: "no overrides — inject both vars",
1421+
instanceRegion: "us-west-2",
1422+
wantRegion: "us-west-2",
1423+
wantDefaultReg: "us-west-2",
1424+
},
1425+
{
1426+
name: "nil image region keys — injection still proceeds",
1427+
instanceRegion: "us-west-2",
1428+
wantRegion: "us-west-2",
1429+
wantDefaultReg: "us-west-2",
1430+
},
1431+
{
1432+
name: "empty region — nothing injected",
1433+
instanceRegion: "",
1434+
wantRegion: "",
1435+
wantDefaultReg: "",
1436+
},
1437+
{
1438+
name: "internal container (CNI pause) — skip injection",
1439+
containerType: apicontainer.ContainerCNIPause,
1440+
instanceRegion: "us-west-2",
1441+
wantRegion: "",
1442+
wantDefaultReg: "",
1443+
},
1444+
{
1445+
// Image already has AWS_REGION: neither var is injected (both would be or neither).
1446+
name: "image has AWS_REGION — no injection",
1447+
imageEnvKeys: map[string]bool{awsRegionEnvVar: true},
1448+
instanceRegion: "us-west-2",
1449+
wantRegion: "",
1450+
wantDefaultReg: "",
1451+
},
1452+
{
1453+
// Image already has AWS_DEFAULT_REGION: neither var is injected.
1454+
name: "image has AWS_DEFAULT_REGION — no injection",
1455+
imageEnvKeys: map[string]bool{awsDefaultRegionEnvVar: true},
1456+
instanceRegion: "us-west-2",
1457+
wantRegion: "",
1458+
wantDefaultReg: "",
1459+
},
1460+
{
1461+
// Task def already has AWS_REGION: injection skipped, task def value preserved,
1462+
// AWS_DEFAULT_REGION is not added.
1463+
name: "task definition has AWS_REGION — no injection, task def value preserved",
1464+
taskDefEnvs: map[string]string{awsRegionEnvVar: "eu-west-1"},
1465+
instanceRegion: "us-west-2",
1466+
wantRegion: "eu-west-1",
1467+
wantDefaultReg: "",
1468+
},
1469+
{
1470+
// Task def already has AWS_DEFAULT_REGION: injection skipped, task def value
1471+
// preserved, AWS_REGION is not added.
1472+
name: "task definition has AWS_DEFAULT_REGION — no injection, task def value preserved",
1473+
taskDefEnvs: map[string]string{awsDefaultRegionEnvVar: "ap-southeast-1"},
1474+
instanceRegion: "us-west-2",
1475+
wantRegion: "",
1476+
wantDefaultReg: "ap-southeast-1",
1477+
},
1478+
{
1479+
// Task def check runs before image check: one var in task def suppresses
1480+
// injection entirely, so the other var from the image is also not applied.
1481+
name: "task definition has AWS_REGION, image has AWS_DEFAULT_REGION — no injection",
1482+
taskDefEnvs: map[string]string{awsRegionEnvVar: "eu-west-1"},
1483+
imageEnvKeys: map[string]bool{awsDefaultRegionEnvVar: true},
1484+
instanceRegion: "us-west-2",
1485+
wantRegion: "eu-west-1",
1486+
wantDefaultReg: "",
1487+
},
1488+
}
1489+
1490+
for _, tt := range tests {
1491+
t.Run(tt.name, func(t *testing.T) {
1492+
container := &apicontainer.Container{
1493+
Name: "app",
1494+
Type: tt.containerType,
1495+
Environment: tt.taskDefEnvs,
1496+
}
1497+
task := Task{Containers: []*apicontainer.Container{container}}
1498+
1499+
task.ApplyRegionToContainer(container, tt.instanceRegion, tt.imageEnvKeys)
1500+
1501+
assert.Equal(t, tt.wantRegion, container.Environment[awsRegionEnvVar])
1502+
assert.Equal(t, tt.wantDefaultReg, container.Environment[awsDefaultRegionEnvVar])
1503+
})
1504+
}
1505+
}
1506+
14091507
func TestPostUnmarshalTaskWithLocalVolumes(t *testing.T) {
14101508
// Constants used here are defined in task_unix_test.go and task_windows_test.go
14111509
taskFromACS := ecsacs.Task{

‎agent/engine/common_testutil.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ func validateContainerRunWorkflow(t *testing.T,
191191
client.EXPECT().
192192
TagImage(gomock.Any(), canonicalImageRef.String(), container.Image).Return(nil)
193193
imageManager.EXPECT().RecordContainerReference(container).Return(nil)
194-
imageManager.EXPECT().GetImageStateFromImageName(gomock.Any()).Return(nil, false)
194+
imageManager.EXPECT().GetImageStateFromImageName(gomock.Any()).Return(nil, false).AnyTimes()
195195
client.EXPECT().APIVersion().Return(defaultDockerClientAPIVersion, nil)
196196
dockerConfig, err := task.DockerConfig(container, defaultDockerClientAPIVersion)
197197
if err != nil {

‎agent/engine/docker_image_manager.go‎

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ func (imageManager *dockerImageManager) RecordContainerReference(container *apic
139139
// On agent restart, container ID was retrieved from agent state file
140140
// TODO add setter and getter for modifying this
141141
if container.ImageID != "" {
142-
if !imageManager.addContainerReferenceToExistingImageState(container) {
142+
if !imageManager.addContainerReferenceToExistingImageState(container, nil) {
143143
return fmt.Errorf("Failed to add container to existing image state")
144144
}
145145
return nil
@@ -164,9 +164,14 @@ func (imageManager *dockerImageManager) RecordContainerReference(container *apic
164164
imageDigest := imageManager.fetchRepoDigest(imageInspected, container)
165165
container.SetImageDigest(imageDigest)
166166
}
167-
added := imageManager.addContainerReferenceToExistingImageState(container)
167+
// Capture managed image env keys for use at container create time.
168+
var managedEnvKeys map[string]bool
169+
if imageInspected.Config != nil {
170+
managedEnvKeys = image.ParseManagedEnvKeys(imageInspected.Config.Env)
171+
}
172+
added := imageManager.addContainerReferenceToExistingImageState(container, managedEnvKeys)
168173
if !added {
169-
imageManager.addContainerReferenceToNewImageState(container, imageInspected.Size)
174+
imageManager.addContainerReferenceToNewImageState(container, imageInspected.Size, managedEnvKeys)
170175
}
171176
return nil
172177
}
@@ -191,20 +196,24 @@ func (imageManager *dockerImageManager) fetchRepoDigest(imageInspected *types.Im
191196
return resultRepoDigest
192197
}
193198

194-
func (imageManager *dockerImageManager) addContainerReferenceToExistingImageState(container *apicontainer.Container) bool {
199+
func (imageManager *dockerImageManager) addContainerReferenceToExistingImageState(container *apicontainer.Container, managedEnvKeys map[string]bool) bool {
195200
// this lock is used for reading the image states in the image manager
196201
imageManager.updateLock.RLock()
197202
defer imageManager.updateLock.RUnlock()
198203
imageManager.removeExistingImageNameOfDifferentID(container.Image, container.ImageID)
199204
imageState, ok := imageManager.getImageState(container.ImageID)
200205
if ok {
201206
imageState.UpdateImageState(container)
207+
if managedEnvKeys != nil {
208+
// nil on restart (no inspect); preserves persisted value.
209+
imageState.SetManagedEnvKeys(managedEnvKeys)
210+
}
202211
imageManager.saveImageStateData(imageState)
203212
}
204213
return ok
205214
}
206215

207-
func (imageManager *dockerImageManager) addContainerReferenceToNewImageState(container *apicontainer.Container, imageSize int64) {
216+
func (imageManager *dockerImageManager) addContainerReferenceToNewImageState(container *apicontainer.Container, imageSize int64, managedEnvKeys map[string]bool) {
208217
// this lock is used while creating and adding new image state to image manager
209218
imageManager.updateLock.Lock()
210219
defer imageManager.updateLock.Unlock()
@@ -213,16 +222,20 @@ func (imageManager *dockerImageManager) addContainerReferenceToNewImageState(con
213222
imageState, ok := imageManager.getImageState(container.ImageID)
214223
if ok {
215224
imageState.UpdateImageState(container)
225+
if managedEnvKeys != nil {
226+
imageState.SetManagedEnvKeys(managedEnvKeys)
227+
}
216228
imageManager.saveImageStateData(imageState)
217229
} else {
218230
sourceImage := &image.Image{
219231
ImageID: container.ImageID,
220232
Size: imageSize,
221233
}
222234
sourceImageState := &image.ImageState{
223-
Image: sourceImage,
224-
PulledAt: time.Now(),
225-
LastUsedAt: time.Now(),
235+
Image: sourceImage,
236+
PulledAt: time.Now(),
237+
LastUsedAt: time.Now(),
238+
ManagedEnvKeys: managedEnvKeys,
226239
}
227240
sourceImageState.UpdateImageState(container)
228241
imageManager.addImageState(sourceImageState)

‎agent/engine/docker_image_manager_data_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ func TestAddContainerReferenceToNewImageStateSaveData(t *testing.T) {
6767
dataClient: dataClient,
6868
}
6969

70-
imageManager.addContainerReferenceToNewImageState(testContainerData, 0)
70+
imageManager.addContainerReferenceToNewImageState(testContainerData, 0, nil)
7171
imageStates, err := dataClient.GetImageStates()
7272
assert.NoError(t, err)
7373
assert.Len(t, imageStates, 1)
@@ -80,7 +80,7 @@ func TestAddContainerReferenceToExistingImageStateSaveData(t *testing.T) {
8080
dataClient: dataClient,
8181
}
8282
imageManager.imageStates = append(imageManager.imageStates, testImageStateData)
83-
imageManager.addContainerReferenceToExistingImageState(testContainerData)
83+
imageManager.addContainerReferenceToExistingImageState(testContainerData, nil)
8484
imageStates, err := dataClient.GetImageStates()
8585
assert.NoError(t, err)
8686
assert.Len(t, imageStates, 1)

‎agent/engine/docker_image_manager_test.go‎

Lines changed: 89 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ import (
3535
ec2testutil "github.com/aws/amazon-ecs-agent/agent/utils/test/ec2util"
3636

3737
"github.com/docker/docker/api/types"
38+
dockercontainer "github.com/docker/docker/api/types/container"
3839
"github.com/golang/mock/gomock"
3940
"github.com/stretchr/testify/assert"
4041
"github.com/stretchr/testify/require"
@@ -192,6 +193,90 @@ func TestRecordContainerReferenceInspectError(t *testing.T) {
192193
}
193194
}
194195

196+
func TestRecordContainerReferenceStoresManagedEnvKeysOnImageState(t *testing.T) {
197+
envVars := []string{"PATH=/usr/local/bin", "AWS_DEFAULT_REGION=us-east-1"}
198+
tests := []struct {
199+
name string
200+
config *dockercontainer.Config
201+
wantEnvKeys map[string]bool
202+
}{
203+
{
204+
name: "image config with region env var — key stored on image state",
205+
config: &dockercontainer.Config{Env: envVars},
206+
wantEnvKeys: map[string]bool{"AWS_DEFAULT_REGION": true},
207+
},
208+
{
209+
name: "nil image config — ManagedEnvKeys nil on image state",
210+
config: nil,
211+
wantEnvKeys: nil,
212+
},
213+
}
214+
215+
for _, tt := range tests {
216+
t.Run(tt.name, func(t *testing.T) {
217+
ctrl := gomock.NewController(t)
218+
defer ctrl.Finish()
219+
client := mock_dockerapi.NewMockDockerClient(ctrl)
220+
221+
imageManager := NewImageManager(defaultTestConfig(), client, dockerstate.NewTaskEngineState())
222+
imageManager.SetDataClient(data.NewNoopClient())
223+
224+
container := &apicontainer.Container{
225+
Name: "testContainer",
226+
Image: "testContainerImage",
227+
}
228+
client.EXPECT().InspectImage(container.Image).Return(&types.ImageInspect{
229+
ID: "sha256:qwerty",
230+
Config: tt.config,
231+
}, nil)
232+
233+
err := imageManager.RecordContainerReference(container)
234+
assert.NoError(t, err)
235+
236+
imageState, ok := imageManager.GetImageStateFromImageName(container.Image)
237+
assert.True(t, ok)
238+
assert.Equal(t, tt.wantEnvKeys, imageState.GetManagedEnvKeys())
239+
})
240+
}
241+
}
242+
243+
// TestRecordContainerReferencePreservesManagedEnvKeysOnRestart verifies that when a container
244+
// already has an ImageID set (e.g. on agent restart from saved state), RecordContainerReference
245+
// does not re-inspect the image — the ManagedEnvKeys are already persisted on ImageState.
246+
func TestRecordContainerReferencePreservesManagedEnvKeysOnRestart(t *testing.T) {
247+
ctrl := gomock.NewController(t)
248+
defer ctrl.Finish()
249+
client := mock_dockerapi.NewMockDockerClient(ctrl)
250+
251+
imageManager := NewImageManager(defaultTestConfig(), client, dockerstate.NewTaskEngineState())
252+
imageManager.SetDataClient(data.NewNoopClient())
253+
254+
const imageID = "sha256:qwerty"
255+
existingKeys := map[string]bool{"AWS_DEFAULT_REGION": true}
256+
257+
// Seed the image state with persisted ManagedEnvKeys (as if loaded from state file).
258+
sourceImageState := &image.ImageState{
259+
Image: &image.Image{ImageID: imageID},
260+
ManagedEnvKeys: existingKeys,
261+
}
262+
imageManager.(*dockerImageManager).addImageState(sourceImageState)
263+
264+
container := &apicontainer.Container{
265+
Name: "testContainer",
266+
Image: "testContainerImage",
267+
ImageID: imageID, // pre-populated, as it would be after agent restart
268+
}
269+
270+
// No InspectImage call expected — keys are already persisted.
271+
err := imageManager.RecordContainerReference(container)
272+
assert.NoError(t, err)
273+
274+
// Verify the keys are still on the image state.
275+
imageState, ok := imageManager.GetImageStateFromImageName(container.Image)
276+
assert.True(t, ok)
277+
assert.Equal(t, existingKeys, imageState.GetManagedEnvKeys())
278+
}
279+
195280
func TestRecordContainerReferenceWithNoImageName(t *testing.T) {
196281
ctrl := gomock.NewController(t)
197282
defer ctrl.Finish()
@@ -282,7 +367,7 @@ func TestAddContainerReferenceToExistingImageState(t *testing.T) {
282367
sourceImageState1.AddImageName("testContainerImage")
283368
imageManager.addImageState(sourceImageState)
284369
imageManager.addImageState(sourceImageState1)
285-
if !imageManager.addContainerReferenceToExistingImageState(container) {
370+
if !imageManager.addContainerReferenceToExistingImageState(container, nil) {
286371
t.Error("Error in adding container to an already existing image state")
287372
}
288373
if !reflect.DeepEqual(sourceImageState.Containers[0], container) {
@@ -371,7 +456,7 @@ func TestAddContainerReferenceToExistingImageStateNoState(t *testing.T) {
371456
Image: "testContainerImage",
372457
ImageID: "sha256:qwerty",
373458
}
374-
if imageManager.addContainerReferenceToExistingImageState(container) {
459+
if imageManager.addContainerReferenceToExistingImageState(container, nil) {
375460
t.Error("Error adding container to an incorrect existing image state")
376461
}
377462
}
@@ -390,7 +475,7 @@ func TestAddContainerReferenceToNewImageState(t *testing.T) {
390475
Image: "testContainerImage",
391476
ImageID: imageID,
392477
}
393-
imageManager.addContainerReferenceToNewImageState(container, imageSize)
478+
imageManager.addContainerReferenceToNewImageState(container, imageSize, nil)
394479
_, ok := imageManager.getImageState(imageID)
395480
if !ok {
396481
t.Error("Error adding container reference to new image state")
@@ -426,7 +511,7 @@ func TestAddContainerReferenceToNewImageStateAddedState(t *testing.T) {
426511
sourceImageState1.AddImageName("testContainerImage")
427512
imageManager.addImageState(sourceImageState)
428513
imageManager.addImageState(sourceImageState1)
429-
imageManager.addContainerReferenceToNewImageState(container, imageSize)
514+
imageManager.addContainerReferenceToNewImageState(container, imageSize, nil)
430515
if !reflect.DeepEqual(sourceImageState.Containers[0], container) {
431516
t.Error("Incorrect container added to an already existing image state")
432517
}

0 commit comments

Comments
 (0)