feat(ecs): allow load balancers to connect to any mapped port

packages/@aws-cdk/aws-ecs/lib/base/base-service.ts

@@ -7,6 +7,7 @@ import cloudmap = require('@aws-cdk/aws-servicediscovery');
 import { Construct, Duration, IResolvable, IResource, Lazy, Resource, Stack } from '@aws-cdk/core';
 import { NetworkMode, TaskDefinition } from '../base/task-definition';
 import { ICluster } from '../cluster';
+import { Protocol } from '../container-definition';
 import { CfnService } from '../ecs.generated';
 import { ScalableTaskCount } from './scalable-task-count';
 
@@ -225,7 +226,7 @@ export abstract class BaseService extends Resource
 
     // Open up security groups. For dynamic port mapping, we won't know the port range
     // in advance so we need to open up all ports.
-    const port = this.taskDefinition.defaultContainer!.ingressPort;
+    const port = targetGroup.targetPort || this.taskDefinition.defaultContainer!.ingressPort;
     const portRange = port === 0 ? EPHEMERAL_PORT_RANGE : ec2.Port.tcp(port);
     targetGroup.registerConnectable(this, portRange);
 
@@ -327,11 +328,7 @@ export abstract class BaseService extends Resource
       throw new Error("Cannot use a load balancer if NetworkMode is None. Use Bridge, Host or AwsVpc instead.");
     }
 
-    this.loadBalancers.push({
-      targetGroupArn: targetGroup.targetGroupArn,
-      containerName: this.taskDefinition.defaultContainer!.containerName,
-      containerPort: this.taskDefinition.defaultContainer!.containerPort,
-    });
+    this.loadBalancers.push(this.makeLoadBalancerProperty(targetGroup.targetGroupArn, targetGroup.targetPort));
 
     // Service creation can only happen after the load balancer has
     // been associated with our target group(s), so add ordering dependency.
@@ -341,6 +338,31 @@ export abstract class BaseService extends Resource
     return { targetType };
   }
 
+  /**
+   * Generate a load balancer property given an optional host port
+   */
+  private makeLoadBalancerProperty(targetGroupArn: string, targetPort?: number): CfnService.LoadBalancerProperty {
+    if (targetPort === undefined) {
+      return {
+        targetGroupArn,
+        containerPort: this.taskDefinition.defaultContainer!.containerPort,
+        containerName: this.taskDefinition.defaultContainer!.containerName,
+      };
+    }
+
+    const container = this.taskDefinition._findContainerByHostPort(targetPort, Protocol.TCP);
+    if (container === undefined) {
+      throw new Error("Cannot attach a load balancer to an unmapped container port.");
+    }
+
+    const portMapping = container._findPortMapping(targetPort, Protocol.TCP);
+    return {
+      targetGroupArn,
+      containerPort: portMapping!.containerPort,
+      containerName: container.containerName,
+    };
+  }
+
   /**
    * Generate the role that will be used for autoscaling this service
    */

packages/@aws-cdk/aws-ecs/lib/base/task-definition.ts

@@ -1,6 +1,6 @@
 import iam = require('@aws-cdk/aws-iam');
 import { Construct, IResource, Lazy, Resource } from '@aws-cdk/core';
-import { ContainerDefinition, ContainerDefinitionOptions } from '../container-definition';
+import { ContainerDefinition, ContainerDefinitionOptions, Protocol } from '../container-definition';
 import { CfnTaskDefinition } from '../ecs.generated';
 import { PlacementConstraint } from '../placement';
 
@@ -325,6 +325,19 @@ export class TaskDefinition extends TaskDefinitionBase {
     }
   }
 
+  /**
+   * Returns the container that listens to a given host port if present.
+   * @internal
+   */
+  public _findContainerByHostPort(port: number, protocol: Protocol): ContainerDefinition | undefined {
+    for (const container of this.containers) {
+      if (container._findPortMapping(port, protocol) !== undefined) {
+        return container;
+      }
+    }
+    return undefined;
+  }
+
   /**
    * Adds a volume to the task definition.
    */

packages/@aws-cdk/aws-ecs/lib/container-definition.ts

@@ -479,6 +479,21 @@ export class ContainerDefinition extends cdk.Construct {
     return defaultPortMapping.containerPort;
   }
 
