Amazon WebSocket API Gateway authoriser in Rust

In today’s technology landscape, high-performance, real-time communication has become a requirement for many applications. While ultra-low latency solutions exist for niche use cases, WebSockets have emerged as an attractive option for most real-time communication needs. However, implementing and scaling a WebSocket-based service can be significantly more complex compared to traditional request-response based HTTP APIs.

WebSocket API in Amazon API Gateway is a fully managed service that greatly simplifies the complexity of implementing and scaling WebSocket-based services. By seamlessly integrating with AWS Lambda, it provides a serverless solution for real-time communication that can effortlessly scale to handle millions of concurrent connections.

A key feature of API Gateway is built-in authentication and authorisation. However, while REST APIs and HTTP APIs in API Gateway offer built-in JWT and Amazon Cognito authorisers, WebSocket APIs have a more limited set of options. As outlined in the API Gateway documentation, WebSocket APIs only support two types of authorisers: IAM authorisation and Lambda request authorisers.

When using a Lambda authoriser, a Lambda function is triggered whenever a new connection is established. This function has access to all the connection details, including headers and query parameters. Although Lambda authorisers require a bit more effort to set up, they provide the most flexible option by allowing you to implement custom authorisation logic within the Lambda function, tailored to your application’s specific requirements.

In this article, we’ll demonstrate how to create a WebSocket request authoriser using Rust that integrates with JWT tokens issued by Amazon Cognito.

However, the ultimate aim of this article extends beyond the specific implementation. By walking through this exercise, we hope to provide you with a sense of the strengths of the Rust ecosystem for serverless development on AWS in a broader context (if you aren’t already persuaded).

The code for this article is on GitHub.

Project outline

In this article, we will implement two Lambda functions using Rust. The first is a simple WebSocket message handler that logs incoming messages. The more interesting function is the authoriser function, which accepts a JWT token provided as a query parameter. This authoriser function decodes and verifies the token against the JSON Web Keys obtained from the authorisation server (Cognito in our case).

The two Lambda functions are implemented as separate Rust binaries, leveraging a Cargo Workspace to maintain a clear separation between the two crates. The Rust ecosystem provides excellent support for AWS Lambda development. We utilise Cargo Lambda for seamless compilation and take advantage of crates such as lambda_runtime and aws_lambda_events to streamline the implementation process. These tools and libraries significantly simplify the development workflow and enable efficient integration with the AWS Lambda platform.

In addition to the Lambda function implementations, the GitHub repository includes an infrastructure folder housing the AWS Cloud Development Kit (CDK) code responsible for deploying the Lambda functions and the WebSocket API Gateway. While we won’t delve into the intricacies of the CDK code within the scope of this article, it serves as a valuable reference for those interested in leveraging CDK for their deployment needs. The CDK code in the repository provides a practical example of how to define and manage the necessary AWS resources using infrastructure as code principles.

The dummy message handler

Although the dummy message handler doesn’t do much, it’s worthwhile going through it as it contains the fundamental building blocks of a Lambda function.

use aws_lambda_events::apigw::ApiGatewayWebsocketProxyRequest;
use lambda_runtime::tracing::Level;
use lambda_runtime::{service_fn, Error, LambdaEvent};
use serde::Serialize;
use tracing::info;

#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt()
        .with_max_level(Level::INFO)
        .json()
        .init();
    let func = service_fn(func);
    lambda_runtime::run(func).await?;
    Ok(())
}

#[derive(Debug, Serialize)]
struct Response {
    #[serde(rename = "statusCode")]
    status_code: i32,
}

async fn func(
    event: LambdaEvent<ApiGatewayWebsocketProxyRequest>,
) -> Result<Response, Error> {
    let (event, _context) = event.into_parts();
    info!(event = debug(event), "handling event");

    Ok(Response { status_code: 200 })
}

The main function follows the standard structure of a Tokio runtime’s entry point. Before invoking the lambda_runtime::run function to start the Lambda runtime, the main function provides an opportunity to perform any necessary setup tasks. This can include configuring the tracing subscriber, initialising dependencies, or applying global settings. By utilising the main function for setup, we can ensure that the Lambda function is properly initialised and configured before it begins processing events.

