Lamba function for NAT Gateway custom resource & Cloudformation template to use it

gateway.json

{
  "AWSTemplateFormatVersion": "2010-09-09",
  "Description": "VPC setup",
  "Parameters": {},
  "Mappings": {},
  "Conditions": {},
  "Resources": {
    "VPC": {
      "Type": "AWS::EC2::VPC",
      "Properties": {
        "CidrBlock": "10.0.0.0/16",
        "Tags": [
          {
            "Key": "Name",
            "Value": "nat gateway demo vpc"
          }
        ]
      }
    },
    "InternetGateway": {
      "Type": "AWS::EC2::InternetGateway"
    },
    "InternetGatewayAttachement": {
      "Type": "AWS::EC2::VPCGatewayAttachment",
      "Properties": {
        "InternetGatewayId": {
          "Ref": "InternetGateway"
        },
        "VpcId": {
          "Ref": "VPC"
        }
      }
    },
    "NatIP": {
      "Type": "AWS::EC2::EIP",
      "DependsOn": "InternetGatewayAttachement",
      "Properties": {
        "Domain": "vpc"
      }
    },
    "NatSubnet": {
      "Type": "AWS::EC2::Subnet",
      "Properties": {
        "AvailabilityZone": {
          "Fn::Join": [
            "",
            [
              {
                "Ref": "AWS::Region"
              },
              "a"
            ]
          ]
        },
        "CidrBlock": "10.0.0.0/24",
        "VpcId": {
          "Ref": "VPC"
        },
        "Tags": [
          {
            "Key": "Name",
            "Value": "nat subnet"
          }
        ]

      }
    },

    "GatewayWaitHandle" : {
     "Type" : "AWS::CloudFormation::WaitConditionHandle",
     "Properties" : {
       }
    },

    "NatGateway": {
      "Type": "Custom::NatGateway",
      "Properties": {
        "ServiceToken": {
          "Fn::GetAtt": [ "CustomResourceFunction", "Arn" ]
        },
        "SubnetId": { "Ref": "NatSubnet" },
        "AllocationId": {
          "Fn::GetAtt": ["NatIP", "AllocationId"]
        },
        "WaitHandle": {"Ref": "GatewayWaitHandle"}
      }
    },
    "GatewayWaitCondition" : {
      "Type" : "AWS::CloudFormation::WaitCondition",
      "DependsOn" : "NatGateway",
      "Properties" : {
        "Handle" : { "Ref" : "GatewayWaitHandle" },
        "Timeout" : "240"
      }
    },
    "PrivateSubnet": {
      "Type": "AWS::EC2::Subnet",
      "Properties": {
        "AvailabilityZone": {
          "Fn::Join": [
            "",
            [
              {
                "Ref": "AWS::Region"
              },
              "a"
            ]
          ]
        },
        "CidrBlock": "10.0.1.0/24",
        "VpcId": {
          "Ref": "VPC"
        },
        "Tags": [
          {
            "Key": "Name",
            "Value": "Private subnet"
          }
        ]

      }
    },
    "PublicSubnetsRouteTable": {
      "Type": "AWS::EC2::RouteTable",
      "Properties": {
        "VpcId": {
          "Ref": "VPC"
        }
      }
    },
    "PrivateSubnetsRouteTable": {
      "Type": "AWS::EC2::RouteTable",
      "Properties": {
        "VpcId": {
          "Ref": "VPC"
        }
      }
    },
    "InternetRoute": {
      "Type": "AWS::EC2::Route",
      "DependsOn": "InternetGateway",
      "Properties": {
        "RouteTableId": {
          "Ref": "PublicSubnetsRouteTable"
        },
        "DestinationCidrBlock": "0.0.0.0/0",
        "GatewayId": {
          "Ref": "InternetGateway"
        }
      }
    },
    "NatRoute": {
      "Type": "Custom::NatGatewayRoute",
      "DependsOn": [
        "GatewayWaitCondition",
        "AssociateRouteTableWithNatSubnet"
      ],
      "Properties": {
        "ServiceToken": {
          "Fn::GetAtt": [ "CustomResourceFunction", "Arn" ]
        },
        "RouteTableId": {
          "Ref": "PrivateSubnetsRouteTable"
        },
        "DestinationCidrBlock": "0.0.0.0/0",
        "NatGatewayId": {
          "Ref": "NatGateway"
        }
      }
    },
    "AssociateRouteTableWithNatSubnet": {
      "Type": "AWS::EC2::SubnetRouteTableAssociation",
      "Properties": {
        "RouteTableId": {
          "Ref": "PublicSubnetsRouteTable"
        },
        "SubnetId": {
          "Ref": "NatSubnet"
        }
      }
    },
    "AssociateRouteTableWithPrivateSubnet": {
      "Type": "AWS::EC2::SubnetRouteTableAssociation",
      "Properties": {
        "RouteTableId": {
          "Ref": "PrivateSubnetsRouteTable"
        },
        "SubnetId": {
          "Ref": "PrivateSubnet"
        }
      }
    },
    "LambdaExecutionRole": {
      "Type": "AWS::IAM::Role",
      "Properties": {
        "AssumeRolePolicyDocument": {
          "Version": "2012-10-17",
          "Statement": [
            {
              "Effect": "Allow",
              "Principal": {
                "Service": [
                  "lambda.amazonaws.com"
                ]
              },
              "Action": [
                "sts:AssumeRole"
              ]
            }
          ]
        },
        "Path": "/",
        "Policies": [
          {
            "PolicyName": "root",
            "PolicyDocument": {
              "Version": "2012-10-17",
              "Statement": [
                {
                  "Effect": "Allow",
                  "Action": [
                    "logs:CreateLogGroup",
                    "logs:CreateLogStream",
                    "logs:PutLogEvents"
                  ],
                  "Resource": "arn:aws:logs:*:*:*"
                },
                {
                  "Effect": "Allow",
                  "Action": [
                    "ec2:DescribeNatGateways",
                    "ec2:DeleteRoute",
                    "ec2:ReplaceRoute",
                    "ec2:CreateRoute",
                    "ec2:DeleteNatGateway",
                    "ec2:CreateNatGateway"
                  ],
                  "Resource": "*"
                }
              ]
            }
          }
        ]
      }
    },
    "CustomResourceFunction": {
      "Type": "AWS::Lambda::Function",
      "Properties": {
        "Description": "Lambda function for using nat gateways with CloudFormation",
        "Code": {
          "S3Bucket": "fred-lambda",
          "S3Key": "nat_gateway.zip"
        },
        "Handler": "nat_gateway.handler",
        "Runtime": "nodejs",
        "Timeout": "240",
        "Role": {
          "Fn::GetAtt": [
            "LambdaExecutionRole",
            "Arn"
          ]
        }
      }
    }
  },
  "Outputs": {
    "VPC": {
      "Value": {
        "Ref": "VPC"
      },
      "Description": "The id of the created vpc"
    },
    "NatIP": {
      "Value": {
        "Ref": "NatIP"
      },
      "Description": "The elastic ip used by the nat"
    }
  }
}

