Exposing AWS KMS Asymmetric Keys as a JWKS

Brian Maloney
Benchling Engineering
7 min readFeb 2, 2023

--

Here at Benchling, interaction with services is a large part of our business, from employees interacting with the software-as-a-service products with which we conduct our daily business, all the way down to interactions between the services that make up the Benchling application platform itself. Secure authentication and authorization to services is a long-standing issue in the industry, but one that has been improving in recent years due to the widespread adoption of modern standards such as OAuth 2.0 and OpenID Connect (OIDC).

One specific use case for service-to-service authentication that is important to Benchling Security is connecting our Threat Detection Pipeline to our enterprise identity services vendor. We use this connection to connect log and other data provided by the vendor to our centralized Threat Detection Platform, where we correlate this with other sources of intelligence to detect risky or suspicious user activity in near real-time.

Modern Authentication with OIDC

Our specific identity services vendor offers two options for authenticating to its API: either an API token that an administrator can generate, or interaction by acting as an Application. API tokens, while very easy to use, are a poor choice for two reasons:

  1. First, they are a static secret that must be handled carefully and rotated frequently to mitigate the risk of a leaked key, which causes significant management overhead.
  2. Second, the identity services vendor links the privileges and identity of an API token inextricably to the administrator who generated it. This causes actions using the key to be attributed to the administrator and also makes it impossible to implement the principle of least privilege.

Client authentication when acting as an Application allows the use of OIDC, and this vendor specifically requires the use of the private_key_jwt Client Authentication method. Enforcing this requirement is a good choice on the part of the vendor — by using public-key encryption, no secrets need to be shared, only public keys need to be exchanged, and neither party can impersonate the other.

At this point you may be thinking, “Even though secrets don’t need to be exchanged, isn’t there still overhead for rotating the public key? And what is the best way to manage and safeguard the private keys in a modern cloud environment?” These are legitimate concerns and both have relatively simple solutions.

For managing private keys in a cloud environment, AWS was kind enough to solve this for us when they added asymmetric key functionality to KMS in 2019. The KMS asymmetric key functionality allows you to provision public/private keypairs using the same cloud infrastructure tooling you’re already using. The private key never leaves AWS infrastructure, but the public key can be exported and shared. Signing and verification operations must therefore be done using AWS APIs rather than the traditional SSL toolkits, however this functionality is readily available via existing SDKs and command-line tools.

Client token flow using private_key_jwt with KMS Asymmetric Keys

The Challenge: Public Key Rotation

The first issue described above is a bit more complex, and isn’t solved for us by our cloud provider. The KMS Asymmetric key must still be rotated for maximum security, but that means that the new public key needs to be shared with the identity provider. The OIDC Discovery standard provides a common method for discovery of information about OpenID Providers using standard web technologies, including providing a URL to a JSON Web Key Set (JWKS defined in RFC7517) containing the public keys for the OpenID Provider. Our identity services vendor supports this approach — when creating a new API Services Application, public keys can either be uploaded into the provider or a JWKS URL can be provided. Because our vendor caches the JWKS, it is only loaded infrequently to update the cache, or when a new key is used.

Given the flexibility gained from dynamic JWKS sharing, we strongly preferred this approach for our integration. With the design selected, implementation can be broken into two major steps:

  • Enable signing of JWTs using AWS KMS Asymmetric Keys in the client
  • Dynamically generate and serve the JWKS for the relevant KMS Keys

Signing of JWTs using KMS Asymmetric Keys is a common use case — there are already multiple examples of how to do this available (Python, Node.js). Integrating signing into your client workflow is fairly straightforward following one of these examples, so we won’t dig deeper into that in this article.

JWKS Construction and Serving

While this isn’t that complex of a task, there are numerous ways to accomplish it with different capabilities and performance characteristics. This article will cover the way we tackled the problem to meet Benchling Security’s specific needs without being prescriptive. Although for this post we are using a very simple design, there are still some challenges to overcome on the way to a functioning solution.

The basic design of our JWKS service is a Lambda function, fronted by the relatively new Function URLs feature of AWS Lambda. Function URLs allow a single-function microservice (like our JWKS exporter) to be served without the additional infrastructure of an API Gateway. This greatly reduces the amount of infrastructure we have to build to provide this service.