The handler function passed to the runtime takes a LambdaEvent (a type provided by lambda_runtime), which can wrap an arbitrary serde Deserialize type. The return type must be a Result with the Ok variant wrapping a serde Serialize type.

The statusCode here isn’t used, but it demonstrates the use of serde types for responses.

The authoriser Lambda

The authoriser Lambda function retrieves configuration parameters from the environment variables. These parameters are then utilised to validate the provided JWT token. Based on the token’s claims and the verification result, the Lambda function generates an IAM policy response. This response determines the permissions and access rights associated with the authenticated user.

To gain a better understanding of the authoriser Lambda’s implementation, let’s examine each module in detail in a bottom up manner.

Loading the configuration

To successfully decode and verify the JWT token, we require two crucial pieces of information. The first is the URL that provides access to the JSON Web Key Set (JWKS) associated with the Cognito user pool. The JWKS contains the necessary key(s) to verify the integrity and authenticity of the JWT. The second piece of information is the expected audience for which the token should have been issued. In the context of Cognito, the audience corresponds to the ID of the user pool client.

Environment variables offer a flexible and widely-used mechanism for configuring Lambda functions with dynamic parameters. However, retrieving environment variables individually using the std::env module can be cumbersome and lead to repetitive code.

The serde_env crate provides a convenient and efficient way to deserialise environment variables directly into structured data types, eliminating the need for manual variable retrieval and parsing.

use serde::Deserialize;
use std::sync::OnceLock;

#[derive(Clone, Debug, Deserialize)]
pub struct Config {
    pub jwks_url: String,
    pub audience: String,
}

static CONFIG: OnceLock<Config> = OnceLock::new();

pub fn get_config() -> &'static Config {
    CONFIG.get_or_init(|| serde_env::from_env::<Config>().expect("config to load fine"))
}

Decoding the token

The decode.rs module plays a crucial role in the authentication process by handling the complex task of decoding and verifying the JWT token. It leverages the functionality provided by the jsonwebtoken crate to ensure the token’s validity and extract relevant information. The decoding and verification logic itself is not specific to AWS API Gateways, making it reusable in other contexts. However, one aspect that is particular to Cognito is the extraction of the username from the cognito:username field within the token’s payload.

use anyhow::{bail, Result};
use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Validation};
use serde::Deserialize;
use serde_json::Value;
use tokio::sync::OnceCell;
use tracing::info;

use crate::config::get_config;

static CACHED_KEYS: OnceCell<Vec<Jwk>> = OnceCell::const_new();

#[derive(Debug, Deserialize)]
pub struct Jwk {
    pub kid: String,
    pub e: String,
    pub n: String,
}

#[derive(Debug, Deserialize)]
pub struct Claims {
    #[serde(rename = "cognito:username")]
    pub username: String,
}

pub async fn verify_claims(token: &str) -> Result<Claims> {
    let keys = keys().await?;

    let header = decode_header(token)?;
    let kid = match header.kid {
        Some(k) => k,
        None => bail!("token header has no kid"),
    };
    let key = match keys.iter().find(|&k| k.kid == kid) {
        Some(key) => key,
        None => bail!("none of the keys match token kid"),
    };

    info!(key = debug(key), "found appropriate key");

    let mut validation = Validation::new(Algorithm::RS256);
    let audience = &get_config().audience;
    validation.set_audience(&[audience]);

    let token_data = decode::<Claims>(
        token,
        &DecodingKey::from_rsa_components(&key.n, &key.e)?,
        &validation,
    )?;

    Ok(token_data.claims)
}

pub async fn keys() -> Result<&'static Vec<Jwk>> {
    CACHED_KEYS.get_or_try_init(fetch_keys).await
}

