agent/app: register task-eni-block-instance-metadata capability when …

…AWSVPCBlockInstanceMetdata config is enabled
aws · Sep 7, 2017 · 654e8bf · 654e8bf
1 parent 63734d3
commit 654e8bf
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 16 deletions.
diff --git a/agent/app/agent_capability.go b/agent/app/agent_capability.go
@@ -23,12 +23,13 @@ import (
 )
 
 const (
-	capabilityPrefix             = "com.amazonaws.ecs.capability."
-	capabilityTaskIAMRole        = "task-iam-role"
-	capabilityTaskIAMRoleNetHost = "task-iam-role-network-host"
-	attributePrefix              = "ecs.capability."
-	taskENIAttributeSuffix       = "task-eni"
-	cniPluginVersionSuffix       = "cni-plugin-version"
+	capabilityPrefix                            = "com.amazonaws.ecs.capability."
+	capabilityTaskIAMRole                       = "task-iam-role"
+	capabilityTaskIAMRoleNetHost                = "task-iam-role-network-host"
+	attributePrefix                             = "ecs.capability."
+	taskENIAttributeSuffix                      = "task-eni"
+	taskENIBlockInstanceMetadataAttributeSuffix = "task-eni-block-instance-metadata"
+	cniPluginVersionSuffix                      = "cni-plugin-version"
 )
 
 // capabilities returns the supported capabilities of this agent / docker-client pair.
@@ -49,7 +50,8 @@ const (
 //    com.amazonaws.ecs.capability.ecr-auth
 //    com.amazonaws.ecs.capability.task-iam-role
 //    com.amazonaws.ecs.capability.task-iam-role-network-host
-//    ecs.capability.task-eni.0.1.0
+//    ecs.capability.task-eni
+//    ecs.capability.task-eni-block-instance-metadata
 func (agent *ecsAgent) capabilities() []*ecs.Attribute {
 	var capabilities []*ecs.Attribute
 
@@ -124,6 +126,14 @@ func (agent *ecsAgent) capabilities() []*ecs.Attribute {
 		capabilities = append(capabilities, taskENIVersionAttribute)
 	}
 
+	if agent.cfg.AWSVPCBlockInstanceMetdata {
+		// If the Block Instance Metadata flag is set for AWS VPC networking mode, register a capability
+		// indicating the same
+		capabilities = append(capabilities, &ecs.Attribute{
+			Name: aws.String(attributePrefix + taskENIBlockInstanceMetadataAttributeSuffix),
+		})
+	}
+
 	return capabilities
 }
 

diff --git a/agent/app/agent_capability_test.go b/agent/app/agent_capability_test.go
@@ -47,11 +47,12 @@ func TestCapabilities(t *testing.T) {
 			dockerclient.GelfDriver,
 			dockerclient.FluentdDriver,
 		},
-		PrivilegedDisabled:      false,
-		SELinuxCapable:          true,
-		AppArmorCapable:         true,
-		TaskENIEnabled:          true,
-		TaskCleanupWaitDuration: config.DefaultConfig().TaskCleanupWaitDuration,
+		PrivilegedDisabled:         false,
+		SELinuxCapable:             true,
+		AppArmorCapable:            true,
+		TaskENIEnabled:             true,
+		AWSVPCBlockInstanceMetdata: true,
+		TaskCleanupWaitDuration:    config.DefaultConfig().TaskCleanupWaitDuration,
 	}
 
 	gomock.InOrder(
@@ -85,10 +86,15 @@ func TestCapabilities(t *testing.T) {
 			&ecs.Attribute{Name: aws.String(name)})
 	}
 	expectedCapabilities = append(expectedCapabilities,
-		&ecs.Attribute{
-			Name:  aws.String(attributePrefix + cniPluginVersionSuffix),
-			Value: aws.String("v1"),
-		})
+		[]*ecs.Attribute{
+			{
+				Name:  aws.String(attributePrefix + cniPluginVersionSuffix),
+				Value: aws.String("v1"),
+			},
+			{
+				Name: aws.String(attributePrefix + taskENIBlockInstanceMetadataAttributeSuffix),
+			},
+		}...)
 
 	ctx, cancel := context.WithCancel(context.TODO())
 	// Cancel the context to cancel async routines