Using the AWS Parameters and Secrets Lambda Extension

How to use the newly released AWS Parameters and Secrets Lambda Extension in Python

Posted by Steven Tan on 27th October 2022

What is the new AWS Parameters and Secrets Lambda Extension?

The newly released AWS Parameters and Secrets Lambda Extension is an AWS managed lambda layer that allows you to easily call a localhost HTTP API to retrieve Secrets Manager and Parameter Store. This allows for simple retrieval of secrets for use within your lambdas without having to rely on the language specific aws-sdk libraries.

The local HTTP API is exposed on port 2773 by default (configurable via the PARAMETERS_SECRETS_EXTENSION_HTTP_PORT environment variable). Other options related to the extension can configured via other environment variables outlined in the AWS documentation.

To add a level of security to the extension, you will need to pass the environment variable AWS_SESSION_TOKEN as a value to the X-Aws-Parameters-Secrets-Token HTTP header. This is to mitigate any unwanted calls to the lambda extension HTTP API from calls originating outside of the lambda environment.

Why should I use this over the standard AWS SDKs?

When writing your Python code using a library such as boto3, you will need to import the entire library and then make a new HTTP call to the AWS API every time you call the client.get_parameter or client.get_secret_value methods.

The Lambda extension improves your developer / operational experience by removing the need to import the entire AWS SDK libraries when needing to retrieve your secrets. To use the extension, you simply need to make a HTTP call to the localhost API and retrieve the value. When making the HTTP call, the lambda extension also caches the secrets for you (which defaults at 5 minutes and can be configured via the SECRETS_MANAGER_TTL environment variable), helping to reduce the number of calls made to the AWS API and therefore also reducing the overall cost of your application. This also means you do not need to build your own secrets caching mechanism.

Lambda Extensions API Flow

From a single execution of a lambda, I noticed that the Max Memory Used when using the extension was 61MB compared to using boto3 which was at 90MB. Although this figure doesn't seem like much despite a 30% decrease in memory consumption, it can start to add up when using other libraries which may tip you over the edge of the 128MB memory limit and into the next tier.

Code Sample

To use the lambda extension, I've written some very simple helper functions to retrieve the secrets without importing any additional requirements like the requests library.

import json
import os

from urllib.parse import urlencode, quote_plus
import urllib3

from typing import Optional

http = urllib3.PoolManager()
PARAMETERS_SECRETS_EXTENSION_HTTP_PORT = os.environ.get(
    "PARAMETERS_SECRETS_EXTENSION_HTTP_PORT", 2773
)

def get_ssm_paramater(
    parameter_name: str,
    version: Optional[str] = "1",
    label: Optional[str] = None,
    decrypt: Optional[bool] = False,
):
    """
    Retrieve a parameter from the AWS Systems Manager Service

    Params:
        parameter_name (str): The name of the parameter to retrieve or it's AWS ARN
        version (str): The version of the parameter to retrieve
        label (str): The label of the parameter to retrieve
        decrypt (bool): Whether to decrypt the parameter value or not

    Returns:
        The value stored in SSM Parameter Store
    """        
    parameters = {"name": parameter_name, "withDecryption": decrypt}
    if version:
        parameters["version"] = version
    elif label:
        parameters["label"] = label
    url_root = f"http://localhost:{PARAMETERS_SECRETS_EXTENSION_HTTP_PORT}/systemsmanager/parameters/get/"
    r = http.request(
        "GET",
        f"{url_root}?{urlencode(parameters, quote_via=quote_plus)}",
        headers={"X-Aws-Parameters-Secrets-Token": os.environ.get("AWS_SESSION_TOKEN")},
    )
    return json.loads(r.data.decode("utf-8"))

def get_secret(
    secret_id: str,
    version_id: Optional[str] = None,
    version_stage: Optional[str] = None,
):
    """
    Retrieve a secret from the AWS Secrets Manager Service

    Params:
        secret_id: The name of the secret to retrieve or it's AWS ARN
        version_id: The version of the secret to retrieve
        version_stage: The version stage of the secret to retrieve

    Returns:
        The value stored in AWS Secrets Manager
    """
    parameters = {
        "secretId": secret_id,
    }
    if version_id and version_stage:
        raise Exception("Cannot specify both version and stage")
    if version_id:
        parameters["versionId"] = version_id
    elif version_stage:
        parameters["versionStage"] = version_stage
    url_root = (
        f"http://localhost:{PARAMETERS_SECRETS_EXTENSION_HTTP_PORT}/secretsmanager/get"
    )
    r = http.request(
        "GET",
        f"{url_root}?{urlencode(parameters, quote_via=quote_plus)}",
        headers={"X-Aws-Parameters-Secrets-Token": os.environ.get("AWS_SESSION_TOKEN")},
    )
    return json.loads(r.data.decode("utf-8"))

