Repository: pulumi/pulumi-awsx (ts)

Issue:
Validation when using multiple subnets of the same type
```
const vpc = new awsx.ec2.Vpc("custom", {
  subnetSpecs:[{
    type: awsx.ec2.SubnetType.Public,
    cidrMask: 22,
  }, {
    type: awsx.ec2.SubnetType.Private,
    cidrMask: 21,
  },
  {
    type: awsx.ec2.SubnetType.Isolated,
    cidrMask: 24,
    name: "db"
  },
  {
    type: awsx.ec2.SubnetType.Isolated,
    cidrMask: 24,
    name: "redis"
  }],
});
```

produces

```
    error: Error: The following subnets overlap with at least one other subnet. Make the CIDR for the VPC larger, reduce the size of the subnets per AZ, or use less Availability Zones:

    1. custom-db-1: 10.0.12.0/24
    2. custom-redis-1: 10.0.12.0/24
    3. custom-db-2: 10.0.76.0/24
    4. custom-redis-2: 10.0.76.0/24
    5. custom-db-3: 10.0.140.0/24
    6. custom-redis-3: 10.0.140.0/24
```

Pull request description:
Fix up CIDR block calculation for more advanced subnet configs.
Fixes #810

Patch:
diff --git a/awsx/ec2/subnetDistributor.ts b/awsx/ec2/subnetDistributor.ts
index 1e4e770a..65fc65de 100644
--- a/awsx/ec2/subnetDistributor.ts
+++ b/awsx/ec2/subnetDistributor.ts
@@ -44,105 +44,109 @@ export function getSubnetSpecs(
     const ipv4 = new ipAddress.Address4(azBases[0]);
     const baseSubnetMask = ipv4.subnetMask;
 
-    const privateSubnets: SubnetSpec[] = [];
-    const publicSubnets: SubnetSpec[] = [];
-    const isolatedSubnets: SubnetSpec[] = [];
+    const privateSubnetsIn = subnetInputs.filter((x) => x.type === "Private");
+    const publicSubnetsIn = subnetInputs.filter((x) => x.type === "Public");
+    const isolatedSubnetsIn = subnetInputs.filter((x) => x.type === "Isolated");
+
+    let subnetsOut: SubnetSpec[] = [];
 
     for (let i = 0; i < azNames.length; i++) {
-        subnetInputs
-            .filter((x) => x.type === "Private")
-            .forEach((privateSubnet) => {
-                const newBits = privateSubnet.cidrMask - baseSubnetMask;
-                const privateSubnetCidrBlock = cidrSubnetV4(
-                    azBases[i],
-                    newBits,
-                    0,
-                );
-                privateSubnets.push({
-                    azName: azNames[i],
-                    cidrBlock: privateSubnetCidrBlock,
-                    type: "Private",
-                    subnetName: `${vpcName}-${
-                        privateSubnet.name ?? "private"
-                    }-${i + 1}`,
-                });
+        const privateSubnetsOut: SubnetSpec[] = [];
+        const publicSubnetsOut: SubnetSpec[] = [];
+        const isolatedSubnetsOut: SubnetSpec[] = [];
+
+        for (let j = 0; j < privateSubnetsIn.length; j++) {
+            const newBits = privateSubnetsIn[j].cidrMask - baseSubnetMask;
+
+            // The "j" input to cidrSubnetV4 below covers the case where we have
+            // multiple subnets of each type per AZ. In this case, we need the
+            // nth subnet of size newBits in the block for the AZ. The other
+            // subnet types below have similar logic.
+            const privateSubnetCidrBlock = cidrSubnetV4(
+                azBases[i],
+                newBits,
+                j,
+            );
+            privateSubnetsOut.push({
+                azName: azNames[i],
+                cidrBlock: privateSubnetCidrBlock,
+                type: "Private",
+                subnetName: `${vpcName}-${privateSubnetsIn[j].name ?? "private"}-${i + 1}`,
             });
-
-        subnetInputs
-            .filter((x) => x.type === "Public")
-            .forEach((publicSubnet) => {
-                const baseCidr =
-                    privateSubnets.length > 0
-                        ? privateSubnets[i].cidrBlock
-                        : azBases[i];
-
-                const baseIp = new ipAddress.Address4(baseCidr);
-
-                const splitBase =
-                    privateSubnets.length > 0
-                        ? cidrSubnetV4(baseCidr, 0, 1)
-                        : azBases[i];
-                const newBits = publicSubnet.cidrMask - baseIp.subnetMask;
-
-                const publicSubnetCidrBlock = cidrSubnetV4(
-                    splitBase,
-                    newBits,
-                    0,
-                );
-                publicSubnets.push({
-                    azName: azNames[i],
-                    cidrBlock: publicSubnetCidrBlock,
-                    type: "Public",
-                    subnetName: `${vpcName}-${publicSubnet.name ?? "public"}-${
-                        i + 1
-                    }`,
-                });
+        }
+
+        for (let j = 0; j < publicSubnetsIn.length; j++) {
+            // If we declared any private subnets in this AZ, we want our public
+            // subnet to start after the last private subnet. If not, we can
+            // start with the CIDR block for the AZ:
+            const baseCidr =
+                privateSubnetsOut.length > 0
+                    ? privateSubnetsOut[privateSubnetsOut.length - 1].cidrBlock
+                    : azBases[i];
+
+            const baseIp = new ipAddress.Address4(baseCidr);
+
+            const splitBase =
+                privateSubnetsOut.length > 0
+                    ? cidrSubnetV4(baseCidr, 0, 1)
+                    : azBases[i];
+            const newBits = publicSubnetsIn[j].cidrMask - baseIp.subnetMask;
+
+            const publicSubnetCidrBlock = cidrSubnetV4(
+                splitBase,
+                newBits,
+                j,
+            );
+            publicSubnetsOut.push({
+                azName: azNames[i],
+                cidrBlock: publicSubnetCidrBlock,
+                type: "Public",
+                subnetName: `${vpcName}-${publicSubnetsIn[j].name ?? "public"}-${i + 1}`,
             });
-
-        subnetInputs
-            .filter((x) => x.type === "Isolated")
-            .forEach((isolatedSubnet) => {
-                let baseCidr: string;
-                if (publicSubnets.length > 0) {
-                    baseCidr = publicSubnets[i].cidrBlock;
-                } else if (privateSubnets.length > 0) {
-                    baseCidr = privateSubnets[i].cidrBlock;
-                } else {
-                    baseCidr = azBases[i];
-                }
-
-                const baseIp = new ipAddress.Address4(baseCidr);
-
-                let splitBase: string;
-                if (publicSubnets.length > 0 || privateSubnets.length > 0) {
-                    splitBase = cidrSubnetV4(baseCidr, 0, 1);
-                } else {
-                    splitBase = azBases[i];
-                }
-
-                const newBits = isolatedSubnet.cidrMask - baseIp.subnetMask;
-
-                const isolatedSubnetCidrBlock = cidrSubnetV4(
-                    splitBase,
-                    newBits,
-                    0,
-                );
-                isolatedSubnets.push({
-                    azName: azNames[i],
-                    cidrBlock: isolatedSubnetCidrBlock,
-                    type: "Isolated",
-                    subnetName: `${vpcName}-${
-                        isolatedSubnet.name ?? "isolated"
-                    }-${i + 1}`,
-                });
+        }
+
+        for (let j = 0; j < isolatedSubnetsIn.length; j++) {
+            // Similar to above, if we declared any public subnets in this AZ,
+            // we want our isolated subnet to start after the last public
+            // subnet. If no public subnets, then after the last private subnet.
+            // If no private subnets either, use the base CIDR block for the AZ:
+            let baseCidr: string;
+            if (publicSubnetsOut.length > 0) {
+                baseCidr = publicSubnetsOut[publicSubnetsOut.length - 1].cidrBlock;
+            } else if (privateSubnetsOut.length > 0) {
+                baseCidr = privateSubnetsOut[privateSubnetsOut.length - 1].cidrBlock;
+            } else {
+                baseCidr = azBases[i];
+            }
+
+            const baseIp = new ipAddress.Address4(baseCidr);
+
+            let splitBase: string;
+            if (publicSubnetsOut.length > 0 || privateSubnetsOut.length > 0) {
+                splitBase = cidrSubnetV4(baseCidr, 0, 1);
+            } else {
+                splitBase = azBases[i];
+            }
+
+            const newBits = isolatedSubnetsIn[j].cidrMask - baseIp.subnetMask;
+
+            const isolatedSubnetCidrBlock = cidrSubnetV4(
+                splitBase,
+                newBits,
+                j,
+            );
+            isolatedSubnetsOut.push({
+                azName: azNames[i],
+                cidrBlock: isolatedSubnetCidrBlock,
+                type: "Isolated",
+                subnetName: `${vpcName}-${isolatedSubnetsIn[j].name ?? "isolated"}-${i + 1}`,
             });
+        }
+
+        subnetsOut = subnetsOut.concat(privateSubnetsOut, publicSubnetsOut, isolatedSubnetsOut);
     }
 
-    return Array.prototype.concat(
-        privateSubnets,
-        publicSubnets,
-        isolatedSubnets,
-    );
+    return subnetsOut;
 }
 
 function generateDefaultSubnets(
@@ -192,7 +196,7 @@ function cidrSubnetV4(
     if (newSubnetMask > 32) {
         throw new Error(
             `Requested ${newBits} new bits, but ` +
-                `only ${32 - ipv4.subnetMask} are available.`,
+            `only ${32 - ipv4.subnetMask} are available.`,
         );
     }
 
diff --git a/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/Pulumi.yaml b/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/Pulumi.yaml
new file mode 100644
index 00000000..a7761ba0
--- /dev/null
+++ b/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/Pulumi.yaml
@@ -0,0 +1,3 @@
+name: nodejs
+runtime: nodejs
+description: AWSX VPC - specifying multiple similar subnet types
diff --git a/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/index.ts b/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/index.ts
new file mode 100644
index 00000000..b774aaab
--- /dev/null
+++ b/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/index.ts
@@ -0,0 +1,29 @@
+import * as awsx from "@pulumi/awsx";
+
+const myVpc = new awsx.ec2.Vpc("awsx-nodejs-multiple-subnets", {
+    subnetSpecs: [
+        {
+            type: awsx.ec2.SubnetType.Public,
+            cidrMask: 22,
+        },
+        {
+            type: awsx.ec2.SubnetType.Private,
+            cidrMask: 21,
+        },
+        {
+            type: awsx.ec2.SubnetType.Isolated,
+            cidrMask: 24,
+            name: "db",
+        },
+        {
+            type: awsx.ec2.SubnetType.Isolated,
+            cidrMask: 24,
+            name: "redis",
+        },
+    ],
+});
+
+export const vpcId = myVpc.vpcId;
+export const publicSubnetIds = myVpc.publicSubnetIds;
+export const privateSubnetIds = myVpc.privateSubnetIds;
+export const isolatedSubnetIds = myVpc.isolatedSubnetIds;
diff --git a/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/package.json b/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/package.json
new file mode 100644
index 00000000..36f7bc2b
--- /dev/null
+++ b/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/package.json
@@ -0,0 +1,9 @@
+{
+    "name": "nodejs",
+    "devDependencies": {
+        "@types/node": "^14"
+    },
+    "dependencies": {
+        "@pulumi/pulumi": "^3.0.0"
+    }
+}
diff --git a/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/tsconfig.json b/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/tsconfig.json
new file mode 100644
index 00000000..2dd89fb9
--- /dev/null
+++ b/examples/vpc/nodejs/vpc-multiple-similar-subnet-types/tsconfig.json
@@ -0,0 +1,18 @@
+{
+    "compilerOptions": {
+        "strict": true,
+        "outDir": "bin",
+        "target": "es2016",
+        "module": "commonjs",
+        "moduleResolution": "node",
+        "sourceMap": true,
+        "experimentalDecorators": true,
+        "pretty": true,
+        "noFallthroughCasesInSwitch": true,
+        "noImplicitReturns": true,
+        "forceConsistentCasingInFileNames": true
+    },
+    "files": [
+      "index.ts"
+    ]
+}


Test patch:
diff --git a/awsx/ec2/subnetDistributor.test.ts b/awsx/ec2/subnetDistributor.test.ts
index b3fe36b1..af3846e6 100644
--- a/awsx/ec2/subnetDistributor.test.ts
+++ b/awsx/ec2/subnetDistributor.test.ts
@@ -21,8 +21,8 @@ describe("getSubnetSpecs", () => {
     const vpcName = "vpcname";
 
     it("should return the default subnets with no parameters and 3 AZs", () => {
-       const result = getSubnetSpecs(vpcName, vpcCidr, azs);
-       const expected: SubnetSpec[] = [
+        const result = getSubnetSpecs(vpcName, vpcCidr, azs);
+        const expected: SubnetSpec[] = [
             {
                 type: "Private",
                 cidrBlock: "10.0.0.0/19",
@@ -60,7 +60,7 @@ describe("getSubnetSpecs", () => {
                 subnetName: "vpcname-public-3",
             },
         ];
-       expect(result).toEqual(expected);
+        expect(result).toEqual(expected);
     });
 
     describe
@@ -127,30 +127,30 @@ describe("getSubnetSpecs", () => {
                     azName: "us-east-1a",
                     subnetName: "vpcname-foo-1",
                 },
-                {
-                    type: slash19Type,
-                    cidrBlock: "10.0.64.0/19",
-                    azName: "us-east-1b",
-                    subnetName: "vpcname-foo-2",
-                },
-                {
-                    type: slash19Type,
-                    cidrBlock: "10.0.128.0/19",
-                    azName: "us-east-1c",
-                    subnetName: "vpcname-foo-3",
-                },
                 {
                     type: slash20Type,
                     cidrBlock: "10.0.32.0/20",
                     azName: "us-east-1a",
                     subnetName: "vpcname-bar-1",
                 },
+                {
+                    type: slash19Type,
+                    cidrBlock: "10.0.64.0/19",
+                    azName: "us-east-1b",
+                    subnetName: "vpcname-foo-2",
+                },
                 {
                     type: slash20Type,
                     cidrBlock: "10.0.96.0/20",
                     azName: "us-east-1b",
                     subnetName: "vpcname-bar-2",
                 },
+                {
+                    type: slash19Type,
+                    cidrBlock: "10.0.128.0/19",
+                    azName: "us-east-1c",
+                    subnetName: "vpcname-foo-3",
+                },
                 {
                     type: slash20Type,
                     cidrBlock: "10.0.160.0/20",
diff --git a/awsx/ec2/subnetsIntegration.test.ts b/awsx/ec2/subnetsIntegration.test.ts
new file mode 100644
index 00000000..dafc2415
--- /dev/null
+++ b/awsx/ec2/subnetsIntegration.test.ts
@@ -0,0 +1,155 @@
+// Copyright 2016-2022, Pulumi Corporation.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { SubnetSpecInputs } from "../schema-types";
+import { getSubnetSpecs } from "./subnetDistributor";
+import { getOverlappingSubnets } from "./vpc";
+
+describe("subnet creation integration tests", () => {
+  it("should return no overlapping subnets for the default subnet specs for this component", () => {
+    const defaultSpecs = getSubnetSpecs("dummy", "10.0.0.0/16", [
+      "us-east-1a",
+      "us-east-1b",
+      "us-east-1c",
+    ]);
+    expect(getOverlappingSubnets(defaultSpecs)).toEqual([]);
+  });
+
+  // TODO: Consider moving these back to unit tests since they are simple enough to make assertions on:
+  it("should return no overlapping subnets 2 public subnets per AZ", () => {
+    const inputs: SubnetSpecInputs[] = [
+      {
+        type: "Public",
+        cidrMask: 22,
+      },
+      {
+        type: "Public",
+        cidrMask: 22,
+      },
+    ];
+
+    const subnetSpecs = getSubnetSpecs("dummy", "10.0.0.0/16", ["us-east-1a", "us-east-1b", "us-east-1c"], inputs);
+
+    expect(getOverlappingSubnets(subnetSpecs)).toEqual([]);
+  });
+
+  it("should return no overlapping subnets 2 private subnets per AZ", () => {
+    const inputs: SubnetSpecInputs[] = [
+      {
+        type: "Private",
+        cidrMask: 22,
+      },
+      {
+        type: "Private",
+        cidrMask: 22,
+      },
+    ];
+
+    const subnetSpecs = getSubnetSpecs("dummy", "10.0.0.0/16", ["us-east-1a", "us-east-1b", "us-east-1c"], inputs);
+
+    expect(getOverlappingSubnets(subnetSpecs)).toEqual([]);
+  });
+
+  it("should return no overlapping subnets 2 isolated subnets per AZ", () => {
+    const inputs: SubnetSpecInputs[] = [
+      {
+        type: "Isolated",
+        cidrMask: 22,
+      },
+      {
+        type: "Isolated",
+        cidrMask: 22,
+      },
+    ];
+
+    const subnetSpecs = getSubnetSpecs("dummy", "10.0.0.0/16", ["us-east-1a", "us-east-1b", "us-east-1c"], inputs);
+
+    expect(getOverlappingSubnets(subnetSpecs)).toEqual([]);
+  });
+
+  it("should return no overlapping subnets with 2 privates and 1 public per AZ", () => {
+    const inputs: SubnetSpecInputs[] = [
+      {
+        type: "Private",
+        cidrMask: 22,
+      },
+      {
+        type: "Private",
+        cidrMask: 22,
+      },
+      {
+        type: "Public",
+        cidrMask: 21,
+      },
+    ];
+
+    const subnetSpecs = getSubnetSpecs("dummy", "10.0.0.0/16", ["us-east-1a", "us-east-1b", "us-east-1c"], inputs);
+
+    expect(getOverlappingSubnets(subnetSpecs)).toEqual([]);
+  });
+
+  it("should return no overlapping subnets with 2 publics and 1 isolated per AZ", () => {
+    const inputs: SubnetSpecInputs[] = [
+      {
+        type: "Public",
+        cidrMask: 22,
+      },
+      {
+        type: "Public",
+        cidrMask: 22,
+      },
+      {
+        type: "Isolated",
+        cidrMask: 21,
+      },
+    ];
+
+    const subnetSpecs = getSubnetSpecs("dummy", "10.0.0.0/16", ["us-east-1a", "us-east-1b", "us-east-1c"], inputs);
+
+    expect(getOverlappingSubnets(subnetSpecs)).toEqual([]);
+  });
+
+  it("should return no overlapping subnets with 2 of each type per AZ", () => {
+    const inputs: SubnetSpecInputs[] = [
+      {
+        type: "Public",
+        cidrMask: 24,
+      },
+      {
+        type: "Public",
+        cidrMask: 24,
+      },
+      {
+        type: "Private",
+        cidrMask: 24,
+      },
+      {
+        type: "Private",
+        cidrMask: 24,
+      },
+      {
+        type: "Isolated",
+        cidrMask: 24,
+      },
+      {
+        type: "Isolated",
+        cidrMask: 24,
+      },
+    ];
+
+    const subnetSpecs = getSubnetSpecs("dummy", "10.0.0.0/16", ["us-east-1a", "us-east-1b", "us-east-1c"], inputs);
+
+    expect(getOverlappingSubnets(subnetSpecs)).toEqual([]);
+  });
+});
diff --git a/awsx/ec2/vpc.test.ts b/awsx/ec2/vpc.test.ts
index f7ca5003..79736442 100644
--- a/awsx/ec2/vpc.test.ts
+++ b/awsx/ec2/vpc.test.ts
@@ -244,16 +244,6 @@ describe("getOverlappingSubnets", () => {
             ]),
         ).toEqual([subnet2a, subnet1, subnet1a, subnet2]);
     });
-
-    // This is technically an integration test, but it's probably better to just call getSubnetSpecs than it is to restate the default subnet specs as a literal:
-    it("should return no overlapping subnets for the default subnet specs for this component", () => {
-        const defaultSpecs = getSubnetSpecs("dummy", "10.0.0.0/16", [
-            "us-east-1a",
-            "us-east-1b",
-            "us-east-1c",
-        ]);
-        expect(getOverlappingSubnets(defaultSpecs)).toEqual([]);
-    });
 });
 
 describe("validateSubnets", () => {
diff --git a/examples/examples_nodejs_test.go b/examples/examples_nodejs_test.go
index d0f4d1ac..2f1abdd9 100644
--- a/examples/examples_nodejs_test.go
+++ b/examples/examples_nodejs_test.go
@@ -124,6 +124,16 @@ func TestVpcSpecificSubnetSpecArgs(t *testing.T) {
 	integration.ProgramTest(t, &test)
 }
 
+func TestVpcMultipleSimilarSubnetSpecArgs(t *testing.T) {
+	test := getNodeJSBaseOptions(t).
+		With(integration.ProgramTestOptions{
+			RunUpdateTest: false,
+			Dir:           filepath.Join(getCwd(t), "vpc", "nodejs", "vpc-multiple-similar-subnet-types"),
+		})
+
+	integration.ProgramTest(t, &test)
+}
+
 func getNodeJSBaseOptions(t *testing.T) integration.ProgramTestOptions {
 	base := getBaseOptions(t)
 	nodeBase := base.With(integration.ProgramTestOptions{