nat_gateway.js

var aws = require('aws-sdk');
var ec2 = new aws.EC2();

exports.handler = function(event, context) {
    if (event.ResourceType === 'Custom::NatGateway') {
        handleGateway(event, context);
    } else if (event.ResourceType === 'Custom::NatGatewayRoute') {
        handleRoute(event, context);
    } else {
        console.log("Unknown resource type: " + event.ResourceType);
        response.send(event, context, {
            Error: "unknown resource type: " + event.ResourceType
        }, response.FAILED);
    }
};

var handleRoute = function(event, context) {
    var destinationCidrBlock = event.ResourceProperties.DestinationCidrBlock;
    var routeTableId = event.ResourceProperties.RouteTableId;
    var responseData = {};
    if (!destinationCidrBlock) {
        responseData = {
            Error: "missing parameter DestinationCidrBlock "
        };
        console.log(responseData.Error);
        response.send(event, context, response.FAILED, responseData);
        return;
    }
    else {
        if (!routeTableId) {
          responseData = {
              Error: "missing parameter RouteTableId "
          };
          console.log(responseData.Error);
          response.send(event, context, response.FAILED, responseData);
          return;
      }
    }

    if (event.RequestType === 'Delete') {
        deleteRoute(event, context);
    } else if (event.RequestType === 'Create') {
        createRoute(event, context);
    } else if (event.RequestType === 'Update') {
        if (event.ResourceProperties.DestinationCIDRBlock === event.OldResourceProperties.DestinationCIDRBlock &&
            event.ResourceProperties.RouteTableId === event.OldResourceProperties.RouteTableId) {
            replaceRoute(event, context);
        } else {
            createRoute(event, context);
        }
    } else {
        console.log("Unknown request type " + event.RequestType);
        response.send(event, context, {
            Error: "unknown request type: " + event.RequestType
        }, response.FAILED);
    }
};

