feat(ec2): access gateways created by NatProvider

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

@@ -6,6 +6,22 @@ import { Port } from './port';
 import { SecurityGroup } from './security-group';
 import { PrivateSubnet, PublicSubnet, RouterType, Vpc } from './vpc';
 
+/**
+ * Pair represents a gateway created by NAT Provider
+ */
+export interface GatewayConfig {
+
+  /**
+   * Availability Zone
+   */
+  readonly az: string
+
+  /**
+   * Identity of gateway spawned by the provider
+   */
+  readonly gatewayId: string
+}
+
 /**
  * NAT providers
  *
@@ -41,10 +57,20 @@ export abstract class NatProvider {
     return new NatInstance(props);
   }
 
+  /**
+   * Return list of gateways spawned by the provider
+   */
+  public abstract readonly configuredGateways: GatewayConfig[];
+
   /**
    * Called by the VPC to configure NAT
    */
   public abstract configureNat(options: ConfigureNatOptions): void;
+
+  /**
+   * Configures subnet with the gateway
+   */
+  public abstract configureSubnet(subnet: PrivateSubnet): void;
 }
 
 /**
@@ -111,33 +137,45 @@ export interface NatInstanceProps {
 }
 
 class NatGateway extends NatProvider {
+  private gateways: PrefSet<string> = new PrefSet<string>();
+
   public configureNat(options: ConfigureNatOptions) {
     // Create the NAT gateways
-    const gatewayIds = new PrefSet<string>();
     for (const sub of options.natSubnets) {
       const gateway = sub.addNatGateway();
-      gatewayIds.add(sub.availabilityZone, gateway.ref);
+      this.gateways.add(sub.availabilityZone, gateway.ref);
     }
 
     // Add routes to them in the private subnets
     for (const sub of options.privateSubnets) {
-      sub.addRoute('DefaultRoute', {
-        routerType: RouterType.NAT_GATEWAY,
-        routerId: gatewayIds.pick(sub.availabilityZone),
-        enablesInternetConnectivity: true,
-      });
+      this.configureSubnet(sub);
     }
   }
+
+  public configureSubnet(subnet: PrivateSubnet) {
+    const az = subnet.availabilityZone;
+    const gatewayId = this.gateways.pick(az);
+    subnet.addRoute('DefaultRoute', {
+      routerType: RouterType.NAT_GATEWAY,
+      routerId: gatewayId,
+      enablesInternetConnectivity: true,
+    });
+  }
+
+  public get configuredGateways(): GatewayConfig[] {
+    return this.gateways.values().map(x => ({az: x[0], gatewayId: x[1]}));
+  }
 }
 
 class NatInstance extends NatProvider {
+  private gateways: PrefSet<Instance> = new PrefSet<Instance>();
+
   constructor(private readonly props: NatInstanceProps) {
     super();
   }
 
   public configureNat(options: ConfigureNatOptions) {
     // Create the NAT instances. They can share a security group and a Role.
-    const instances = new PrefSet<Instance>();
     const machineImage = this.props.machineImage || new NatInstanceImage();
     const sg = new SecurityGroup(options.vpc, 'NatSecurityGroup', {
       vpc: options.vpc,
@@ -163,18 +201,28 @@ class NatInstance extends NatProvider {
         keyName: this.props.keyName
       });
       // NAT instance routes all traffic, both ways
-      instances.add(sub.availabilityZone, natInstance);
+      this.gateways.add(sub.availabilityZone, natInstance);
     }
 
     // Add routes to them in the private subnets
     for (const sub of options.privateSubnets) {
-      sub.addRoute('DefaultRoute', {
-        routerType: RouterType.INSTANCE,
-        routerId: instances.pick(sub.availabilityZone).instanceId,
-        enablesInternetConnectivity: true,
-      });
+      this.configureSubnet(sub);
     }
   }
+
+  public configureSubnet(subnet: PrivateSubnet) {
+    const az = subnet.availabilityZone;
+    const gatewayId = this.gateways.pick(az).instanceId;
+    subnet.addRoute('DefaultRoute', {
+      routerType: RouterType.INSTANCE,
+      routerId: gatewayId,
+      enablesInternetConnectivity: true,
+    });
+  }
+
+  public get configuredGateways(): GatewayConfig[] {
+    return this.gateways.values().map(x => ({az: x[0], gatewayId: x[1].instanceId}));
+  }
 }
 
 /**
@@ -185,12 +233,12 @@ class NatInstance extends NatProvider {
  */
 class PrefSet<A> {
   private readonly map: Record<string, A> = {};
-  private readonly vals = new Array<A>();
+  private readonly vals = new Array<[string, A]>();
   private next: number = 0;
 
   public add(pref: string, value: A) {
     this.map[pref] = value;
-    this.vals.push(value);
+    this.vals.push([pref, value]);
   }
 
   public pick(pref: string): A {
@@ -199,7 +247,11 @@ class PrefSet<A> {
     }
 
     if (pref in this.map) { return this.map[pref]; }
-    return this.vals[this.next++ % this.vals.length];
+    return this.vals[this.next++ % this.vals.length][1];
+  }
+
+  public values(): Array<[string, A]> {
+    return this.vals;
   }
 }
 
