feat(vpc): allow Vpc.fromLookup() to discover asymmetric subnets (#4544)

Previously, Vpc.fromLookup() required every subnet group to be in the same Availability Zones,
and for all subnets in all groups to cover all of those Availability Zones.
This loosens this requirement, so that now Vpc.fromLookup() works for VPCs of all shapes.

This also add a way to customize which tag of the subnet is considered when grouping subnets into groups when calling fromLookup().

Fixes #3407

packages/@aws-cdk/aws-ec2/lib/vpc-lookup.ts

@@ -38,4 +38,14 @@ export interface VpcLookupOptions {
    * @default Don't care whether we return the default VPC
    */
   readonly isDefault?: boolean;
+
+  /**
+   * Optional tag for subnet group name.
+   * If not provided, we'll look at the aws-cdk:subnet-name tag.
+   * If the subnet does not have the specified tag,
+   * we'll use its type as the name.
+   *
+   * @default aws-cdk:subnet-name
+   */
+  readonly subnetGroupNameTag?: string;
 }

packages/@aws-cdk/aws-ec2/lib/vpc.ts

@@ -847,13 +847,17 @@ export class Vpc extends VpcBase {
       filter.isDefault = options.isDefault ? 'true' : 'false';
     }
 
-    const attributes = ContextProvider.getValue(scope, {
+    const attributes: cxapi.VpcContextResponse = ContextProvider.getValue(scope, {
       provider: cxapi.VPC_PROVIDER,
-      props: { filter } as cxapi.VpcContextQuery,
-      dummyValue: undefined
+      props: {
+        filter,
+        returnAsymmetricSubnets: true,
+        subnetGroupNameTag: options.subnetGroupNameTag,
+      } as cxapi.VpcContextQuery,
+      dummyValue: undefined,
     }).value;
 
-    return new ImportedVpc(scope, id, attributes || DUMMY_VPC_PROPS, attributes === undefined);
+    return new LookedUpVpc(scope, id, attributes || DUMMY_VPC_PROPS, attributes === undefined);
 
     /**
      * Prefixes all keys in the argument with `tag:`.`
@@ -1486,6 +1490,61 @@ class ImportedVpc extends VpcBase {
   }
 }
 
+class LookedUpVpc extends VpcBase {
+  public readonly vpcId: string;
+  public readonly vpnGatewayId?: string;
+  public readonly internetConnectivityEstablished: IDependable = new ConcreteDependable();
+  public readonly availabilityZones: string[];
+  public readonly publicSubnets: ISubnet[];
+  public readonly privateSubnets: ISubnet[];
+  public readonly isolatedSubnets: ISubnet[];
+
+  constructor(scope: Construct, id: string, props: cxapi.VpcContextResponse, isIncomplete: boolean) {
+    super(scope, id);
+
+    this.vpcId = props.vpcId;
+    this.vpnGatewayId = props.vpnGatewayId;
+    this.incompleteSubnetDefinition = isIncomplete;
+
+    const subnetGroups = props.subnetGroups || [];
+    const availabilityZones = Array.from(new Set<string>(flatMap(subnetGroups, subnetGroup => {
+      return subnetGroup.subnets.map(subnet => subnet.availabilityZone);
+    })));
+    availabilityZones.sort((az1, az2) => az1.localeCompare(az2));
+    this.availabilityZones = availabilityZones;
+
+    this.publicSubnets = this.extractSubnetsOfType(subnetGroups, cxapi.VpcSubnetGroupType.PUBLIC);
+    this.privateSubnets = this.extractSubnetsOfType(subnetGroups, cxapi.VpcSubnetGroupType.PRIVATE);
+    this.isolatedSubnets = this.extractSubnetsOfType(subnetGroups, cxapi.VpcSubnetGroupType.ISOLATED);
+  }
+
+  private extractSubnetsOfType(subnetGroups: cxapi.VpcSubnetGroup[], subnetGroupType: cxapi.VpcSubnetGroupType): ISubnet[] {
+    return flatMap(subnetGroups.filter(subnetGroup => subnetGroup.type === subnetGroupType),
+        subnetGroup => this.subnetGroupToSubnets(subnetGroup));
+  }
+
+  private subnetGroupToSubnets(subnetGroup: cxapi.VpcSubnetGroup): ISubnet[] {
+    const ret = new Array<ISubnet>();
+    for (let i = 0; i < subnetGroup.subnets.length; i++) {
+      const vpcSubnet = subnetGroup.subnets[i];
+      ret.push(Subnet.fromSubnetAttributes(this, `${subnetGroup.name}Subnet${i + 1}`, {
+        availabilityZone: vpcSubnet.availabilityZone,
+        subnetId: vpcSubnet.subnetId,
+        routeTableId: vpcSubnet.routeTableId,
+      }));
+    }
+    return ret;
+  }
+}
+
+function flatMap<T, U>(xs: T[], fn: (x: T) => U[]): U[] {
+  const ret = new Array<U>();
+  for (const x of xs) {
+    ret.push(...fn(x));
+  }
+  return ret;
+}
+
 class CompositeDependable implements IDependable {
   private readonly dependables = new Array<IDependable>();
 
@@ -1589,10 +1648,49 @@ function determineNatGatewayCount(requestedCount: number | undefined, subnetConf
  * It's only used for testing and on the first run-through.
  */
 const DUMMY_VPC_PROPS: cxapi.VpcContextResponse = {
-  availabilityZones: ['dummy-1a', 'dummy-1b'],
+  availabilityZones: [],
+  isolatedSubnetIds: undefined,
+  isolatedSubnetNames: undefined,
+  isolatedSubnetRouteTableIds: undefined,
+  privateSubnetIds: undefined,
+  privateSubnetNames: undefined,
+  privateSubnetRouteTableIds: undefined,
+  publicSubnetIds: undefined,
+  publicSubnetNames: undefined,
+  publicSubnetRouteTableIds: undefined,
+  subnetGroups: [
+    {
+      name: 'Public',
+      type: cxapi.VpcSubnetGroupType.PUBLIC,
+      subnets: [
+        {
+          availabilityZone: 'dummy-1a',
+          subnetId: 's-12345',
+          routeTableId: 'rtb-12345s',
+        },
+        {
+          availabilityZone: 'dummy-1b',
+          subnetId: 's-67890',
+          routeTableId: 'rtb-67890s',
+        },
+      ],
+    },
+    {
+      name: 'Private',
+      type: cxapi.VpcSubnetGroupType.PRIVATE,
+      subnets: [
+        {
+          availabilityZone: 'dummy-1a',
+          subnetId: 'p-12345',
+          routeTableId: 'rtb-12345p',
+        },
+        {
+          availabilityZone: 'dummy-1b',
+          subnetId: 'p-67890',
+          routeTableId: 'rtb-57890p',
+        },
+      ],
+    },
+  ],
   vpcId: 'vpc-12345',
-  publicSubnetIds: ['s-12345', 's-67890'],
-  publicSubnetRouteTableIds: ['rtb-12345s', 'rtb-67890s'],
-  privateSubnetIds: ['p-12345', 'p-67890'],
-  privateSubnetRouteTableIds: ['rtb-12345p', 'rtb-57890p'],
 };

packages/@aws-cdk/aws-ec2/test/test.vpc.from-lookup.ts

@@ -0,0 +1,148 @@
+import { Construct, ContextProvider, GetContextValueOptions, GetContextValueResult, Lazy, Stack } from "@aws-cdk/core";
+import cxapi = require('@aws-cdk/cx-api');
+import { Test } from 'nodeunit';
+import { Vpc } from "../lib";
+
+export = {
+  'Vpc.fromLookup()': {
+    'requires concrete values'(test: Test) {
+      // GIVEN
+      const stack = new Stack();
+
+      test.throws(() => {
+        Vpc.fromLookup(stack, 'Vpc', {
+          vpcId: Lazy.stringValue({ produce: () => 'some-id' })
+        });
+
+      }, 'All arguments to Vpc.fromLookup() must be concrete');
+
+      test.done();
+    },
+
+    'selecting subnets by name from a looked-up VPC does not throw'(test: Test) {
+      // GIVEN
+      const stack = new Stack(undefined, undefined, { env: { region: 'us-east-1', account: '123456789012' }});
+      const vpc = Vpc.fromLookup(stack, 'VPC', {
+        vpcId: 'vpc-1234'
+      });
+
+      // WHEN
+      vpc.selectSubnets({ subnetName: 'Bleep' });
+
+      // THEN: no exception
+
+      test.done();
+    },
+
+    'accepts asymmetric subnets'(test: Test) {
+      const previous = mockVpcContextProviderWith(test, {
+        vpcId: 'vpc-1234',
+        subnetGroups: [
+          {
+            name: 'Public',
+            type: cxapi.VpcSubnetGroupType.PUBLIC,
+            subnets: [
+              {
+                subnetId: 'pub-sub-in-us-east-1a',
+                availabilityZone: 'us-east-1a',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pub-sub-in-us-east-1b',
+                availabilityZone: 'us-east-1b',
+                routeTableId: 'rt-123',
+              },
+            ],
+          },
+          {
+            name: 'Private',
+            type: cxapi.VpcSubnetGroupType.PRIVATE,
+            subnets: [
+              {
+                subnetId: 'pri-sub-1-in-us-east-1c',
+                availabilityZone: 'us-east-1c',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pri-sub-2-in-us-east-1c',
+                availabilityZone: 'us-east-1c',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pri-sub-1-in-us-east-1d',
+                availabilityZone: 'us-east-1d',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pri-sub-2-in-us-east-1d',
+                availabilityZone: 'us-east-1d',
+                routeTableId: 'rt-123',
+              },
+            ],
+          },
+        ],
+      }, options => {
+        test.deepEqual(options.filter, {
+          isDefault: 'true',
+        });
+
+        test.equal(options.subnetGroupNameTag, undefined);
+      });
+
+      const stack = new Stack();
+      const vpc = Vpc.fromLookup(stack, 'Vpc', {
+        isDefault: true,
+      });
+
+      test.deepEqual(vpc.availabilityZones, ['us-east-1a', 'us-east-1b', 'us-east-1c', 'us-east-1d']);
+      test.equal(vpc.publicSubnets.length, 2);
+      test.equal(vpc.privateSubnets.length, 4);
+      test.equal(vpc.isolatedSubnets.length, 0);
+
+      restoreContextProvider(previous);
+      test.done();
+    },
+  },
+};
+
+interface MockVcpContextResponse {
+  readonly vpcId: string;
+  readonly subnetGroups: cxapi.VpcSubnetGroup[];
+}
+
+function mockVpcContextProviderWith(test: Test, response: MockVcpContextResponse,
+                                    paramValidator?: (options: cxapi.VpcContextQuery) => void) {
+  const previous = ContextProvider.getValue;
+  ContextProvider.getValue = (_scope: Construct, options: GetContextValueOptions) => {
+    // do some basic sanity checks
+    test.equal(options.provider, cxapi.VPC_PROVIDER,
+      `Expected provider to be: '${cxapi.VPC_PROVIDER}', got: '${options.provider}'`);
+    test.equal((options.props || {}).returnAsymmetricSubnets, true,
+      `Expected options.props.returnAsymmetricSubnets to be true, got: '${(options.props || {}).returnAsymmetricSubnets}'`);
+
+    if (paramValidator) {
+      paramValidator(options.props as any);
+    }
+
+    return {
+      value: {
+        availabilityZones: [],
+        isolatedSubnetIds: undefined,
+        isolatedSubnetNames: undefined,
+        isolatedSubnetRouteTableIds: undefined,
+        privateSubnetIds: undefined,
+        privateSubnetNames: undefined,
+        privateSubnetRouteTableIds: undefined,
+        publicSubnetIds: undefined,
+        publicSubnetNames: undefined,
+        publicSubnetRouteTableIds: undefined,
+        ...response,
+      } as cxapi.VpcContextResponse,
+    };
+  };
+  return previous;
+}
+
+function restoreContextProvider(previous: (scope: Construct, options: GetContextValueOptions) => GetContextValueResult): void {
+  ContextProvider.getValue = previous;
+}

packages/@aws-cdk/aws-ec2/test/test.vpc.ts

@@ -907,35 +907,6 @@ export = {
       test.done();
     }
   },
-
-  'fromLookup() requires concrete values'(test: Test) {
-    // GIVEN
-    const stack = new Stack();
-
-    test.throws(() => {
-      Vpc.fromLookup(stack, 'Vpc', {
-        vpcId: Lazy.stringValue({ produce: () => 'some-id' })
-      });
-
-    }, 'All arguments to Vpc.fromLookup() must be concrete');
-
-    test.done();
-  },
-
-  'selecting subnets by name from a looked-up VPC does not throw'(test: Test) {
-    // GIVEN
-    const stack = new Stack(undefined, undefined, { env: { region: 'us-east-1', account: '123456789012' }});
-    const vpc = Vpc.fromLookup(stack, 'VPC', {
-      vpcId: 'vpc-1234'
-    });
-
-    // WHEN
-    vpc.selectSubnets({ subnetName: 'Bleep' });
-
-    // THEN: no exception
-
-    test.done();
-  },
 };
 
 function getTestStack(): Stack {

packages/@aws-cdk/core/lib/context-provider.ts

@@ -137,6 +137,11 @@ function propsToArray(props: {[key: string]: any}, keyPrefix = ''): string[] {
   const ret: string[] = [];
 
   for (const key of Object.keys(props)) {
+    // skip undefined values
+    if (props[key] === undefined) {
+      continue;
+    }
+
     switch (typeof props[key]) {
       case 'object': {
         ret.push(...propsToArray(props[key], `${keyPrefix}${key}.`));

packages/@aws-cdk/core/test/test.context.ts

@@ -112,6 +112,33 @@ export = {
     test.done();
   },
 
+  'Keys with undefined values are not serialized'(test: Test) {
+    // GIVEN
+    const stack = new Stack(undefined, 'TestStack', { env: { account: '12345', region: 'us-east-1' } });
+
+    // WHEN
+    const result = ContextProvider.getKey(stack, {
+      provider: 'provider',
+      props: {
+        p1: 42,
+        p2: undefined,
+      },
+    });
+
+    // THEN
+    test.deepEqual(result, {
+      key: 'provider:account=12345:p1=42:region=us-east-1',
+      props: {
+        account: '12345',
+        region: 'us-east-1',
+        p1: 42,
+        p2: undefined,
+      },
+    });
+
+    test.done();
+  },
+
   'context provider errors are attached to tree'(test: Test) {
     const contextProps = { provider: 'bloop' };
     const contextKey = 'bloop:account=12345:region=us-east-1';  // Depends on the mangling algo