JKWS Lambda Flow

Key Selection

In any AWS region, there may be multiple KMS Asymmetric Keys defined, but we don’t need to export every key in the region in this JWKS. AWS has some great ways to group, select, and filter cloud resources, tags being the best example. Our initial expectation was to tag each of the keys for each use case and construct the JWKS from only those keys. Unfortunately, the response from the AWS API’s ListKeys function response does not include the tags defined on the keys. The only way to filter the list of keys based on tags would be to iterate over each key in the region, calling ListResourceTags on each key, which isn’t scalable for a function that needs to return within a few seconds.

Because of these limitations, the only reasonable approach is to use key aliases to identify the keys to be exported in the JWKS. Using aliases allows the use of the ListAliases API function and filter the results to just the keys which should be exported.

Rendering Public Keys as JWKs

Now that we have a list of Asymmetric key resources within KMS that we want to export, the public keys need to be converted into JWK structures which can then be combined into a JWKS. Because simplicity is a design goal, we want to use the fewest possible steps to convert the public key returned by the AWS GetPublicKey API into a JWK. There are numerous Python libraries available that implement the JOSE (Javascript Object Signing and Encryption) standards, including JWK. Unfortunately, many of the popular options on PyPI have limited JWK functionality. For example, some libraries can only handle JWKs which are already in JWK format but cannot convert an existing public key or certificate to a JWK. Fortunately, the Python ecosystem is large and JWCrypto has a full suite of JWK-handling functions, including conversion from PEM.

The only remaining piece of the puzzle is conversion from the DER key delivered by GetPublicKey into a PEM formatted public key. While this is a simple operation to do by hand, this functionality is provided by the very popular python Cryptography library, which is already used by JWCrypto.

Finally, all that’s needed is to gather the keys into a JWKS structure and render that as JSON, and write that as the body of your Lambda response.

Putting it all together

Now that we know how to do all the individual steps, we can assemble them into a pleasingly simple Python Lambda with only 3 functions:

import os
import logging
import boto3
from cryptography.hazmat.primitives import serialization
from jwcrypto import jwkkms = boto3.client('kms')for var in ["ALIAS_START"]:
if not var in os.environ:
logging.critical(f'{var} environment variable not set')
exit()
# Search for and return the list of enabled keys with aliases
# that start with our desited string
def find_enabled_keys(starts_with):
alias_paginator = kms.get_paginator('list_aliases')
alias_iterator = alias_paginator.paginate().search(f'Aliases[?TargetKeyId != null && starts_with(AliasName, `{starts_with}`)].TargetKeyId')
key_ids = [page for page in alias_iterator] enabled_keys = [
key_id
for key_id in key_ids
if kms.describe_key(KeyId=key_id)['KeyMetadata']['KeyState']
== 'Enabled'
]
return enabled_keys# Return a single key, formatted as a JWCrypto JWK object
def get_jwk(key_id):
response = kms.get_public_key(KeyId=key_id)
pubkey = serialization.\
load_der_public_key(response['PublicKey'])
pub_pem = pubkey.public_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PublicFormat.SubjectPublicKeyInfo
)
key = jwk.JWK.from_pem(pub_pem)
key.update(use=response['KeyUsage'][0:3].lower(), kid=key_id)
return keydef lambda_handler(event, context):
# Construct JWKS structure from keys that match ALIAS_START
jwks = {
"keys": [
get_jwk(key_id)
for key_id
in find_enabled_keys(os.environ['ALIAS_START'])
]
}
# Construct Lambda return structure
return({
'statusCode': 200,
'headers': {
"Content-Type": "application/json",
},
'body': jwks
})

Enhancements for Production Use

The code presented in this article should be considered a poof-of-concept only — if you intend to use this pattern in production, you should consider your use case. If your needs require frequent reloading of this JWKS, it would be more scalable to write this out to an object store when changes occur, and serve the JWKS directly from the object store or a CDN. If you do not have production-level requirements, we still recommend building additional resilience into the function by expanding error handling.

--

--