var deleteRoute = function(event, context) {
    var responseData = {};
    var destinationCidrBlock = event.ResourceProperties.DestinationCidrBlock;
    var routeTableId = event.ResourceProperties.RouteTableId;

    if(event.PhysicalResourceId.match(/^gateway-route-/)){

      ec2.deleteRoute({
          RouteTableId: routeTableId,
          DestinationCidrBlock: destinationCidrBlock
      }, function(err, data) {
          if (err) {
              responseData = {
                  Error: "delete route failed " + err
              };
              console.log(responseData.Error);
              response.send(event, context, response.FAILED, responseData);

          } else {
              response.send(event, context, response.SUCCESS, {}, physicalId(event.ResourceProperties));
          }
      });
    }else{
      console.log("unexpected physical id for route " + event.PhysicalResourceId + " - ignoring");
      response.send(event, context, response.SUCCESS, {});
    }
};


var createRoute = function(event, context) {
    var responseData = {};
    var destinationCidrBlock = event.ResourceProperties.DestinationCidrBlock;
    var routeTableId = event.ResourceProperties.RouteTableId;
    var natGatewayId = event.ResourceProperties.NatGatewayId;

    if (natGatewayId) {
        ec2.createRoute({
            RouteTableId: routeTableId,
            DestinationCidrBlock: destinationCidrBlock,
            NatGatewayId: natGatewayId
        }, function(err, data) {
            if (err) {
                responseData = {
                    Error: "create route failed " + err
                };
                console.log(responseData.Error);
                response.send(event, context, response.FAILED, responseData);

            } else {
                response.send(event, context, response.SUCCESS, {}, physicalId(event.ResourceProperties));
            }
        });
    } else {
        responseData = {
            Error: "missing parameter natGatewayId "
        };
        console.log(responseData.Error);
        response.send(event, context, response.FAILED, responseData);
        return;
    }
};

var replaceRoute = function(event, context) {
    var responseData = {};
    var destinationCidrBlock = event.ResourceProperties.DestinationCidrBlock;
    var routeTableId = event.ResourceProperties.RouteTableId;
    var natGatewayId = event.ResourceProperties.NatGatewayId;

    if (natGatewayId) {
        ec2.replaceRoute({
            RouteTableId: routeTableId,
            DestinationCidrBlock: destinationCidrBlock,
            NatGatewayId: natGatewayId
        }, function(err, data) {
            if (err) {
                responseData = {
                    Error: "create route failed " + err
                };
                console.log(responseData.Error);
                response.send(event, context, response.FAILED, responseData);

            } else {
                response.send(event, context, response.SUCCESS, {}, physicalId(event.ResourceProperties));
            }
        });
    } else {
        responseData = {
            Error: "missing parameter natGatewayId "
        };
        console.log(responseData.Error);
        response.send(event, context, response.FAILED, responseData);
        return;
    }
};


