Add the ability to strip certain response headers

README.md

@@ -105,8 +105,9 @@ The readiness check port/path and traffic port can be configured using environme
 | 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  |
-| AWS_LWA_ERROR_STATUS_CODES                                  | comma-separated list of HTTP status codes that will cause Lambda invocations to fail (e.g. "500,502-504,422") | None  |
-| AWS_LWA_LAMBDA_RUNTIME_API_PROXY                              | overwrites `AWS_LAMBDA_RUNTIME_API` to allow proxying request (not affecting registration)                    | None       |
+| AWS_LWA_ERROR_STATUS_CODES                                   | comma-separated list of HTTP status codes that will cause Lambda invocations to fail (e.g. "500,502-504,422") | None  |
+| AWS_LWA_LAMBDA_RUNTIME_API_PROXY                             | overwrites `AWS_LAMBDA_RUNTIME_API` to allow proxying request (not affecting registration)                    | None       |
+| AWS_LWA_STRIP_RESPONSE_HEADERS                               | comma-separated list of HTTP response headers to remove from responses               | 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.
@@ -133,6 +134,8 @@ When enabled, this will compress responses unless it's an image as determined by
 **AWS_LWA_INVOKE_MODE** - Lambda function invoke mode, this should match Function Url invoke mode. The default is "buffered". When configured as "response_stream", Lambda Web Adapter will stream response to Lambda service [blog](https://aws.amazon.com/blogs/compute/introducing-aws-lambda-response-streaming/).
 Please check out [FastAPI with Response Streaming](examples/fastapi-response-streaming) example.
 
+**AWS_LWA_STRIP_RESPONSE_HEADERS** - Allows you to specify a list of HTTP response headers that should be removed from responses before they are sent back to the client. This is useful when you want to prevent certain headers from being exposed or when you need to override headers set by your web application framework. For example, setting `AWS_LWA_STRIP_RESPONSE_HEADERS=x-powered-by,server` will remove the `X-Powered-By` and `Server` headers from all responses. Multiple headers should be specified as a comma-separated list.
+
 **AWS_LWA_READINESS_CHECK_MIN_UNHEALTHY_STATUS** - allows you to customize which HTTP status codes are considered healthy and which ones are not
 
 **AWS_LWA_PASS_THROUGH_PATH** - Path to receive events payloads passed through from non-http event triggers. The default is "/events".

src/lib.rs

@@ -81,6 +81,7 @@ pub struct AdapterOptions {
     pub invoke_mode: LambdaInvokeMode,
     pub authorization_source: Option<String>,
     pub error_status_codes: Option<Vec<u16>>,
+    pub strip_response_headers: Option<Vec<String>>,
 }
 
 impl Default for AdapterOptions {
@@ -122,6 +123,9 @@ impl Default for AdapterOptions {
             error_status_codes: env::var("AWS_LWA_ERROR_STATUS_CODES")
                 .ok()
                 .map(|codes| parse_status_codes(&codes)),
+            strip_response_headers: env::var("AWS_LWA_STRIP_RESPONSE_HEADERS")
+                .ok()
+                .map(|h| h.split(',').map(|s| s.trim().to_string()).collect()),
         }
     }
 }
@@ -170,6 +174,7 @@ pub struct Adapter<C, B> {
     invoke_mode: LambdaInvokeMode,
     authorization_source: Option<String>,
     error_status_codes: Option<Vec<u16>>,
+    strip_response_headers: Option<Vec<String>>,
 }
 
 impl Adapter<HttpConnector, Body> {
@@ -208,6 +213,7 @@ impl Adapter<HttpConnector, Body> {
             invoke_mode: options.invoke_mode,
             authorization_source: options.authorization_source.clone(),
             error_status_codes: options.error_status_codes.clone(),
+            strip_response_headers: options.strip_response_headers.clone(),
         }
     }
 }
@@ -405,6 +411,17 @@ impl Adapter<HttpConnector, Body> {
         // remove "transfer-encoding" from the response to support "sam local start-api"
         app_response.headers_mut().remove("transfer-encoding");
 
+        // Remove any configured headers from the response
+        if let Some(headers_to_strip) = &self.strip_response_headers {
+            for header_name in headers_to_strip {
+                if let Ok(name) = HeaderName::from_bytes(header_name.as_bytes()) {
+                    app_response.headers_mut().remove(&name);
+                } else {
+                    tracing::warn!("Invalid header name to strip: {}", header_name);
+                }
+            }
+        }
+
         tracing::debug!(status = %app_response.status(), body_size = ?app_response.body().size_hint().lower(),
             app_headers = ?app_response.headers().clone(), "responding to lambda event");
 

tests/integ_tests/main.rs

@@ -750,6 +750,62 @@ async fn test_http_context_multi_headers() {
     assert_eq!("OK", body_to_string(response).await);
 }
 
+#[tokio::test]
+async fn test_http_strip_response_headers() {
+    // Start app server
+    let app_server = MockServer::start();
+
+    // An endpoint that returns multiple headers
+    let test_endpoint = app_server.mock(|when, then| {
+        when.method(GET).path("/");
+        then.status(200)
+            .header("x-custom-header", "value")
+            .header("server", "test-server")
+            .header("content-type", "application/json")
+            .body("{}");
+    });
+
+    // Initialize adapter with headers to strip
+    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(),
+        strip_response_headers: Some(vec!["x-custom-header".to_string(), "server".to_string()]),
+        ..Default::default()
+    });
+
+    // Prepare request
+    let req = LambdaEventBuilder::new().with_path("/").build();
+
+    // Convert to Request object and add 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();
+
+    // Verify headers were stripped
+    assert!(
+        !response.headers().contains_key("x-custom-header"),
+        "x-custom-header should be stripped"
+    );
+    assert!(
+        !response.headers().contains_key("server"),
+        "server header should be stripped"
+    );
+
+    // Verify other headers remain
+    assert!(
+        response.headers().contains_key("content-type"),
+        "content-type header should remain"
+    );
+    assert_eq!("{}", body_to_string(response).await);
+}
+
 async fn body_to_string(res: Response<Incoming>) -> String {
     let body_bytes = res.collect().await.unwrap().to_bytes();
     String::from_utf8_lossy(&body_bytes).to_string()