@@ -215,4 +267,4 @@ export class NatInstanceImage extends LookupMachineImage {
       owners: ['amazon'],
     });
   }
-}
+}

packages/@aws-cdk/aws-ec2/test/integ.nat-instances.lit.ts

@@ -7,17 +7,20 @@ class NatInstanceStack extends cdk.Stack {
 
     /// !show
     // Configure the `natGatewayProvider` when defining a Vpc
+    const natGatewayProvider = ec2.NatProvider.instance({
+      instanceType: new ec2.InstanceType('t3.small')
+    });
+
     const vpc = new ec2.Vpc(this, 'MyVpc', {
-      natGatewayProvider: ec2.NatProvider.instance({
-        instanceType: new ec2.InstanceType('t3.small')
-      }),
+      natGatewayProvider,
 
       // The 'natGateways' parameter now controls the number of NAT instances
       natGateways: 2,
     });
     /// !hide
 
     Array.isArray(vpc);
+    Array.isArray(natGatewayProvider.configuredGateways);
   }
 }
 
@@ -28,4 +31,4 @@ new NatInstanceStack(app, 'aws-cdk-vpc-nat-instances', {
     region: process.env.CDK_INTEG_REGION || process.env.CDK_DEFAULT_REGION,
   }
 });
-app.synth();
+app.synth();

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

@@ -618,6 +618,16 @@ export = {
 
       test.done();
     },
+
+    'Default NAT gateway provider'(test: Test) {
+      const stack = new Stack();
+      const natGatewayProvider = NatProvider.gateway();
+      new Vpc(stack, 'VpcNetwork', { natGatewayProvider });
+
+      test.ok(natGatewayProvider.configuredGateways.length > 0);
+
+      test.done();
+    }
   },
 
   'NAT instances': {
@@ -626,14 +636,13 @@ export = {
       const stack = getTestStack();
 
       // WHEN
-      new Vpc(stack, 'TheVPC', {
-        natGatewayProvider: NatProvider.instance({
-          instanceType: new InstanceType('q86.mega'),
-          machineImage: new GenericLinuxImage({
-            'us-east-1': 'ami-1'
-          })
+      const natGatewayProvider = NatProvider.instance({
+        instanceType: new InstanceType('q86.mega'),
+        machineImage: new GenericLinuxImage({
+          'us-east-1': 'ami-1'
         })
       });
+      new Vpc(stack, 'TheVPC', { natGatewayProvider });
 
       // THEN
       expect(stack).to(countResources('AWS::EC2::Instance', 3));