API Gateway Custom Authorization With Lambda, DynamoDB, and CloudFormation

API Gateway provides an HTTP API endpoint that is fully configurable. You define the HTTP resources (like /user), the HTTP methods on that resources (like POST, GET, DELETE, etc.) and the integration (e.g. Lambda functions) that should be called to process the request. A Lambda function can then run whatever logic is needed to answer the request. The Lambda function returns its result to the API Gateway. The API Gateway responds to the caller. The following figure demonstrates this flow.

You could include the authentication and authorization logic into the Lambda function that handles the request. But you can also separate concerns, make use of API Gateway caching mechanism, and go for Custom Authorization. API Gateway will invoke another Lambda function (Auth Lambda Function) for the first request and caches that result for a configurable duration. Caching will reduce the overhead (latency and DynamoDB charges) for authentication and authorization to a minimum.

You can use whatever logic you like to decide if a request is allowed or not. In this blog post, I will implement an API token mechanism. All HTTP requests from clients must pass an Authorization: xyz header. The Auth Lambda Function will take this token to query a DynamoDB table. The request is allowed or denied depending on if the query matches.

As usual in this blog, I will use CloudFormation to setup all the needed resources.

CloudFormation Setup

Getting Started

The first thing you need is an API Gateway.

"RestApi": {
    "Type": "AWS::ApiGateway::RestApi",
    "Properties": {
        "Name": "Api"
    }
}

With the logical group in place, you can continue to add all the necessary parts to your API.

Auth Lambda Function

The code for the Auth Lambda Function is responsible for looking up the token. The Authorization HTTP header field is used to transmit the token. You can use Node.js and the AWS SDK for JavaScript to implement this logic. API Gateway will pass an event to our function like this:

{
    "type":"TOKEN",
    "authorizationToken":"<caller-supplied-token>",
    "methodArn":"arn:aws:execute-api:<regionId>:<accountId>:<apiId>/<stage>/<method>/<resourcePath>"
}

API Gateway expects that we respond in the following way:

{
    "principalId": "xyz",
    "policyDocument": {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Action": "execute-api:Invoke",
                "Effect": "Allow", // or Deny
                "Resource": "arn:aws:execute-api:<regionId>:<accountId>:<apiId>/<stage>/<method>/<resourcePath>"
            }
        ]
    }
}

A simple implementation looks like this:

var AWS = require('aws-sdk');
var dynamodb = new AWS.DynamoDB();

function generatePolicy(principalId, effect, resource) {
    return {
        'principalId': principalId,
        'policyDocument': {
            'Version': '2012-10-17',
            'Statement': [{
                'Action': 'execute-api:Invoke',
                'Effect': effect,
                'Resource': resource
            }]
        }
    };
}

exports.handler = function(event, context, cb) {
    var token = event.authorizationToken;
    dynamodb.getItem({
        "Key": {
            "token": {"S": token}
        },
        "TableName": "auth-token"
    }, function(err, data) {
        if (err) {
            cb(err);
        } else {
            if (data.Item === undefined) {
                cb(null, generatePolicy('user', 'Deny', event.methodArn));
            } else {
                cb(null, generatePolicy('user', 'Allow', event.methodArn));
            }
        }
    });
};

Let’s turn this into a CloudFormation template, including the DynamoDB table for the tokens and the IAM Role with permissions to talk to DynamoDB.

"TokenTable": {
    "Type": "AWS::DynamoDB::Table",
    "Properties": {
        "AttributeDefinitions": [{
            "AttributeName": "token",
            "AttributeType": "S"
        }],
       "KeySchema": [{
           "AttributeName": "token",
           "KeyType": "HASH"
        }],
       "ProvisionedThroughput": {
           "ReadCapacityUnits": 1,
           "WriteCapacityUnits": 1
        },
        "TableName": "auth-token"
    }
},
"AuthLambdaRole": {
    "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": "dynamodb",
            "PolicyDocument": {
                "Version": "2012-10-17",
                "Statement": [{
                    "Sid": "1",
                    "Effect": "Allow",
                    "Action": [
                        "dynamodb:GetItem"
                    ],
                    "Resource": [
                        {"Fn::Join" : ["", ["arn:aws:dynamodb:", {"Ref": "AWS::Region"}, ":", {"Ref": "AWS::AccountId"}, ":table/auth-token"]]}
                    ]
                }]
            }
        }]
    }
},
"AuthLambda": {
    "Type": "AWS::Lambda::Function",
    "Properties": {
        "Code": {
            "ZipFile": {"Fn::Join" : ["\n", [
                "// the code, line by line from above"
            ]]}
        },
        "Handler": "index.handler",
        "MemorySize": 128,
        "Role": {"Fn::GetAtt": ["AuthLambdaRole", "Arn"]},
        "Runtime": "nodejs4.3",
        "Timeout": 60
    }
},

Now we have our Auth Lambda Function. It’s time to integrate with our API Gateway.

The resource that we need is called an Authorizer. The Authorizer defines:

• The Lambda function to execute.
• The header field that is passed to the Lambda function including a RegEx to validate the input value.
• The TTL of the result.

