Streamline CloudFormation Custom Resource Development with a Wrapper Lambda Function
Tags: aws, cloud, devops, serverless
Published: 2023-03-31 02:31:09

AWS CloudFormation and AWS Lambda are powerful tools available to all AWS customers. AWS CloudFormation enables you to define and manage your AWS infrastructure using code, while AWS Lambda allows you to run code in response to events without managing backend infrastructure. A powerful feature is CloudFormation Custom Resources, which enables you to implement a Lambda function triggered by a CloudFormation template. A CloudFormation template can achieve more possibilities by utilizing this feature.

Errors during the development of custom resources can result in unresponsive stacks and delays in the development process. This blog post explores how a wrapper Lambda function can enhance your CloudFormation custom resource development experience, preventing unresponsive stacks and boosting overall productivity. The AWS provided cfn-response module assists in sending responses from your Lambda function, which serves as a CloudFormation Custom Resource, back to CloudFormation. Unresponsive stacks typically arise from coding mistakes, such as typos or incorrect permissions, which hinder CloudFormation from receiving the necessary response from the Lambda function. When a stack becomes unresponsive, it disrupts the development process. A common workaround to an unresponsive stack is selecting a new stack name and continuing development while remembering to clean up the unresponsive stack after its one-hour timeout period expires.

Before we get further into our solution, let's demonstrate how easy it is to create an unresponsive stack accidentally. In this demo, we will create an AWS Lambda function that will try to list all DynamoDB tables in the account without being given permission. Another common cause is a simple typo in the code. You may want to avoid following along in your AWS account because this will create an unresponsive stack, so only perform this yourself if you're ready to deal with the unresponsive stack afterward.

Example in JavaScript:
```
AWSTemplateFormatVersion: "2010-09-09"
Resources:
  LambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          const { DynamoDBClient, ListTablesCommand } = require("@aws-sdk/client-dynamodb");
          const cfnresponse = require("cfn-response");

          const asyncCfnResponseSend = (event, context, responseStatus, responseData, physicalResourceId, noEcho) => {
            return new Promise((resolve, reject) => {
              cfnresponse.send(event, context, responseStatus, responseData, physicalResourceId, noEcho, (err, res) => {
                if (err) {
                  reject(err);
                } else {
                  resolve(res);
                }
              });
            });
          };

          exports.handler = async (event, context) => {
            const client = new DynamoDBClient();
            const response = await client.send(new ListTablesCommand({}));
            await asyncCfnResponseSend(event, context, cfnresponse.SUCCESS, response);
          };

      Handler: index.handler
      Role: !GetAtt LambdaFunctionExecutionRole.Arn
      Runtime: nodejs18.x
      Timeout: 10

  LambdaFunctionExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

  TestInvocationOfCustomResource:
    Type: Custom::TestInvocation
    Properties:
      ServiceToken: !GetAtt LambdaFunction.Arn
```

Example in Python:
```
AWSTemplateFormatVersion: "2010-09-09"
Resources:
  LambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          import boto3
          import cfnresponse

          def handler(event, context):
              client = boto3.client("dynamodb")
              response = client.list_tables()
              cfnresponse.send(event, context, cfnresponse.SUCCESS, response)

      Handler: index.handler
      Role: !GetAtt LambdaFunctionExecutionRole.Arn
      Runtime: python3.9
      Timeout: 10

  LambdaFunctionExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

  TestInvocationOfCustomResource:
    Type: Custom::TestInvocation
    Properties:
      ServiceToken: !GetAtt LambdaFunction.Arn
```

A screenshot of a unresponsive stack:
![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/9rn3lmdczktxl461pqa0.png)

As you see, simply trying to perform an SDK action without proper permissions can create an unresponsive stack.  Of course that above can be avoided by wrapping the SDK action in a try/catch, sometimes when we're quickly developing we may forget that or we may make a typo.  It is detrimental to productivity when your process is interrupted by a unresponsive stack.  To avoid this problem, we can use a wrapper Lambda function to handle exceptions and ensure proper communication back to the CloudFormation service. Our wrapper Lambda function is a reliable approach to avoid creating an unresponsive CloudFormation stack due to quick mistakes during the development of custom resources. This wrapper function acts as a handler that invokes the Lambda being developed and communicates a response to CloudFormation, ensuring that errors are handled gracefully and presenting information of items that require attention for the developers to see. This approach streamlines the development process, making it more efficient and dependable. As a result, custom resources can be developed with confidence, knowing that communication with CloudFormation is reliable and error-free.