+  /**
+   * Returns the container port for the requested host port if it exists
+   * @internal
+   */
+  public _findPortMapping(hostPort: number, protocol: Protocol): PortMapping | undefined {
+    for (const portMapping of this.portMappings) {
+      const p = portMapping.protocol || Protocol.TCP;
+      const h = portMapping.hostPort || portMapping.containerPort;
+      if (protocol === p && hostPort === h) {
+        return portMapping;
+      }
+    }
+    return undefined;
+  }
+
   /**
    * Render this container definition to a CloudFormation object
    *

packages/@aws-cdk/aws-ecs/test/ec2/test.ec2-service.ts

@@ -1,4 +1,4 @@
-import { expect, haveResource } from '@aws-cdk/assert';
+import { expect, haveResource, haveResourceLike } from '@aws-cdk/assert';
 import ec2 = require('@aws-cdk/aws-ec2');
 import elb = require('@aws-cdk/aws-elasticloadbalancing');
 import elbv2 = require("@aws-cdk/aws-elasticloadbalancingv2");
@@ -1061,7 +1061,78 @@ export = {
       test.throws(() => {
         service.attachToApplicationTargetGroup(targetGroup);
       }, /Cannot use a load balancer if NetworkMode is None. Use Bridge, Host or AwsVpc instead./);
+      test.done();
+    },
+
+    "allows connection to an explicit target port"(test: Test) {
+      // GIVEN
+      const stack = new cdk.Stack();
+      const vpc = new ec2.Vpc(stack, 'MyVpc', {});
+      const cluster = new ecs.Cluster(stack, 'EcsCluster', { vpc });
+      cluster.addCapacity('EcsCapacity', { instanceType: new ec2.InstanceType('t2.micro')});
+      const taskDefinition = new ecs.Ec2TaskDefinition(stack, 'Ec2TaskDef', { networkMode: ecs.NetworkMode.AWS_VPC });
+      const container = taskDefinition.addContainer('MainContainer', {
+        image: ContainerImage.fromRegistry('hello'),
+        memoryLimitMiB: 2048,
+      });
+      container.addPortMappings({ containerPort: 8000 }, { containerPort: 8080 });
+
+      const service = new ecs.Ec2Service(stack, 'Service', {
+        cluster,
+        taskDefinition
+      });
+
+      const lb = new elbv2.ApplicationLoadBalancer(stack, "lb", { vpc });
+      const listener = lb.addListener("listener", { port: 80 });
+      const targetGroup = listener.addTargets("target", {
+        port: 80,
+        targetPort: 8080,
+      });
+
+      // THEN
+      service.attachToApplicationTargetGroup(targetGroup);
+      test.equal(container.ingressPort, 8000);
+      expect(stack).to(haveResourceLike("AWS::ECS::Service", {
+        LoadBalancers: [
+          {
+            ContainerName: "MainContainer",
+            ContainerPort: 8080,
+          }
+        ]
+      }));
+      expect(stack).to(haveResourceLike("AWS::EC2::SecurityGroupIngress", {
+        ToPort: 8080
+      }));
+      test.done();
+    },
+
+    "throws if connection specifies a target port that doesn't exist"(test: Test) {
+      // GIVEN
+      const stack = new cdk.Stack();
+      const vpc = new ec2.Vpc(stack, 'MyVpc', {});
+      const cluster = new ecs.Cluster(stack, 'EcsCluster', { vpc });
+      const taskDefinition = new ecs.Ec2TaskDefinition(stack, 'Ec2TaskDef', { networkMode: ecs.NetworkMode.AWS_VPC });
+      const container = taskDefinition.addContainer('MainContainer', {
+        image: ContainerImage.fromRegistry('hello'),
+      });
+      container.addPortMappings({ containerPort: 8000 });
 
+      const service = new ecs.Ec2Service(stack, 'Service', {
+        cluster,
+        taskDefinition
+      });
+
+      const lb = new elbv2.ApplicationLoadBalancer(stack, "lb", { vpc });
+      const listener = lb.addListener("listener", { port: 80 });
+      const targetGroup = listener.addTargets("target", {
+        port: 80,
+        targetPort: 8080,
+      });
+
+      // THEN
+      test.throws(() => {
+        service.attachToApplicationTargetGroup(targetGroup);
+      }, /Cannot attach a load balancer to an unmapped container port./);
       test.done();
     }
   },
@@ -1122,6 +1193,75 @@ export = {
         service.attachToNetworkTargetGroup(targetGroup);
       }, /Cannot use a load balancer if NetworkMode is None. Use Bridge, Host or AwsVpc instead./);
 
+      test.done();
+    },
+
+    "allows connection to an explicit target port"(test: Test) {
+      // GIVEN
+      const stack = new cdk.Stack();
+      const vpc = new ec2.Vpc(stack, 'MyVpc', {});
+      const cluster = new ecs.Cluster(stack, 'EcsCluster', { vpc });
+      cluster.addCapacity('EcsCapacity', { instanceType: new ec2.InstanceType('t2.micro') });
+      const taskDefinition = new ecs.Ec2TaskDefinition(stack, 'Ec2TaskDef', { networkMode: ecs.NetworkMode.AWS_VPC });
+      const container = taskDefinition.addContainer('MainContainer', {
+        image: ContainerImage.fromRegistry('hello'),
+        memoryLimitMiB: 2048,
+      });
+      container.addPortMappings({ containerPort: 8000 }, { containerPort: 8080 });
+
+      const service = new ecs.Ec2Service(stack, 'Service', {
+        cluster,
+        taskDefinition
+      });
+
+      const lb = new elbv2.NetworkLoadBalancer(stack, "lb", { vpc });
+      const listener = lb.addListener("listener", { port: 80 });
+      const targetGroup = listener.addTargets("target", {
+        port: 80,
+        targetPort: 8080,
+      });
+
+      // THEN
+      service.attachToNetworkTargetGroup(targetGroup);
+      test.equal(container.ingressPort, 8000);
+      expect(stack).to(haveResourceLike("AWS::ECS::Service", {
+        LoadBalancers: [
+          {
+            ContainerName: "MainContainer",
+            ContainerPort: 8080,
+          }
+        ]
+      }));
+      test.done();
+    },
+
+    "throws if connection specifies a target port that doesn't exist"(test: Test) {
+      // GIVEN
+      const stack = new cdk.Stack();
+      const vpc = new ec2.Vpc(stack, 'MyVpc', {});
+      const cluster = new ecs.Cluster(stack, 'EcsCluster', { vpc });
+      const taskDefinition = new ecs.Ec2TaskDefinition(stack, 'Ec2TaskDef', { networkMode: ecs.NetworkMode.AWS_VPC });
+      const container = taskDefinition.addContainer('MainContainer', {
+        image: ContainerImage.fromRegistry('hello'),
+      });
+      container.addPortMappings({ containerPort: 8000 });
+
+      const service = new ecs.Ec2Service(stack, 'Service', {
+        cluster,
+        taskDefinition
+      });
+
+      const lb = new elbv2.NetworkLoadBalancer(stack, "lb", { vpc });
+      const listener = lb.addListener("listener", { port: 80 });
+      const targetGroup = listener.addTargets("target", {
+        port: 80,
+        targetPort: 8080,
+      });
+
+      // THEN
+      test.throws(() => {
+        service.attachToNetworkTargetGroup(targetGroup);
+      }, /Cannot attach a load balancer to an unmapped container port./);
       test.done();
     }
   },

packages/@aws-cdk/aws-ecs/test/ec2/test.ec2-task-definition.ts

@@ -994,6 +994,59 @@ export = {
 
       test.done();
     },
+  },
+  "_findContainerByHostPort works correctly": {
+    'when no containers are present'(test: Test) {
+      // GIVEN
+      const stack = new cdk.Stack();
+
+      // WHEN
+      const taskDefinition = new ecs.Ec2TaskDefinition(stack, 'Ec2TaskDef');
+
+      // THEN
+      test.equal(taskDefinition._findContainerByHostPort(8000, ecs.Protocol.TCP), undefined);
+      test.done();
+    },
+
+    'when no matching container port is available'(test: Test) {
+      // GIVEN
+      const stack = new cdk.Stack();
+
+      // WHEN
+      const taskDefinition = new ecs.Ec2TaskDefinition(stack, 'Ec2TaskDef');
+      const container = taskDefinition.addContainer("web", {
+        image: ecs.ContainerImage.fromRegistry("amazon/amazon-ecs-sample"),
+        memoryLimitMiB: 512
+      });
+      container.addPortMappings({ containerPort: 8080 });
+
+      // THEN
+      test.equal(taskDefinition._findContainerByHostPort(8000, ecs.Protocol.TCP), undefined);
+      test.done();
+    },
+
+    'when a matching container port exists'(test: Test) {
+      // GIVEN
+      const stack = new cdk.Stack();
+
+      // WHEN
+      const taskDefinition = new ecs.Ec2TaskDefinition(stack, 'Ec2TaskDef');
+      const primary = taskDefinition.addContainer("web", {
+        image: ecs.ContainerImage.fromRegistry("amazon/amazon-ecs-sample"),
+        memoryLimitMiB: 512
+      });
+      primary.addPortMappings({ containerPort: 8080 });
+      const secondary = taskDefinition.addContainer("sidecar", {
+        image: ecs.ContainerImage.fromRegistry("sidecar"),
+        memoryLimitMiB: 512
+      });
+      secondary.addPortMappings({ containerPort: 8000 });
 
+      // THEN
+      test.equal(taskDefinition.defaultContainer, primary);
+      test.equal(taskDefinition._findContainerByHostPort(8000, ecs.Protocol.TCP), secondary);
+      test.equal(taskDefinition._findContainerByHostPort(8080, ecs.Protocol.TCP), primary);
+      test.done();
+    }
   }
 };