var physicalId = function(properties) {
    return 'gateway-route-' + properties.RouteTableId + '-' + properties.DestinationCIDRBlock;
};


var handleGateway = function(event, context) {
    if (event.RequestType === 'Delete') {
        deleteGateway(event, context);
    } else if (event.RequestType === 'Update' || event.RequestType === 'Create') {
        createGateway(event, context);
    } else {
        response.send(event, context, {
            Error: "unknown type: " + event.RequestType
        }, response.FAILED);
    }
};

var createGateway = function(event, context) {
    var responseData = {};
    var subnetId = event.ResourceProperties.SubnetId;
    var allocationId = event.ResourceProperties.AllocationId;
    var waitHandle = event.ResourceProperties.WaitHandle;

    if (subnetId && allocationId) {
        ec2.createNatGateway({
            AllocationId: allocationId,
            SubnetId: subnetId
        }, function(err, data) {
            if (err) {
                responseData = {
                    Error: "create gateway failed " + err
                };
                console.log(responseData.Error);
                response.send(event, context, response.FAILED, responseData);
            } else {
                responseData = {
                }
                response.send(event, context, response.SUCCESS, responseData, data.NatGateway.NatGatewayId, true);

                waitForGatewayStateChange(data.NatGateway.NatGatewayId, ['available', 'failed'], function(state){
                  if(waitHandle){
                    signalData = {
                      "Status": state == 'available' ? 'SUCCESS' : 'FAILURE',
                      "UniqueId": data.NatGateway.NatGatewayId,
                      "Data": "Gateway has state " + state,
                      "Reason": ""
                    }
                    sendSignal(waitHandle, context, signalData);
                  }else{
                    if(state != 'available'){
                      console.log("gateway state is not available");
                    }
                    context.done();
                  }
                });
            }
        })
    } else {
        if (!subnetId) {
            responseData = {
                Error: 'subnet id not specified'
            };
            console.log(responseData.Error);
            response.send(event, context, response.FAILED, responseData);
        } else {
            responseData = {
                Error: 'allocationId not specified'
            };
            console.log(responseData.Error);
            response.send(event, context, response.FAILED, responseData);
        }
    }
};

var waitForGatewayStateChange = function (id, states, onComplete){
  ec2.describeNatGateways({NatGatewayIds: [id], Filter: [{Name: "state", Values: states}]}, function(err, data){
    if(err){
      console.log("could not describeNatGateways " + err);
      onComplete('failed');
    }else{
      if(data.NatGateways.length > 0){
        onComplete(data.NatGateways[0].State)
      }else{
        console.log("gateway not ready; waiting");
        setTimeout(function(){ waitForGatewayStateChange(id, states, onComplete);}, 15000);
      }
    }
  });
};

var deleteGateway = function(event, context) {
    var responseData = {};
    if (event.PhysicalResourceId && event.PhysicalResourceId.match(/^nat-/)) {
        ec2.deleteNatGateway({
            NatGatewayId: event.PhysicalResourceId
        }, function(err, data) {
            if (err) {
                responseData = {
                    Error: "delete gateway failed " + err
                };
                console.log(responseData.Error);
                response.send(event, context, response.FAILED, responseData, event.PhysicalResourceId);
            } else {
                waitForGatewayStateChange(event.PhysicalResourceId, ['deleted'], function(state){
                  response.send(event, context, response.SUCCESS, {}, event.PhysicalResourceId);
                });
            }
        })
    } else {
        console.log("No valid physical resource id passed to destroy - ignoring " + event.PhysicalResourceId);
        response.send(event, context, response.SUCCESS, responseData, event.PhysicalResourceId);
    }
}


