add authorization_source feature (awslabs#478)

* add authorization_source feature

* use option

Author: tmokmss

README.md

@@ -104,6 +104,7 @@ The readiness check port/path and traffic port can be configured using environme
 | AWS_LWA_ENABLE_COMPRESSION                                   | enable gzip compression for response body                                            | "false"    |
 | AWS_LWA_INVOKE_MODE                                          | Lambda function invoke mode: "buffered" or "response_stream", default is "buffered"  | "buffered" |
 | AWS_LWA_PASS_THROUGH_PATH                                    | the path for receiving event payloads that are passed through from non-http triggers | "/events"  |
+| AWS_LWA_AUTHORIZATION_SOURCE                                 | a header name to be replaced to `Authorization` | None  |
 
 > **Note:**
 > We use "AWS_LWA_" prefix to namespacing all environment variables used by Lambda Web Adapter. The original ones will be supported until we reach version 1.0.
@@ -134,6 +135,8 @@ Please check out [FastAPI with Response Streaming](examples/fastapi-response-str
 
 **AWS_LWA_PASS_THROUGH_PATH** - Path to receive events payloads passed through from non-http event triggers. The default is "/events".
 
+**AWS_LWA_AUTHORIZATION_SOURCE** - When set, Lambda Web Adapter replaces the specified header name to `Authorization` before proxying a request. This is useful when you use Lambda function URL with [IAM auth type](https://docs.aws.amazon.com/lambda/latest/dg/urls-auth.html), which reserves Authorization header for IAM authentication, but you want to still use Authorization header for your backend apps. This feature is disabled by default.
+
 ## Request Context
 
 **Request Context** is metadata API Gateway sends to Lambda for a request. It usually contains requestId, requestTime, apiId, identity, and authorizer. Identity and authorizer are useful to get client identity for authorization. API Gateway Developer Guide contains more details [here](https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-input-format).  

src/lib.rs

@@ -77,6 +77,7 @@ pub struct AdapterOptions {
     pub async_init: bool,
     pub compression: bool,
     pub invoke_mode: LambdaInvokeMode,
+    pub authorization_source: Option<String>,
 }
 
 impl Default for AdapterOptions {
@@ -114,6 +115,7 @@ impl Default for AdapterOptions {
                 .unwrap_or("buffered".to_string())
                 .as_str()
                 .into(),
+            authorization_source: env::var("AWS_LWA_AUTHORIZATION_SOURCE").ok(),
         }
     }
 }
@@ -131,6 +133,7 @@ pub struct Adapter<C, B> {
     path_through_path: String,
     compression: bool,
     invoke_mode: LambdaInvokeMode,
+    authorization_source: Option<String>,
 }
 
 impl Adapter<HttpConnector, Body> {
@@ -167,6 +170,7 @@ impl Adapter<HttpConnector, Body> {
             ready_at_init: Arc::new(AtomicBool::new(false)),
             compression: options.compression,
             invoke_mode: options.invoke_mode,
+            authorization_source: options.authorization_source.clone(),
         }
     }
 }
@@ -312,6 +316,14 @@ impl Adapter<HttpConnector, Body> {
             HeaderName::from_static("x-amzn-lambda-context"),
             HeaderValue::from_bytes(serde_json::to_string(&lambda_context)?.as_bytes())?,
         );
+
+        if let Some(authorization_source) = self.authorization_source.as_deref() {
+            if req_headers.contains_key(authorization_source) {
+                let original = req_headers.remove(authorization_source).unwrap();
+                req_headers.insert("authorization", original);
+            }
+        }
+
         let mut app_url = self.domain.clone();
         app_url.set_path(path);
         app_url.set_query(parts.uri.query());

tests/integ_tests/main.rs

@@ -41,6 +41,7 @@ fn test_adapter_options_from_env() {
     env::set_var("AWS_LWA_TLS_SERVER_NAME", "api.example.com");
     env::remove_var("AWS_LWA_TLS_CERT_FILE");
     env::set_var("AWS_LWA_INVOKE_MODE", "buffered");
+    env::set_var("AWS_LWA_AUTHORIZATION_SOURCE", "auth-token");
 
     // Initialize adapter with env options
     let options = AdapterOptions::default();
@@ -55,6 +56,7 @@ fn test_adapter_options_from_env() {
     assert!(options.async_init);
     assert!(options.compression);
     assert_eq!(LambdaInvokeMode::Buffered, options.invoke_mode);
+    assert_eq!(Some("auth-token".into()), options.authorization_source);
 }
 
 #[test]
@@ -69,6 +71,7 @@ fn test_adapter_options_from_namespaced_env() {
     env::set_var("AWS_LWA_ASYNC_INIT", "true");
     env::set_var("AWS_LWA_ENABLE_COMPRESSION", "true");
     env::set_var("AWS_LWA_INVOKE_MODE", "response_stream");
+    env::set_var("AWS_LWA_AUTHORIZATION_SOURCE", "auth-token");
 
     // Initialize adapter with env options
     let options = AdapterOptions::default();
@@ -84,6 +87,7 @@ fn test_adapter_options_from_namespaced_env() {
     assert!(options.async_init);
     assert!(options.compression);
     assert_eq!(LambdaInvokeMode::ResponseStream, options.invoke_mode);
+    assert_eq!(Some("auth-token".into()), options.authorization_source);
 }
 
 #[test]
@@ -607,6 +611,47 @@ async fn test_http_content_encoding_suffix() {
     assert_eq!(json_data.to_owned(), body_to_string(response).await);
 }
 
+#[tokio::test]
+async fn test_http_authorization_source() {
+    // Start app server
+    let app_server = MockServer::start();
+    let hello = app_server.mock(|when, then| {
+        when.method(GET).path("/hello").header_exists("Authorization");
+        then.status(200).body("Hello World");
+    });
+
+    // Initialize adapter
+    let mut adapter = Adapter::new(&AdapterOptions {
+        host: app_server.host(),
+        port: app_server.port().to_string(),
+        readiness_check_port: app_server.port().to_string(),
+        readiness_check_path: "/healthcheck".to_string(),
+        authorization_source: Some("auth-token".to_string()),
+        ..Default::default()
+    });
+
+    // // Call the adapter service with basic request
+    let req = LambdaEventBuilder::new()
+        .with_path("/hello")
+        .with_header("auth-token", "Bearer token")
+        .build();
+
+    // We convert to Request object because it allows us to add
+    // the lambda Context
+    let mut request = Request::from(req);
+    add_lambda_context_to_request(&mut request);
+
+    let response = adapter.call(request).await.expect("Request failed");
+
+    // Assert endpoint was called once
+    hello.assert();
+
+    // and response has expected content
+    assert_eq!(200, response.status());
+    assert_eq!(response.headers().get("content-length").unwrap(), "11");
+    assert_eq!("Hello World", body_to_string(response).await);
+}
+
 #[tokio::test]
 async fn test_http_context_multi_headers() {
     // Start app server