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
For the application type, I have OIDC and single-page app.
We need a user, so I added one manually
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
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.
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
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:
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
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
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-...."
}
]
}
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>,
}
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)
}
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>,
}
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)?);
}
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
}
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)
}
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,
}
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);
}
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)>
}
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);
}
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.
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
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
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.
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)
A really cool example! Thanks!