var sendSignal = function(handle, context, data){
    var body = JSON.stringify(data);
    var https = require("https");
    var url = require("url");
    console.log("signal body:\n", body);

    var parsedUrl = url.parse(handle);
    var options = {
        hostname: parsedUrl.hostname,
        port: 443,
        path: parsedUrl.path,
        method: "PUT",
        headers: {
            "content-type": "",
            "content-length": body.length
        }
    };

    var request = https.request(options, function(response) {
        console.log("Status code: " + response.statusCode);
        console.log("Status message: " + response.statusMessage);
        context.done();
    });

    request.on("error", function(error) {
        console.log("sendSignal(..) failed executing https.request(..): " + error);
        context.done();
    });

    request.write(body);
    request.end();
};
/* The below section is adapted from the cfn-response module, as published at:
   
   http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-lambda-function-code.html

   */

/* Copyright 2015 Amazon Web Services, Inc. or its affiliates. All Rights Reserved.
   This file is licensed to you under the AWS Customer Agreement (the "License").
   You may not use this file except in compliance with the License.
   A copy of the License is located at http://aws.amazon.com/agreement/.
   This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied.
   See the License for the specific language governing permissions and limitations under the License. */
var response = {}

response.SUCCESS = "SUCCESS";
response.FAILED = "FAILED";

response.send = function(event, context, responseStatus, responseData, physicalResourceId, continueFuncton) {

    var responseBody = JSON.stringify({
        Status: responseStatus,
        Reason: "See the details in CloudWatch Log Stream: " + context.logStreamName,
        PhysicalResourceId: physicalResourceId || context.logStreamName,
        StackId: event.StackId,
        RequestId: event.RequestId,
        LogicalResourceId: event.LogicalResourceId,
        Data: responseData
    });

    console.log("Response body:\n", responseBody);

    var https = require("https");
    var url = require("url");

    var parsedUrl = url.parse(event.ResponseURL);
    var options = {
        hostname: parsedUrl.hostname,
        port: 443,
        path: parsedUrl.path,
        method: "PUT",
        headers: {
            "content-type": "",
            "content-length": responseBody.length
        }
    };

    var request = https.request(options, function(response) {
        console.log("Status code: " + response.statusCode);
        console.log("Status message: " + response.statusMessage);
        if(!continueFuncton){
          context.done();
        }
    });

    request.on("error", function(error) {
        console.log("send(..) failed executing https.request(..): " + error);
        context.done();
    });

    request.write(responseBody);
    request.end();
}

Hi @fcheung,
thank you so much for sharing this. I've been planning on doing the same thing so this is a great starting point.
I have a question: Why did you go the route of using a GatewayWaitHandle/GatewayWaitCondition? If you're waiting for the NAT Gateway to become available inside the Lambda function anyway, why can't you reduce the overall complexity and do a response.send() instead depend on the result of the custom resource instead (that's what you're doing when deleting the gateway)?

Thanks for the work. You may want to checkout: https://www.npmjs.com/package/cfn-lambda. It simplifies the process of building and deploying custom resources on top of Lambda.

this is great work and very useful. thank you for showing a way to make this work. I've been testing it out and noticed problems when I delete a stack that uses this method. The EIP on the NAT Gateway still shows as attached even after the gateway custom resource has been deleted. This prevents me from deleting the VPC. Any thoughts on how to resolve that? Perhaps explicitly delete the EIP as part of the custom delete handling?

Great tutorial - thanks. I get an issue when I delete the stack:

AWS::EC2::VPCGatewayAttachment
InternetGatewayAttachement
Network vpc-77665812 has some mapped public address(es). Please unmap those public address(es) before detaching the gateway.

It appears to try deleting the attachment before the gateway has been deleted. Not sure how to specify the dependency.

Awesome, thanks much for the tutorial. Great to see the code in github so nicely. You rock!