async fn fetch_keys() -> Result<Vec<Jwk>> {
    let url = &get_config().jwks_url;
    info!(url, "fetching jwks");
    let client = reqwest::Client::builder().use_rustls_tls().build()?;
    let res = client.get(url).send().await?;

    let jwk_text = res.text().await?;

    let keys_value = match serde_json::from_str::<Value>(&jwk_text)? {
        Value::Object(mut obj) => match obj.get_mut("keys") {
            Some(val) => val.take(),
            None => bail!("no keys found in JWK JSON"),
        },
        _ => bail!("JWK is not a mapping for keys"),
    };

    let keys: Vec<Jwk> = serde_json::from_value(keys_value)?;
    Ok(keys)
}

The JWK keys are fetched using the reqwest crate from the endpoint provided by Cognito, which is specific to the user pool. Since the keys remain relatively stable and do not change frequently, it is acceptable to cache them within the Lambda function to optimise performance and reduce unnecessary network calls.

The verify_claims function utilises the retrieved keys to validate the integrity and authenticity of the token. Additionally, it performs crucial checks to ensure that the token has not expired and that it was issued for the intended audience. In this implementation, we focus on extracting the Cognito username from the token’s claims. However, with minor modifications to the code, we can easily extract any other relevant field that Cognito includes in the token, enabling flexibility and extensibility based on the specific requirements of the application.

Generating an authoriser response

API Gateway expects the authoriser to return an IAM policy document that specifies the actions the user is permitted to perform. In our implementation, we will construct a straightforward policy that includes either a blanket Allow or Deny statement, depending on the outcome of the token verification process.

The response object returned by the authoriser consists of two key components: the IAM policy document and a custom context. The custom context is derived from the claims extracted from the JWT during the verification process. This additional contextual information can be utilised by the API Gateway and downstream services to make informed decisions based on the user’s attributes and permissions.

use aws_lambda_events::apigw::ApiGatewayV2CustomAuthorizerIamPolicyResponse;
use aws_lambda_events::event::apigw::ApiGatewayCustomAuthorizerPolicy;
use aws_lambda_events::event::iam::IamPolicyStatement;
use aws_lambda_events::iam::IamPolicyEffect;
use serde::{Deserialize, Serialize};
use tracing::error;

use crate::decode::verify_claims;

#[derive(Debug, Deserialize, Serialize)]
pub struct Context {
    username: String,
}

pub async fn authorise(token: &str) -> ApiGatewayV2CustomAuthorizerIamPolicyResponse<Context> {
    match verify_claims(token).await {
        Ok(claims) => {
            let context = Context {
                username: claims.username,
            };

            ApiGatewayV2CustomAuthorizerIamPolicyResponse {
                principal_id: None,
                policy_document: generate_policy(IamPolicyEffect::Allow),
                context,
            }
        }
        Err(err) => {
            error!(error = debug(err), "failed to authenticate connection");
            generate_deny_response()
        }
    }
}

pub fn generate_deny_response() -> ApiGatewayV2CustomAuthorizerIamPolicyResponse<Context> {
    let context = Context {
        username: "".to_string(),
    };

    ApiGatewayV2CustomAuthorizerIamPolicyResponse {
        principal_id: None,
        policy_document: generate_policy(IamPolicyEffect::Deny),
        context,
    }
}

fn generate_policy(effect: IamPolicyEffect) -> ApiGatewayCustomAuthorizerPolicy {
    let statement = IamPolicyStatement {
        action: vec!["execute-api:Invoke".to_string()],
        effect,
        resource: vec!["*".to_string()],
        condition: None,
    };
    ApiGatewayCustomAuthorizerPolicy {
        version: Some("2012-10-17".to_owned()),
        statement: vec![statement],
    }
}

The authorise function accepts a token as input and leverages the verify_claims function, which we previously implemented, to validate the token’s authenticity and integrity. Upon successful verification, the function maps the extracted claims to an appropriate response object. In the event of any errors during the verification process, the function generates a Deny response, indicating that the user is not authorised to perform the requested actions.

