Pass Lambda Context to web app in a new header (#248)

* Pass Lambda Context to web app in a new header

* Added tests to make sure x-amzn-lambda-context header is present

Author: mbfreder

src/lib.rs

@@ -332,6 +332,7 @@ where
         }
 
         let request_context = event.request_context();
+        let lambda_context = &event.lambda_context();
         let path = event.raw_http_path().to_string();
         let mut path = path.as_str();
         let (parts, body) = event.into_parts();
@@ -349,6 +350,11 @@ where
             HeaderValue::from_bytes(serde_json::to_string(&request_context)?.as_bytes())?,
         );
 
+        // include lambda context in http header "x-amzn-lambda-context"
+        req_headers.append(
+            HeaderName::from_static("x-amzn-lambda-context"),
+            HeaderValue::from_bytes(serde_json::to_string(&lambda_context)?.as_bytes())?,
+        );
         let mut app_url = self.domain.clone();
         app_url.set_path(path);
         app_url.set_query(parts.uri.query());

tests/integ_tests/main.rs

@@ -5,12 +5,15 @@ use std::io;
 use std::io::prelude::*;
 
 use crate::common::LambdaEventBuilder;
+use http::HeaderMap;
+use http::Request;
 use http::{Method, Response};
 use httpmock::{
     Method::{DELETE, GET, POST, PUT},
     MockServer,
 };
 use hyper::{body, Body};
+use lambda_http::Context;
 use lambda_web_adapter::{Adapter, AdapterOptions, LambdaInvokeMode, Protocol};
 use tower::{Service, ServiceBuilder};
 
@@ -158,7 +161,13 @@ async fn test_http_basic_request() {
 
     // // Call the adapter service with basic request
     let req = LambdaEventBuilder::new().with_path("/hello").build();
-    let response = adapter.call(req.into()).await.expect("Request failed");
+
+    // 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();
@@ -202,8 +211,13 @@ async fn test_http_headers() {
         .with_header("foo", "bar")
         .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);
+
     // Call the adapter service with request
-    let response = adapter.call(req.into()).await.expect("Request failed");
+    let response = adapter.call(request).await.expect("Request failed");
 
     // Assert endpoint was called once
     test_endpoint.assert();
@@ -245,8 +259,13 @@ async fn test_http_path_encoding() {
     // Prepare request
     let req = LambdaEventBuilder::new().with_path("/Año_1234").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);
+
     // Call the adapter service with request
-    let response = adapter.call(req.into()).await.expect("Request failed");
+    let response = adapter.call(request).await.expect("Request failed");
 
     // Assert endpoint was called once
     test_endpoint.assert();
@@ -293,8 +312,13 @@ async fn test_http_query_params() {
         .with_query("fizz", "buzz")
         .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);
+
     // Call the adapter service with request
-    let response = adapter.call(req.into()).await.expect("Request failed");
+    let response = adapter.call(request).await.expect("Request failed");
 
     // Assert endpoint was called once
     test_endpoint.assert();
@@ -354,11 +378,20 @@ async fn test_http_post_put_delete() {
         .with_method(Method::DELETE)
         .with_path("/")
         .build();
+    // We convert to Request object because it allows us to add the lambda Context
+    let mut post_request = Request::from(post_req);
+    add_lambda_context_to_request(&mut post_request);
+
+    let mut put_request = Request::from(put_req);
+    add_lambda_context_to_request(&mut put_request);
+
+    let mut delete_request = Request::from(delete_req);
+    add_lambda_context_to_request(&mut delete_request);
 
     // Call the adapter service with requests
-    let post_response = adapter.call(post_req.into()).await.expect("Request failed");
-    let put_response = adapter.call(put_req.into()).await.expect("Request failed");
-    let delete_response = adapter.call(delete_req.into()).await.expect("Request failed");
+    let post_response = adapter.call(post_request).await.expect("Request failed");
+    let put_response = adapter.call(put_request).await.expect("Request failed");
+    let delete_response = adapter.call(delete_request).await.expect("Request failed");
 
     // Assert endpoints were called
     post_endpoint.assert();
@@ -407,7 +440,13 @@ async fn test_http_compress() {
         .with_path("/hello")
         .with_header("accept-encoding", "gzip")
         .build();
-    let response = svc.call(req.into()).await.expect("Request failed");
+
+    // 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 = svc.call(request).await.expect("Request failed");
 
     // Assert endpoint was called once
     hello.assert();
@@ -453,7 +492,13 @@ async fn test_http_compress_disallowed_type() {
         .with_path("/hello")
         .with_header("accept-encoding", "gzip")
         .build();
-    let response = adapter.call(req.into()).await.expect("Request failed");
+
+    // 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();
@@ -506,7 +551,13 @@ async fn test_http_compress_already_compressed() {
         .with_path("/hello")
         .with_header("accept-encoding", "gzip")
         .build();
-    let response = svc.call(req.into()).await.expect("Request failed");
+
+    // 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 = svc.call(request).await.expect("Request failed");
 
     // Assert endpoint was called once
     hello.assert();
@@ -521,6 +572,54 @@ async fn test_http_compress_already_compressed() {
     );
 }
 
+#[tokio::test]
+async fn test_http_lambda_context_header() {
+    // Start app server
+    let app_server = MockServer::start();
+
+    // An endpoint that expects and returns headers
+    let test_endpoint = app_server.mock(|when, then| {
+        when.method(GET).path("/").header_exists("x-amzn-lambda-context");
+        then.status(200).header("fizz", "buzz").body("OK");
+    });
+
+    // Initialize adapter and do readiness check
+    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(),
+        readiness_check_protocol: Protocol::Http,
+        async_init: false,
+        base_path: None,
+        compression: false,
+        enable_tls: false,
+        tls_server_name: None,
+        tls_cert_file: None,
+        invoke_mode: LambdaInvokeMode::Buffered,
+    });
+
+    // Prepare request
+    let req = LambdaEventBuilder::new().with_path("/").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);
+
+    // Call the adapter service with request
+    let response = adapter.call(request).await.expect("Request failed");
+
+    // Assert endpoint was called once
+    test_endpoint.assert();
+
+    // and response has expected content
+    assert_eq!(200, response.status());
+    assert!(response.headers().contains_key("fizz"));
+    assert_eq!("buzz", response.headers().get("fizz").unwrap());
+    assert_eq!("OK", body_to_string(response).await);
+}
+
 async fn body_to_string(res: Response<Body>) -> String {
     let body_bytes = body::to_bytes(res.into_body()).await.unwrap();
     String::from_utf8_lossy(&body_bytes).to_string()
@@ -537,3 +636,17 @@ fn decode_reader(bytes: &[u8]) -> io::Result<String> {
     gz.read_to_string(&mut s)?;
     Ok(s)
 }
+
+fn add_lambda_context_to_request(request: &mut Request<lambda_http::Body>) {
+    // create a HeaderMap to build the lambda context
+    let mut headers = HeaderMap::new();
+    headers.insert("lambda-runtime-aws-request-id", "my_id".parse().unwrap());
+    headers.insert("lambda-runtime-deadline-ms", "123".parse().unwrap());
+    headers.insert("lambda-runtime-client-context", "{}".parse().unwrap());
+
+    // converts HeaderMap to Context
+    let context = Context::try_from(headers).expect("Couldn't convert HeaderMap to Context");
+
+    // add Context to the request
+    request.extensions_mut().insert(context);
+}