DEV Community

szymon-szym
szymon-szym

Posted on

AWS Custom API Gateway Authorizer with Rust and Okta

Introduction

The code for this blog post is here

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 with a free plan and added a new application

Image description

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

Image description

We need a user, so I added one manually

Image description

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



sam init


Enter fullscreen mode Exit fullscreen mode

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 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

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



cd okta-hosted-login 
npm run serve


Enter fullscreen mode Exit fullscreen mode

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

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



cargo lambda new authorizer


Enter fullscreen mode Exit fullscreen mode

Image description

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:



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



Enter fullscreen mode Exit fullscreen mode

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:



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


Enter fullscreen mode Exit fullscreen mode

I create a struct to hold this response:



#[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>,
}


Enter fullscreen mode Exit fullscreen mode

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



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


Enter fullscreen mode Exit fullscreen mode

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



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


Enter fullscreen mode Exit fullscreen mode

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



// ...
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)?);
}


Enter fullscreen mode Exit fullscreen mode

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.



#[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
}


Enter fullscreen mode Exit fullscreen mode

Handler

The function handler looks pretty straightforward



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)
}


Enter fullscreen mode Exit fullscreen mode

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.



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


Enter fullscreen mode Exit fullscreen mode

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, which makes decoding tokens easy.



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);
}


Enter fullscreen mode Exit fullscreen mode

Extracted TokenData has the shape of Claims



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)>
}


Enter fullscreen mode Exit fullscreen mode

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.




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);
}


Enter fullscreen mode Exit fullscreen mode

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

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



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


Enter fullscreen mode Exit fullscreen mode

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

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

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.

Top comments (1)

Collapse
 
rdarrylr profile image
Darryl Ruggles

A really cool example! Thanks!