## Example Usage for SSM 
result = get_ssm_paramater("/sktan/testsecret1", decrypt=True)
## Example Usage for Secrets Manager
result = get_secret("testsecret")

Now since this will only work within the lambda environment, you will need to add some logic to your code to determine the execution environment. This can be done by checking for the following environment variables LAMBDA_TASK_ROOT and AWS_EXECUTION_ENV, which indicates that you are running within a Lambda function. If these environment variables are not found, you can assume that the code is running within a local or CI/CD environment, and should use the boto3 retrieval methods instead. Example Code might look similar to the following:

IS_LAMBDA_ENVIRONMENT = True
if not os.getenv("AWS_EXECUTION_ENV") and not os.getenv("LAMBDA_TASK_ROOT"):
    import boto3
    IS_LAMBDA_ENVIRONMENT = False

def lambda_handler(event, context):
    if IS_LAMBDA_ENVIRONMENT:
        # Use the lambda extension to retrieve secrets
        result = get_secret("testsecret")
    else:
        # Use boto3 to retrieve secrets
        client = boto3.client("secretsmanager")
        result = client.get_secret_value(SecretId="testsecret")

Performance Comparison

To briefly look at what the performance difference of making the API calls to the lambda extension vs boto3, I ran a simple test to retrieve values 50 times using each method and timed the amount of time it takes for the retrieval to complete.

The first set of results are for retrieving a single value for the first time. From what we can see, it looks like the performance varies for the initial call. For the SSM Parameter (Extension) value though, I'm assuming that this is caused by the initial call which probably includes a cold-start for the HTTP extension lambda server as this was the first ever call made by the lambda function.

Secret Retrieval Method Time (ms)
SSM Parameter (Extension) 650ms
SSM (Boto3) 250ms
Secrets Manager (Extension) 101ms
Secrets Manager (Boto3) 213ms

The second set of results are for the overall figures gathered during the test.

Secret Retrieval Method Statistic Time (ms)
SSM Parameter (Extension) AVG 27ms
SSM Parameter (Extension) P90 20ms
SSM Parameter (Boto3) AVG 50ms
SSM Parameter (Boto3) P90 60ms
Secrets Manager (Extension) AVG 17ms
Secrets Manager (Extension) P90 20ms
Secrets Manager (Boto3) AVG 31ms
Secrets Manager (Boto3) P90 40ms

From the results, we can see that overall, it is better to use the Lambda extension for retrieving your secrets from either the SSM Parameter Store or Secrets Manager. In general, there seems to be a 2x performance increase between the 2 methods due to how the extension is caching results.

Performance Testing Code

For the sake of transparency on how I came to the conclusions above, I've also included the code I used to test the performance of the lambda extension vs boto3 and raw results in CSV format.

import time

def get_ssm_parameter_via_boto3(parameter_name: str, decrypt=True):
    try:
        return ssm_client.get_parameter(Name=parameter_name, WithDecryption=decrypt)
    except:
        return None

def get_secret_via_boto3(secret_name: str):
    try:
        return secrets_client.get_secret_value(SecretId=secret_name)
    except:
        return None

def lambda_handler(event, context):
    print("=== TESTING PERFORMANCE ===")

    for i in range(__ATTEMPTS__):
        start = time.perf_counter()
        result = get_ssm_paramater_via_extension("/sktan/testsecret1", decrypt=True)
        end = time.perf_counter() - start
        success = result is not None
        print(f"ssm,{i+1},{end},{success}")

    for i in range(__ATTEMPTS__):
        start = time.perf_counter()
        result = get_ssm_parameter_via_boto3("/sktan/testsecret1", decrypt=True)
        end = time.perf_counter() - start
        success = result is not None
        print(f"ssm,{i+1},{end},{success}")

    for i in range(__ATTEMPTS__):
        start = time.perf_counter()
        result = get_secret_via_extension("testsecret")
        end = time.perf_counter() - start
        success = result is not None
        print(f"secrets,{i+1},{end},{success}")

    for i in range(__ATTEMPTS__):
        start = time.perf_counter()
        result = get_secret_via_boto3("testsecret")
        end = time.perf_counter() - start
        success = result is not None
        print(f"secrets,{i+1},{end},{success}")

Final Thoughts

For applications that require low latency and high throughput such as web APIs, the AWS Parameters and Secrets Lambda Extension is a great addition to the arsenal of tools you can use. In return for a bit of extra setup, you can get a significant performance boost and reducing your overall costs by taking advantage of their caching mechanism.

For Lambda Functions that do not run very often, the benefits offered might not be super significant. Though if you already have started to implement this both from an Infrastructure as Code and application code perspective, I would still use the extension for all Lambda functions to standardise your overall codebase.