From 93a0c3f7c4a6295c8e0cfc3f06283c811a9dbb74 Mon Sep 17 00:00:00 2001
From: David Calavera <david.calavera@gmail.com>
Date: Tue, 20 Feb 2024 21:16:53 -0800
Subject: [PATCH] Add example on how to use Api Gateway authorizers in Axum.

This also shows how to work with the RequestExt trait and the RequestContext object.

Signed-off-by: David Calavera <david.calavera@gmail.com>
---
 .../http-axum-apigw-authorizer/Cargo.toml     | 14 ++++
 examples/http-axum-apigw-authorizer/README.md | 13 +++
 .../http-axum-apigw-authorizer/src/main.rs    | 80 +++++++++++++++++++
 3 files changed, 107 insertions(+)
 create mode 100644 examples/http-axum-apigw-authorizer/Cargo.toml
 create mode 100644 examples/http-axum-apigw-authorizer/README.md
 create mode 100644 examples/http-axum-apigw-authorizer/src/main.rs

diff --git a/examples/http-axum-apigw-authorizer/Cargo.toml b/examples/http-axum-apigw-authorizer/Cargo.toml
new file mode 100644
index 00000000..5c3e806e
--- /dev/null
+++ b/examples/http-axum-apigw-authorizer/Cargo.toml
@@ -0,0 +1,14 @@
+[package]
+name = "http-axum-apigw-authorizer"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+axum = "0.7"
+lambda_http = { path = "../../lambda-http" }
+lambda_runtime = { path = "../../lambda-runtime" }
+serde = "1.0.196"
+serde_json = "1.0"
+tokio = { version = "1", features = ["macros"] }
+tracing = { version = "0.1", features = ["log"] }
+tracing-subscriber = { version = "0.3", default-features = false, features = ["fmt"] }
diff --git a/examples/http-axum-apigw-authorizer/README.md b/examples/http-axum-apigw-authorizer/README.md
new file mode 100644
index 00000000..2d05df59
--- /dev/null
+++ b/examples/http-axum-apigw-authorizer/README.md
@@ -0,0 +1,13 @@
+# Axum example that integrates with Api Gateway authorizers
+
+This example shows how to extract information from the Api Gateway Request Authorizer in an Axum handler.
+
+## Build & Deploy
+
+1. Install [cargo-lambda](https://github.com/cargo-lambda/cargo-lambda#installation)
+2. Build the function with `cargo lambda build --release`
+3. Deploy the function to AWS Lambda with `cargo lambda deploy --iam-role YOUR_ROLE`
+
+## Build for ARM 64
+
+Build the function with `cargo lambda build --release --arm64`
diff --git a/examples/http-axum-apigw-authorizer/src/main.rs b/examples/http-axum-apigw-authorizer/src/main.rs
new file mode 100644
index 00000000..3b8a0810
--- /dev/null
+++ b/examples/http-axum-apigw-authorizer/src/main.rs
@@ -0,0 +1,80 @@
+use axum::{
+    async_trait,
+    extract::{FromRequest, Request},
+    http::StatusCode,
+    response::Json,
+    routing::get,
+    Router,
+};
+use lambda_http::{run, Error, RequestExt};
+use serde_json::{json, Value};
+use std::{collections::HashMap, env::set_var};
+
+struct AuthorizerField(String);
+struct AuthorizerFields(HashMap<String, serde_json::Value>);
+
+#[async_trait]
+impl<S> FromRequest<S> for AuthorizerField
+where
+    S: Send + Sync,
+{
+    type Rejection = (StatusCode, &'static str);
+
+    async fn from_request(req: Request, _state: &S) -> Result<Self, Self::Rejection> {
+        req.request_context_ref()
+            .and_then(|r| r.authorizer())
+            .and_then(|a| a.fields.get("field_name"))
+            .and_then(|f| f.as_str())
+            .map(|v| Self(v.to_string()))
+            .ok_or_else(|| (StatusCode::BAD_REQUEST, "`field_name` authorizer field is missing"))
+    }
+}
+
+#[async_trait]
+impl<S> FromRequest<S> for AuthorizerFields
+where
+    S: Send + Sync,
+{
+    type Rejection = (StatusCode, &'static str);
+
+    async fn from_request(req: Request, _state: &S) -> Result<Self, Self::Rejection> {
+        req.request_context_ref()
+            .and_then(|r| r.authorizer())
+            .map(|a| Self(a.fields.clone()))
+            .ok_or_else(|| (StatusCode::BAD_REQUEST, "authorizer is missing"))
+    }
+}
+
+async fn extract_field(AuthorizerField(field): AuthorizerField) -> Json<Value> {
+    Json(json!({ "field extracted": field }))
+}
+
+async fn extract_all_fields(AuthorizerFields(fields): AuthorizerFields) -> Json<Value> {
+    Json(json!({ "authorizer fields": fields }))
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Error> {
+    // If you use API Gateway stages, the Rust Runtime will include the stage name
+    // as part of the path that your application receives.
+    // Setting the following environment variable, you can remove the stage from the path.
+    // This variable only applies to API Gateway stages,
+    // you can remove it if you don't use them.
+    // i.e with: `GET /test-stage/todo/id/123` without: `GET /todo/id/123`
+    set_var("AWS_LAMBDA_HTTP_IGNORE_STAGE_IN_PATH", "true");
+
+    // required to enable CloudWatch error logging by the runtime
+    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();
+
+    let app = Router::new()
+        .route("/extract-field", get(extract_field))
+        .route("/extract-all-fields", get(extract_all_fields));
+
+    run(app).await
+}