fix: remove user-set x-amzn-*-context headers

Currently, when a client sets their own `x-amzn-*-context` headers, they
will be included in the request to the app server.

This request to the lambda with the adapter installed:

```
$ curl -H "x-amzn-request-context: boom" http://localhost:8000
```

Will result in this request from the adapter to the app server:

```
GET / HTTP/1.1
[...]
x-amzn-request-context: boom
x-amzn-request-context: [json context from adapter]
x-amzn-lambda-context: [json context from adapter]
[...]
```

Many implementations of web frameworks will take the *first* header when
you access a header key.

For example, Starlette/FastAPI does this:
https://github.com/encode/starlette/blob/ad02ee6336faadadf97e0c79dd3a91759f1f32a7/starlette/datastructures.py#L562-L567

This patch removes every user-set `x-amzn-{request,lambda}-context`
headers from the request before adding in JSON context from the adapter.

Signed-off-by: Thitat Auareesuksakul <[email protected]>

Author: fluxth

src/lib.rs

@@ -14,7 +14,7 @@ use std::{
 };
 
 use http::{
-    header::{HeaderName, HeaderValue},
+    header::{Entry, HeaderName, HeaderValue},
     Method, StatusCode,
 };
 use hyper::{
@@ -359,12 +359,22 @@ where
 
         let mut req_headers = parts.headers;
 
+        // remove "x-amzn-request-context" if already set
+        if let Entry::Occupied(entry) = req_headers.entry("x-amzn-request-context") {
+            entry.remove();
+        }
+
         // include request context in http header "x-amzn-request-context"
         req_headers.append(
             HeaderName::from_static("x-amzn-request-context"),
             HeaderValue::from_bytes(serde_json::to_string(&request_context)?.as_bytes())?,
         );
 
+        // remove "x-amzn-lambda-context" if already set
+        if let Entry::Occupied(entry) = req_headers.entry("x-amzn-lambda-context") {
+            entry.remove();
+        }
+
         // include lambda context in http header "x-amzn-lambda-context"
         req_headers.append(
             HeaderName::from_static("x-amzn-lambda-context"),

tests/integ_tests/main.rs

@@ -584,13 +584,16 @@ async fn test_http_compress_already_compressed() {
 }
 
 #[tokio::test]
-async fn test_http_lambda_context_header() {
+async fn test_http_context_headers() {
     // 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");
+        when.method(GET)
+            .path("/")
+            .header_exists("x-amzn-lambda-context")
+            .header_exists("x-amzn-request-context");
         then.status(200).header("fizz", "buzz").body("OK");
     });
 
@@ -632,6 +635,80 @@ async fn test_http_lambda_context_header() {
     assert_eq!("OK", body_to_string(response).await);
 }
 
+#[tokio::test]
+async fn test_http_context_multi_headers() {
+    // 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("/")
+            .matches(|req| {
+                req.headers
+                    .as_ref()
+                    .unwrap()
+                    .iter()
+                    .filter(|(key, _value)| key == "x-amzn-lambda-context")
+                    .count()
+                    == 1
+            })
+            .matches(|req| {
+                req.headers
+                    .as_ref()
+                    .unwrap()
+                    .iter()
+                    .filter(|(key, _value)| key == "x-amzn-request-context")
+                    .count()
+                    == 1
+            });
+        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,
+        readiness_check_min_unhealthy_status: 500,
+        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("/")
+        .with_header("x-amzn-lambda-context", "header_from_client_1")
+        .with_header("x-amzn-lambda-context", "header_from_client_2")
+        .with_header("x-amzn-request-context", "header_from_client_1")
+        .with_header("x-amzn-request-context", "header_from_client_2")
+        .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()