See how this translates to CloudFormation:

"RestApiAuthorizerRole": {
    "Type": "AWS::IAM::Role",
    "Properties": {
        "AssumeRolePolicyDocument": {
            "Version": "2012-10-17",
            "Statement": [{
                "Effect": "Allow",
                "Principal": {
                    "Service": ["apigateway.amazonaws.com"]
                },
                "Action": ["sts:AssumeRole"]
            }]
        },
        "Policies": [{
            "PolicyName": "lambda",
            "PolicyDocument": {
                "Version": "2012-10-17",
                "Statement": [{
                    "Effect": "Allow",
                    "Action": "lambda:invokeFunction",
                    "Resource": {"Fn::GetAtt": ["AuthLambda", "Arn"]}
                }]
            }
        }]
    }
},
"RestApiAuthorizer": {
    "Type": "AWS::ApiGateway::Authorizer",
    "Properties": {
        "AuthorizerCredentials": {"Fn::GetAtt": ["RestApiAuthorizerRole", "Arn"]},
        "AuthorizerResultTtlInSeconds": 300,
        "AuthorizerUri": {"Fn::Join" : ["", ["arn:aws:apigateway:", {"Ref": "AWS::Region"}, ":lambda:path/2015-03-31/functions/", {"Fn::GetAtt": ["AuthLambda", "Arn"]}, "/invocations"]]},
        "IdentitySource": "method.request.header.Authorization",
        "IdentityValidationExpression": "^[a-zA-Z0-9]{3,32}$",
        "Name": "api-authorizer",
        "RestApiId": {"Ref": "RestApi"},
        "Type": "TOKEN"
    }
}

Now we need to define an HTTP Resource with an HTTP method protected by our Custom Authorization.

The PUT /user Endpoint

We will describe a PUT /user endpoint that takes JSON as input in the following form:

{
    "id": "xyz",
    "name": "xyz"
}

To make all this work we need to setup a Lambda function with the permission to be invoked by our API Gateway. We also need a model for our user JSON structure, the /user resource itself and the PUT method.

"Lambda": {
    "Type": "AWS::Lambda::Function",
    "Properties": {
        [...]
    }
},
"LambdaPermission": {
    "Type": "AWS::Lambda::Permission",
    "Properties": {
        "Action": "lambda:invokeFunction",
        "FunctionName": {"Fn::GetAtt": ["Lambda", "Arn"]},
        "Principal": "apigateway.amazonaws.com",
        "SourceArn": {"Fn::Join": ["", ["arn:aws:execute-api:", {"Ref": "AWS::Region"}, ":", {"Ref": "AWS::AccountId"}, ":", {"Ref": "RestApi"}, "/*"]]}
    }
},

"Model": {
    "Type": "AWS::ApiGateway::Model",
    "Properties": {
        "ContentType": "application/json",
        "Name": "user",
        "RestApiId": {"Ref": "RestApi"},
        "Schema": {
            "$schema": "http://json-schema.org/draft-04/schema#",
            "title": "user",
            "type": "object",
            "additionalProperties": false,
            "properties": {
                "id": {
                    "type": "string"
                },
                "name": {
                    "type": "string"
                }
            },
            "required": ["id", "name"]
        }
    }
}
"Resource": {
    "Type": "AWS::ApiGateway::Resource",
    "Properties": {
        "RestApiId": {"Ref": "RestApi"},
        "ParentId": {"Fn::GetAtt": ["RestApi", "RootResourceId"]},
        "PathPart": "user"
    }
},
"Put": {
    "Type": "AWS::ApiGateway::Method",
    "DependsOn": "LambdaPermission",
    "Properties": {
        "AuthorizationType": "CUSTOM",
        "AuthorizerId": {"Ref": "RestApiAuthorizer"},
        "RestApiId": {"Ref": "RestApi"},
        "ResourceId": {"Ref": "Resource"},
        "HttpMethod": "PUT",
        "Integration": {
            "Type": "AWS",
            "IntegrationHttpMethod": "POST",
            "Uri": {"Fn::Join" : ["", ["arn:aws:apigateway:", {"Ref": "AWS::Region"}, ":lambda:path/2015-03-31/functions/", {"Fn::GetAtt": ["Lambda", "Arn"]}, "/invocations"]]},
            "IntegrationResponses": [{
                "StatusCode": 204
            }],
            "RequestTemplates": {
                "application/json": "{\"action\": \"put\", \"body\": $input.json('$')}"
            }
        },
        "RequestModels": {
            "application/json": {"Ref": "Model"}
        },
        "RequestParameters": {
          "method.request.header.Authorization": true
        },
        "MethodResponses": [{
            "ResponseModels": {
                "application/json": "Empty"
            },
            "StatusCode": 204
        }]
    }
}

Now you can use a Lambda function to protect your REST API built with API Gateway.

With API Gateway you can configure a RESTful API. Authorizers can be used to implement Custom Authorization with a Lambda function. The API Gateway will invoke the Auth Lambda Function to check if an HTTP request is allowed. You can use DynamoDB or other databases to store the necessary auth information. As long as the logic can run inside a Lambda function, you can do whatever you want.