To protect ourselves from our development process being interrupted by an unresponsive stack. Let's make some updates to our CloudFormation template. First, we will add a Lambda Function to the Resources block of your CloudFormation template.

In JavaScript:
```
  WrappingLambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          const { LambdaClient, InvokeCommand } = require("@aws-sdk/client-lambda");
          const cfnresponse = require("cfn-response");
          const https = require("https");
          const url = require("url");

          cfnresponse.send = async (event, context, responseStatus, responseData, physicalResourceId, noEcho, reason) => {
            return new Promise((resolve, reject) => {
              try {
                const responseBody = JSON.stringify({
                  Status: responseStatus,
                  Reason: reason || `See the details in CloudWatch Log Stream: ${context.logStreamName}`,
                  PhysicalResourceId: physicalResourceId || context.logStreamName,
                  StackId: event.StackId,
                  RequestId: event.RequestId,
                  LogicalResourceId: event.LogicalResourceId,
                  NoEcho: noEcho || false,
                  Data: responseData,
                });

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

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

                const request = https.request(options, (response) => {
                  console.log("Status code: " + response.statusCode);
                  console.log("Status message: " + response.statusMessage);
                  resolve();
                });

                request.on("error", (error) => {
                  console.log("send(..) failed executing https.request(..): " + error);
                  reject(error);
                });

                request.write(responseBody);
                request.end();
              } catch (error) {
                console.log("send(..) failed: " + error);
                reject(error);
              }
            });
          };

          async function invokeLambda(functionName, event) {
            try {
              const client = new LambdaClient();
              const response = await client.send(
                new InvokeCommand({
                  FunctionName: functionName,
                  Payload: JSON.stringify(event),
                })
              );
              const decoder = new TextDecoder("utf-8");
              const payloadString = decoder.decode(response.Payload);
              const payload = JSON.parse(payloadString);

              return payload;
            } catch (error) {
              throw new Error(`Lambda invocation failed: ${error.message}`);
            }
          }

          exports.handler = async (event, context) => {
            console.log("event:", event);

            const resourceProperties = event.ResourceProperties;
            const functionName = resourceProperties.LambdaFunctionName;

            if (!functionName) {
              const errorMessage = "Lambda invocation failed: LambdaFunctionName must be provided in event.ResourceProperties.";
              console.error(errorMessage);
              await cfnresponse.send(event, context, cfnresponse.FAILED, {}, errorMessage, false, errorMessage);
            } else {
              try {
                const payload = await invokeLambda(functionName, event);

                console.log("lambda response:", payload);

                if (payload.hasOwnProperty("errorMessage")) {
                  const errorMessage = `Lambda invocation failed: ${payload.errorMessage}`;
                  await cfnresponse.send(event, context, cfnresponse.FAILED, {}, errorMessage, false, errorMessage);
                } else {
                  await cfnresponse.send(event, context, cfnresponse.SUCCESS, { Response: JSON.stringify(payload) }, "Lambda invocation succeeded", false, "Lambda invocation succeeded");
                }
              } catch (error) {
                const errorMessage = `Lambda invocation failed: ${error.message}`;
                console.error(errorMessage);
                await cfnresponse.send(event, context, cfnresponse.FAILED, {}, errorMessage, false, errorMessage);
              }
            }
          };

      Handler: index.handler
      Role: !GetAtt WrappingLambdaFunctionExecutionRole.Arn
      Runtime: nodejs18.x
      Timeout: 10
```

