AWS Custom API Gateway Authorizer with Rust and Okta
Tags: aws, lambda, okta, rust
Published: 2023-11-19 17:05:41

# Introduction

[The code for this blog post is here](https://github.com/szymon-szym/rust-okta-authorizer)

One of the main developers' responsibilities is to ensure that the application we are creating is secured. It also happened to be one of the trickiest things to implement. I am not a security specialist by any means, but I try to keep a proper mindset. 
 
For my site project, I am currently experimenting with different ways to authenticate and authorize users. I need to make sure, that only authorized users can access resources related to their accounts. For my REST APIs, I've decided to go with JWT token.

I've created a custom authorizer for AWS API Gateway and I would like to share my learnings.

# Goal

I want to create a dummy endpoint, which is unique for every user: `hello/{userID}`, and every time it is called, I check if the caller has the right to access this specific resource.

# General architecture

I will use JWT from the request's header to check if the user is authorized to call the given endpoint. To do so, I am going to extract the user's ID from the token and grant access only to the endpoint specific to the given user.

# Project structure

- A lambda function for a dummy endpoint
- A lambda function for the custom authorizer
- Sample Vue application to perform end-to-end user authorization flow and receive a token a user

# Projects set up

## Okta

In this post, I will use Okta as my Auth provider. AWS Cognito is the native AWS service I could use, but at this point, I want to check integration with third-party service inside the custom authorizer.

I created an account on the [Okta developer portal](https://developer.okta.com) with a free plan and added a new application


![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/1s3vnqqmvxpjt252awwl.png)

For the application type, I have OIDC and single-page app. 


![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/glgf9wt3grmi0i0c5en5.png)

We need a user, so I added one manually


![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/ajmu5e1rzhqt366t5j1z.png)

The last step is to click on the user and assign our application to her. 

We should be OK with the Okta setup.

## Sample front-end app

To be able to test my authorizer I need a token generated in the user context. There are probably simpler ways to get tokens, but I just run my own web app and log in using it.

Our FE will only run locally at this point. To keep things in one place, I create a root project folder first with 

```bash
sam init
```

Pick Rust Hello World example. 

I didn't want to spend time creating my own web app, so I just cloned [a sample app created by Okta](https://github.com/okta/samples-js-vue) to the root folder.

I only needed to add an issuer and client ID info to the `testenv` file and I was able to run the app locally. 

Please follow the readme in the repo and cross-check if redirect URIs are set up properly on the Okta portal. 

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/f0yw2h5s275v4tl0dwup.png)

Once everything is prepared, I can run `npm i` to install dependencies, then go to the subfolder and run the app

```bash
cd okta-hosted-login 
npm run serve
```

Now let's log in using the user's account we've created in Okta, and finally - we have user tokens stored in the local store: 


![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/bh5dgs7f6hrpm3oxdiwi.png)

## Cloud infrastructure

I use AWS SAM framework. I need API Gateway and lambda functions. It should be quite straightforward. 

I've already generated the sample lambda function by running `sam init`. This function will be responsible for returning a message from the endpoint. 

For the authorizer, I need to create a separate lambda function by running in the root folder

```bash
cargo lambda new authorizer
```


![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/edh7ez6o6xu5r5yjgvu0.png)


Automatically generated `template.yml` needs a few tweaks. First of all, I need to explicitly define API Gateway, to be able to define authorizer. The Authorizer itself needs to be described as well, and finally, I need to update a path in the main `hello world function`. 
There is also a Dynamo DB table created, which will be used later on for caching public keys.

After all those updates my template looks this way:

```yaml
AWSTemplateFormatVersion: "2010-09-09"
Transform: AWS::Serverless-2016-10-31
Description: >
  Dummy project testing custom authorizers

# More info about Globals: https://github.com/awslabs/serverless-application-model/blob/master/docs/globals.rst
Globals:
  Function:
    Timeout: 3
    MemorySize: 128
    Tracing: Active
  Api:
    TracingEnabled: true

Resources:

  MyApi:
    Type: AWS::Serverless::Api
    Properties:
      StageName: Prod
      Auth:
        DefaultAuthorizer: MyLambdaTokenAuthorizer
        Authorizers:
          MyLambdaTokenAuthorizer:
            FunctionArn: !GetAtt CustomAuthorizer.Arn

  HelloWorldFunction:
    Type: AWS::Serverless::Function # More info about Function Resource: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#awsserverlessfunction
    Metadata:
      BuildMethod: rust-cargolambda # More info about Cargo Lambda: https://github.com/cargo-lambda/cargo-lambda
    Properties:
      CodeUri: ./rust_app # Points to dir of Cargo.toml
      Handler: bootstrap # Do not change, as this is the default executable name produced by Cargo Lambda
      Runtime: provided.al2
      Architectures:
        - x86_64
      Events:
        HelloWorld:
          Type: Api # More info about API Event Source: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#api
          Properties:
            RestApiId: !Ref MyApi
            Path: /hello/{userId}
            Method: get
  CustomAuthorizer:
    Type: AWS::Serverless::Function # More info about Function Resource: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#awsserverlessfunction
    Metadata:
      BuildMethod: rust-cargolambda # More info about Cargo Lambda: https://github.com/cargo-lambda/cargo-lambda
    Properties:
      CodeUri: ./authorizer # Points to dir of Cargo.toml
      Handler: bootstrap # Do not change, as this is the default executable name produced by Cargo Lambda
      Runtime: provided.al2
      Environment:
        Variables:
          KEYS_TABLE_NAME: !Ref KeysTable
          OKTA_KEYS_ENDPOINT: "https://dev-56344269.okta.com/oauth2/default/v1/keys"
      Architectures:
        - x86_64
      Policies:
        - DynamoDBCrudPolicy:
            TableName: !Ref KeysTable


  KeysTable:
    Type: AWS::DynamoDB::Table
    Properties:
      AttributeDefinitions:
      - AttributeName: PK
        AttributeType: S
      KeySchema:
      - AttributeName: PK
        KeyType: HASH
      BillingMode: PAY_PER_REQUEST

``` 

# Authorizer

## Main function

Authorizer for API Gateway is a function that returns an IAM Policy, so API GW can decide if the given request should be allowed or denied. In my case, I would extract the user ID from the JWT token and use it to specify the path to be allowed in the IAM Policy.

I extract logic to a few files to keep things more readable. Besides the `main.rs` I have also files for operations related to JWT, IAM policies, and DynamoDB (the last one is not required, in the second I will explain why I need it).

To be able to work with the token passed by the user, I need a key to decode it, validate it, and extract information. To do so, I don't need any secrets, because the public key is, well, public. I'll just grab keys from Okta endpoint 

`https://<YOUR_DOMAIN>.okta.com/oauth2/default/v1/keys`

The response from the endpoint looks like this:

```json
{
  "keys": [
    {
      "kty": "RSA",
      "alg": "RS256",
      "kid": "t6rj1txgY....",
      "use": "sig",
      "e": "...",
      "n": "3jpalT9ek84-...."
    }
  ]
}
```

I create a struct to hold this response:

```rust
#[derive(Deserialize, Serialize, Debug)]
struct JWTK {
    kid: String,
    kty: String,
    alg: String,
    #[serde(rename = "use")]
    uses: String,
    e: String,
    n: String,
}

#[derive(Deserialize, Serialize, Debug)]
struct JWTKResponse {
    keys: Vec<JWTK>,
}
```

I use `reqwest` library to send http request. The function to do so is straightforward 

```rust
async fn get_keys_from_okta(endpoint: String) -> anyhow::Result<JWTKResponse> {
    let result = reqwest::get(endpoint).await?.json::<JWTKResponse>().await?;
    Ok(result)
}
```

As you might expect, I won't call Okta from the function handler to take advantage of the running hot lambda. I call it from the main function and pass the keys to the handler

```rust
#[derive(Serialize, Deserialize, Debug)]
pub struct StoredKeys {
    keys: HashMap<String, JWTK>,
}
```

There is one optimization I wanted to test. Storing public keys in memory is a must-have, but still, on every cold start, there is a need to call the Okta endpoint. During my tests, I observed that it results in ~700ms overall cold start time for the authorizer. We can do much better if we use e.g. DynamoDB as a caching layer shared among lambdas. This will reduce the cold-start significantly.

In the `dynamo_svc.rs` file I put simple logic to get keys stored in Dynamo

```rust
// ...
pub(crate) async fn get_keys_from_dynamo(
    dynamo_client: &aws_sdk_dynamodb::Client,
    table_name: &String,
) -> anyhow::Result<JWTKResponse> {
    let keys_results = dynamo_client
        .get_item()
        .table_name(table_name)
        .key(
            "PK",
            aws_sdk_dynamodb::types::AttributeValue::S("#KEYS".to_string()),
        )
        .send()
        .await?;

    let keys_resp = keys_results.item.unwrap();

    let keys_json = keys_resp.get("keys").unwrap().as_s().unwrap();

    return Ok(serde_json::from_str(&keys_json)?);
}
``` 

And a similar function for putting keys to Dynamo. I didn't spend much time designing Dynamo's table. It has static `PK` and `keys` attributes, which is a stringified JSON object.

At this point, the logic is that in the main function, I check Dynamo for public keys, and go to Okta only if they are missing.

```rust
#[tokio::main]
async fn main() -> Result<(), Error> {

    let table_name = std::env::var("KEYS_TABLE_NAME").unwrap();
    
    let okta_keys_endpoint = std::env::var("OKTA_KEYS_ENDPOINT").unwrap();

    let dynamo_client = dynamo_service::get_dynamo_client().await;

    println!("getting keys from dynamo");

    let keys_from_dynamo = dynamo_service::get_keys_from_dynamo(&dynamo_client, &table_name).await;

    // if keys are present in dynamo - use them
    // get them from Okta and store in dynamo

    let stored_keys: StoredKeys = match keys_from_dynamo {
        Ok(keys_dynamo) => {
            println!("got keys from dynamo");
            jwtk_response_to_map(keys_dynamo)
        }
        Err(_) => {
            println!("no keys in dynamo - getting them from okta and storing in  dynamo");
            let keys_resp = get_keys_from_okta(okta_keys_endpoint).await.unwrap();
            // ignoring result of putting record to dynamo
            let _ =
                dynamo_service::store_keys_in_dynamo(&dynamo_client, &table_name, &keys_resp).await;
            jwtk_response_to_map(keys_resp)
        }
    };

    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        // disable printing the name of the module in every log line.
        .with_target(false)
        // disabling time is handy because CloudWatch will add the ingestion time.
        .without_time()
        .init();

    // run(service_fn(function_handler)).await
    run(service_fn(|event| function_handler(&stored_keys, event))).await
}
```
## Handler

The function handler looks pretty straightforward

```rust
async fn function_handler(
    current_keys: &StoredKeys,
    event: LambdaEvent<ApiGatewayCustomAuthorizerRequest>,
) -> Result<ApiGatewayCustomAuthorizerResponse<AuthContext>, Error> {

    let token: String  = event.payload.authorization_token.unwrap();

    let token_data: Result<jsonwebtoken::TokenData<Claims>, anyhow::Error> = jwt_service::validate_token(&token, current_keys);

    let response: ApiGatewayCustomAuthorizerResponse<AuthContext> = iam_policy:: prepare_response(token_data)?;

    return Ok(response)
}
```
In my case, the `AuthContext` of the `ApiGatewayCustomAuthorizerResponse` stays dummy, but it might be a useful way to pass information from the authorizer to the underlying function.

```rust
#[derive(Serialize, Deserialize)]
pub struct AuthContext {
    text: String,
}
```

The two last pieces of the puzzle are to implement JWT validation and preparation of the IAM policy.

## JWT 

I use a great [jsonwebtoken crate](https://github.com/Keats/jsonwebtoken), which makes decoding tokens easy. 

```rust
pub fn validate_token(
    token: &String,
    current_keys: &StoredKeys,
) -> anyhow::Result<jsonwebtoken::TokenData<Claims>> {

    let token_header: Header = jsonwebtoken::decode_header(&token)?;
    
    let kid: String = token_header.kid.unwrap();
    
    let public_key_to_use: &JWTK = current_keys.keys.get(&kid).unwrap();
    
    let decoding_key: jsonwebtoken::DecodingKey =
        jsonwebtoken::DecodingKey::from_rsa_components(&public_key_to_use.n, &public_key_to_use.e)?;
    
    let expected_aud: String = "api://default".to_string();
    
    let mut validation: jsonwebtoken::Validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::RS256);
    
    validation.set_audience(&[expected_aud]);
    
    let token_data: jsonwebtoken::TokenData<Claims> = jsonwebtoken::decode::<Claims>(&token, &decoding_key, &validation)?;
    
    return Ok(token_data);
}
```

Extracted `TokenData` has the shape of `Claims`

```rust
pub struct Claims {
    aud: String, // Optional. Audience
    exp: usize, // Required (validate_exp defaults to true in validation). Expiration time (as UTC timestamp)
    iat: usize, // Optional. Issued at (as UTC timestamp)
    iss: String, // Optional. Issuer
    uid: String,
    sub: String,      // Optional. Subject (whom token refers to)
    scp: Vec<String>, // Optional. Scopes (permissions)>
}
```

For me, the most interesting field is `uid`, but different auth flows would utilize other fields.

## IAM Policy

Once I validate the token I can use the user ID to prepare IAM Policy. As a function parameter, I expect `Result` type, because I want to prepare an explicit deny in the returned policy in case of validation errors.

```rust

pub fn prepare_response(
    validated_token: anyhow::Result<jsonwebtoken::TokenData<Claims>>,
) -> anyhow::Result<ApiGatewayCustomAuthorizerResponse<AuthContext>> {
    let policy = match validated_token {
        Ok(token_data) => {
            let path_to_allow = format!(
                "arn:aws:execute-api:us-east-1:765444088049:qma7pp9zmf/Prod/GET/hello/{user_id}",
                user_id = token_data.claims.uid
            );

            let statement = vec![IamPolicyStatement {
                effect: Some("Allow".to_string()),
                action: vec!["execute-api:Invoke".to_string()],
                resource: vec![path_to_allow],
            }];

            ApiGatewayCustomAuthorizerPolicy {
                version: Some("2012-10-17".to_string()),
                statement,
            }
        }
        Err(e) => {
            println!("token validation failed with error: {:?}", e);

            let path_to_deny =
                format!("arn:aws:execute-api:us-east-1:765444088049:qma7pp9zmf/Prod/GET/hello/*",);

            let statement = vec![IamPolicyStatement {
                effect: Some("Deny".to_string()),
                action: vec!["execute-api:Invoke".to_string()],
                resource: vec![path_to_deny],
            }];

            ApiGatewayCustomAuthorizerPolicy {
                version: Some("2012-10-17".to_string()),
                statement,
            }
        }
    };
    // Prepare the response
    let resp = ApiGatewayCustomAuthorizerResponse {
        principal_id: Some("12345abc".to_string()),
        policy_document: policy,
        context: AuthContext {
            text: "dummy context".to_string(),
        },
        usage_identifier_key: None,
    };
    return Ok(resp);
}
```
 
# Testing

Ok, let's test the endpoint. The expectation is that, without a valid token, I can't get a response from `https://<API_GATEWY_URL>/Prod/hello/<user_id>`

Let's start by using a random token and random user ID. Of course, it shouldn't work. We expect it to be explicitly denied because the wrong token causes a validation error in the authorizer. 


![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/691nn5svp6dmateoe08l.png)

Ok, now the happy path. I take the correct token from the front end (as described above). 

```bash
cd samples-js-vue/okta-hosted-login
npm run start
```

In the local store, I can find access token and user ID. With this data, I am able to call my endpoint and get a response


![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/ycuugr28zt1aqvyse8ro.png)

The last test - the proper token, but mismatched user ID. The result is 403. The token is valid but the given user can't access this specific endpoint. 


![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/it36ngi24iqt2vprwozj.png)



It works!

In real life, you would definitely want to do a little work on the messages in the 403 responses, but is not needed for my simple scenario

## Performance

For the authorizer lambda function, I've observed ~150ms duration for the cold start with Okta public keys stored in DynamoDB. Before including cache in DynamoDB I saw ~700ms duration for cold starts because public keys were cached only for hot start. 

_I didn't implement any rotation for public keys stored in DynamoDB. Okta rotates keys around four times per year (but it can change) and gives 2 weeks for all applications to update cached keys. In real life, I would set TTL on Dynamo table, or schedule a separate lambda to overwrite stored keys every two weeks or so._