The implementation of the authorise function takes advantage of the types provided by the aws_lambda_events crate to construct a well-formed IAM policy response. The strong type safety offered by Rust, combined with the pre-defined types from the crate, proves invaluable in this context. It significantly reduces the time and effort required to decipher the often sparse and ambiguous documentation provided by AWS regarding the expected request and response schemas for API Gateway authorisers. By leveraging these types, we can ensure that our implementation adheres to the required format and structure, minimising the chances of errors and simplifying the development process.

The entrypoint

The outer structure of the authoriser Lambda is similar to the message handler’s.

23 collapsed lines
mod auth;
mod config;
mod decode;

use aws_lambda_events::apigw::{
    ApiGatewayV2CustomAuthorizerIamPolicyResponse, ApiGatewayWebsocketProxyRequest,
};
use lambda_runtime::tracing::Level;
use lambda_runtime::{service_fn, Error, LambdaEvent};
use tracing::error;

use crate::auth::{authorise, generate_deny_response, Context};

#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt()
        .with_max_level(Level::INFO)
        .json()
        .init();
    let func = service_fn(func);
    lambda_runtime::run(func).await?;
    Ok(())
}

async fn func(
    event: LambdaEvent<ApiGatewayWebsocketProxyRequest>,
) -> Result<ApiGatewayV2CustomAuthorizerIamPolicyResponse<Context>, Error> {
    let (event, _context) = event.into_parts();
    let response = match event.query_string_parameters.first("auth") {
        None => {
            error!("missing auth token in connection request");
            generate_deny_response()
        }
        Some(token) => authorise(token).await,
    };

    Ok(response)
}

The ApiGatewayWebsocketProxyRequest type provides a straightforward way to extract the auth query parameter. This parameter contains the token, which we process using the authorise function we implemented.

The infrastructure code

The infrastructure for this project is defined using AWS Cloud Development Kit (CDK). At the time of writing, CDK does not offer native support for Rust, so the stack is implemented using TypeScript. While providing an exhaustive overview of the entire stack is beyond the scope of this post, there are a few noteworthy points worth highlighting.

Rust functions utilising lambda_runtime run perfectly fine on Amazon Linux runtimes, specifically provided.al2 and provided.al2023. Cargo Lambda handles the cross-compilation process seamlessly, allowing the use of either ARM64 or x86_64 architectures. However, it is crucial to ensure that the architecture specified in Cargo Lambda aligns with the runtime configuration of the Lambda function.

Prior to deploying the stack, it is necessary to recompile the binaries with the latest codebase. While this step could be performed separately, I have opted to invoke Cargo Lambda as a subprocess within the CDK code itself. This approach guarantees that the deployed code is always up to date, eliminating the need for manual intervention.

CDK can be reasonably fast at uploading new code, but for an even faster feedback loop, you may want to leverage the built-in deployment functionality of Cargo Lambda during development. By utilising this feature, the deployment time can be significantly reduced to mere seconds, including the compilation process. This rapid iteration cycle greatly enhances developer productivity and enables faster testing and refinement of the application.

With the code deployed, you can test the WebSocket connection using wscat.

wscat -c "wss://$API_URL/default?auth=$TOKEN"

Conclusion

Rust proves to be a brilliant choice for developing Lambda functions, and its outstanding performance, particularly in terms of minimal cold start times, makes it an ideal fit for building real-time WebSocket APIs. The language’s inherent efficiency and low-level control enable developers to create highly optimised and responsive serverless applications.

Throughout this post, we have delved into the implementation of a JWT authoriser function specifically designed to integrate with Amazon Cognito for use in Amazon WebSocket APIs. This fills a gap in the native functionality of these APIs, which do not provide built-in support for Cognito or JWT authorisation out of the box. By leveraging Rust’s capabilities, we have demonstrated how to overcome this limitation and create a secure and efficient authorisation mechanism.

However, beyond the specific implementation details, the broader aim of this article is to showcase the remarkable strengths of the Rust ecosystem in the context of serverless API development on AWS. Rust’s robustness, performance, and growing set of libraries and frameworks tailored for serverless environments make it a compelling choice for building scalable and reliable APIs. By harnessing the power of Rust, developers can unlock new possibilities and deliver high-quality, performant serverless applications on the AWS platform.