In Python:
```
  WrappingLambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          import json
          import logging
          import boto3
          import cfnresponse

          logger = logging.getLogger(__name__)
          logger.setLevel(logging.INFO)
          client = boto3.client("lambda")

          def invoke_lambda(function_name, event):
              try:
                  response = client.invoke(
                      FunctionName=function_name,
                      Payload=json.dumps(event),
                  )
                  payload_stream = response["Payload"]
                  payload = json.loads(payload_stream.read())
                  payload_stream.close()

                  return payload
              except Exception as error:
                  raise Exception(f"Lambda invocation failed: {error}")

          def handler(event, context):
              logger.info(f"event: {event}")

              resource_properties = event.get("ResourceProperties", {})
              function_name = resource_properties.get("LambdaFunctionName")

              if not function_name:
                  error_message = "Lambda invocation failed: LambdaFunctionName must be provided in event.ResourceProperties."
                  logger.error(error_message)
                  cfnresponse.send(event, context, cfnresponse.FAILED, {}, error_message, False, error_message)
              else:
                  try:
                      payload = invoke_lambda(function_name, event)

                      logger.info(f"lambda response: {payload}")

                      if "errorMessage" in payload:
                          error_message = f"Lambda invocation failed: {payload['errorMessage']}"
                          cfnresponse.send(event, context, cfnresponse.FAILED, {}, error_message, False, error_message)
                      else:
                          cfnresponse.send(event, context, cfnresponse.SUCCESS, { "Response": json.dumps(payload) }, "Lambda invocation succeeded", False, "Lambda invocation succeeded")

                  except Exception as error:
                      error_message = f"Lambda invocation failed: {error}"
                      logger.error(error_message)
                      cfnresponse.send(event, context, cfnresponse.FAILED, {}, error_message, False, error_message)

      Handler: index.handler
      Role: !GetAtt WrappingLambdaFunctionExecutionRole.Arn
      Runtime: python3.9
      Timeout: 10
```

Now that we have added that additional lambda function, let's make some final edits to our CloudFormation template to utilize that wrapper lambda function.  This will be the same for Python and JavaScript.
Add the role for that function we just added:
```
  WrappingLambdaFunctionExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: invoke-lambda-function
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action: lambda:InvokeFunction
                Resource: !GetAtt LambdaFunction.Arn
```
Now from the original lambda function we created:
1. the cfnresponse import can be removed.  In the JavaScript version the asyncCfnResponseSend function can be removed as well.
2. the cfnresponse.send line can be replaced with `return response`
3. In the TestInvocationOfCustomResource resource in the CloudFormation template the `ServiceToken: !GetAtt LambdaFunction.Arn` line should be renamed to `LambdaFunctionName: !Ref LambdaFunction`, then add a new line `ServiceToken: !GetAtt WrappingLambdaFunction.Arn` that will update the CloudFormation template to call the wrapping function which will then call the function in development.
4. You can also add an output, so we can demonstrate how return values can flow from our development lambda function back to CloudFormation, at the end of the template, add:
```
Outputs:
  ResponseFromLambdaFunction:
    Value: !GetAtt TestInvocationOfCustomResource.Response
```

Now that we have added our wrapping function and updated our template, lets deploy and take a look.  My results can be seen in this screenshot.

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/9h5jwuwt2lordtetkxtx.png)
As we expected this function did not succeed, as it did not list all DynamoDB tables, because it does not have permissions.  However the stack is still responsive. Note the message in the Status reason column, we are also returning that same value to the Physical ID column on the Resources tab.  We are returning that value and trying to make it visible to the developer without causing the stack to wait for an hour prior to any further actions being performed.

Side note, while writing this post, I discovered that I prefer cfn-response module's Python implementation over its JavaScript counterpart. Since python typically executes code in a single thread, we don't need to worry about cfn-response not being asynchronous.  The JavaScript version of cfn-response also lacks support for the reason parameter. To address these items in our examples above you will see the JavaScript examples required more code to remedy the lack of support for the reason property and cfn-response's send function not being asynchronous.

If you would like to take this demo a step further, we can remediate the fictitious issue we created for demonstration by adding the appropriate permission to the inner (wrapped) Lambda function to list DynamoDB tables and deploy the template.  If you add the appropriate permission to the inner (wrapped) lambda function to list DynamoDB tables and re-create the template you will see how the data from that call is returned back to the wrapping function, which then in turn makes it available to the CloudFormation stack (as an example it is on the outputs tab of the CloudFormation console.)
```
      Policies:
        - PolicyName: temp
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action: dynamodb:ListTables
                Resource: "*"
```

In conclusion, a wrapper Lambda function can be a valuable tool for streamlining your CloudFormation custom resource development process and preventing unresponsive stacks. By incorporating them into your workflow, you can achieve a more efficient and reliable development experience.  Thanks for reading.