Optional gzip compression of responses (#157)

* Optional gzip compression of responses

* Format with formatOnSave

* Docs, rename env var, add 2 content-types

* Remove `unsafe` from integration test

* Add E2Es for compressed responses

* Code review fixes

Author: huntharo

Cargo.lock

@@ -20,6 +20,7 @@ exclude = ["examples"]
 http = "0.2"
 hyper = { version = "0.14", features = ["client"] }
 lambda_http = "0.7.3"
+flate2 = "1.0.25"
 tokio = { version = "1.24", features = [
     "macros",
     "io-util",

Cargo.toml

@@ -69,14 +69,15 @@ After passing readiness check, Lambda Web Adapter will start Lambda Runtime and
 
 The readiness check port/path and traffic port can be configured using environment variables. These environment variables can be defined either within docker file or as Lambda function configuration.
 
-| Environment Variable     | Description                                                          | Default |
-|--------------------------|----------------------------------------------------------------------|---------|
-| PORT                     | traffic port                                                         | "8080"  |
-| READINESS_CHECK_PORT     | readiness check port, default to the traffic port                    | PORT    |
-| READINESS_CHECK_PATH     | readiness check path                                                 | "/"     |
-| READINESS_CHECK_PROTOCOL | readiness check protocol: "http" or "tcp", default is "http"         | "http"  |
-| ASYNC_INIT               | enable asynchronous initialization for long initialization functions | "false" |
-| REMOVE_BASE_PATH         | (optional) the base path to be removed from request path             | None    |
+| Environment Variable       | Description                                                          | Default |
+| -------------------------- | -------------------------------------------------------------------- | ------- |
+| PORT                       | traffic port                                                         | "8080"  |
+| READINESS_CHECK_PORT       | readiness check port, default to the traffic port                    | PORT    |
+| READINESS_CHECK_PATH       | readiness check path                                                 | "/"     |
+| READINESS_CHECK_PROTOCOL   | readiness check protocol: "http" or "tcp", default is "http"         | "http"  |
+| ASYNC_INIT                 | enable asynchronous initialization for long initialization functions | "false" |
+| REMOVE_BASE_PATH           | (optional) the base path to be removed from request path             | None    |
+| AWS_LWA_ENABLE_COMPRESSION | (optional) enable gzip compression for response body                 | "false" |
 
 **ASYNC_INIT** Lambda managed runtimes offer up to 10 seconds for function initialization. During this period of time, Lambda functions have burst of CPU to accelerate initialization, and it is free. 
 If a lambda function couldn't complete the initialization within 10 seconds, Lambda will restart the function, and bill for the initialization. 
@@ -85,6 +86,10 @@ When this feature is enabled, Lambda Web Adapter performs readiness check up to
 Lambda Web Adapter signals to Lambda service that the init is completed, and continues readiness check in the handler. 
 This feature is disabled by default. Enable it by setting environment variable `ASYNC_INIT` to `true`. 
 
+**AWS_LWA_ENABLE_COMPRESSION** Lambda Web Adapter supports gzip compression for response body. This feature is disabled by default. Enable it by setting environment variable `AWS_LWA_ENABLE_COMPRESSION` to `true`.
+When enabled, Lambda Web Adapter will check the `Accept-Encoding` header in the request, and compress the response body if the header contains `gzip`, if the response body is not already compressed, and if the `Content-Type` starts with `text/` or is `application/javascript`, `application/json`, `application/json+ld`, `application/xml`, `application/xhtml+xml`, `application/x-javascript`, or `image/svg+xml`.
+Note that the `Content-Length` header will be set to the compressed size, not the original size.
+
 **REMOVE_BASE_PATH** - The value of this environment variable tells the adapter whether the application is running under a base path.
 For example, you could have configured your API Gateway to have a /orders/{proxy+} and a /catalog/{proxy+} resource.
 Each resource is handled by a separate Lambda functions. For this reason, the application inside Lambda may not be aware of the fact that the /orders path exists.

README.md

@@ -4,6 +4,7 @@
 use std::{
     env,
     future::Future,
+    io::prelude::*,
     pin::Pin,
     sync::{
         atomic::{AtomicBool, Ordering},
@@ -12,11 +13,14 @@ use std::{
     time::Duration,
 };
 
+use flate2::write::GzEncoder;
+use flate2::Compression;
 use http::{
     header::{HeaderName, HeaderValue},
     Method, StatusCode, Uri,
 };
 use hyper::{
+    body,
     body::HttpBody,
     client::{Client, HttpConnector},
     Body,
@@ -60,6 +64,7 @@ pub struct AdapterOptions {
     pub readiness_check_protocol: Protocol,
     pub base_path: Option<String>,
     pub async_init: bool,
+    pub compression: bool,
 }
 
 impl AdapterOptions {
@@ -79,6 +84,10 @@ impl AdapterOptions {
                 .unwrap_or_else(|_| "false".to_string())
                 .parse()
                 .unwrap_or(false),
+            compression: env::var("AWS_LWA_ENABLE_COMPRESSION")
+                .unwrap_or_else(|_| "false".to_string())
+                .parse()
+                .unwrap_or(false),
         }
     }
 }
@@ -92,6 +101,7 @@ pub struct Adapter {
     ready_at_init: Arc<AtomicBool>,
     domain: Uri,
     base_path: Option<String>,
+    compression: bool,
 }
 
 impl Adapter {
@@ -120,6 +130,7 @@ impl Adapter {
             base_path: options.base_path.clone(),
             async_init: options.async_init,
             ready_at_init: Arc::new(AtomicBool::new(false)),
+            compression: options.compression,
         }
     }
 
@@ -206,6 +217,12 @@ impl Adapter {
             path = path.trim_start_matches(base_path);
         }
 
+        let accepts_gzip = parts
+            .headers
+            .get("accept-encoding")
+            .map(|v| v.to_str().unwrap_or_default().contains("gzip"))
+            .unwrap_or_default();
+
         let mut req_headers = parts.headers;
 
         // include request context in http header "x-amzn-request-context"
@@ -237,6 +254,49 @@ impl Adapter {
         tracing::debug!(status = %app_response.status(), body_size = app_response.body().size_hint().lower(),
             app_headers = ?app_response.headers().clone(), "responding to lambda event");
 
+        let response_compressed = app_response.headers().get("content-encoding").is_some();
+
+        let content_type = if let Some(content_type) = app_response.headers().get("content-type") {
+            content_type.to_str().unwrap()
+        } else {
+            ""
+        };
+
+        let compressable_content_type = content_type.starts_with("text/")
+            || content_type.starts_with("application/json")
+            || content_type.starts_with("application/ld+json")
+            || content_type.starts_with("application/javascript")
+            || content_type.starts_with("image/svg+xml")
+            || content_type.starts_with("application/xhtml+xml")
+            || content_type.starts_with("application/x-javascript")
+            || content_type.starts_with("application/xml");
+
+        // Gzip the response if the client accepts it
+        let app_response = if !self.compression {
+            app_response
+        } else if accepts_gzip && !response_compressed && compressable_content_type {
+            let (parts, body) = app_response.into_parts();
+            let mut builder = hyper::Response::builder().status(parts.status).version(parts.version);
+            if let Some(headers) = builder.headers_mut() {
+                // Remove the content-length header as we can't overwrite it after setting it
+                let mut clean_headers = parts.headers.clone();
+                clean_headers.remove(http::header::CONTENT_LENGTH);
+
+                headers.extend(clean_headers);
+            }
+            let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
+            encoder.write_all(&body::to_bytes(body).await.unwrap())?;
+            let gzipped_body = encoder.finish()?;
+
+            builder
+                // Write the new content-length header
+                .header(http::header::CONTENT_LENGTH, gzipped_body.len().to_string())
+                .header("content-encoding", "gzip")
+                .body(hyper::Body::from(gzipped_body))?
+        } else {
+            app_response
+        };
+
         Ok(app_response)
     }
 }

src/lib.rs

@@ -61,6 +61,9 @@ Resources:
       Handler: bootstrap
       Runtime: provided.al2
       MemorySize: 256
+      Environment:
+        Variables:
+          AWS_LWA_ENABLE_COMPRESSION: 'true'
       Layers:
         - !Ref AdapterLayerArn
       Events:

tests/e2e/fixtures/go-httpbin-zip/template.yaml

@@ -56,6 +56,9 @@ Resources:
     Properties:
       PackageType: Image
       MemorySize: 256
+      Environment:
+        Variables:
+          AWS_LWA_ENABLE_COMPRESSION: 'true'
       Events:
         HttpAPIEvent:
           Type: HttpApi

tests/e2e/fixtures/go-httpbin/template.yaml

@@ -1,9 +1,12 @@
+use flate2::read::GzDecoder;
 use http::Uri;
 use hyper::{Body, Client, Method, Request};
 use hyper_tls::HttpsConnector;
 use lambda_http::aws_lambda_events::serde_json;
 use lambda_http::aws_lambda_events::serde_json::Value;
 use std::env;
+use std::io;
+use std::io::prelude::*;
 
 fn get_endpoints() -> Vec<Option<String>> {
     let configurations = [
@@ -20,6 +23,13 @@ fn get_endpoints() -> Vec<Option<String>> {
     configurations.iter().map(|e| env::var(e).ok()).collect()
 }
 
+fn decode_reader(bytes: &Vec<u8>) -> io::Result<String> {
+    let mut gz = GzDecoder::new(&bytes[..]);
+    let mut s = String::new();
+    gz.read_to_string(&mut s)?;
+    Ok(s)
+}
+
 #[ignore]
 #[tokio::test]
 async fn test_http_basic_request() {
@@ -88,3 +98,33 @@ async fn test_http_query_params() {
         }
     }
 }
+
+#[ignore]
+#[tokio::test]
+async fn test_http_compress() {
+    for endpoint in get_endpoints().iter() {
+        if let Some(endpoint) = endpoint {
+            let client = Client::builder().build::<_, hyper::Body>(HttpsConnector::new());
+            let parts = endpoint.parse::<Uri>().unwrap().into_parts();
+            let uri = Uri::builder()
+                .scheme(parts.scheme.unwrap())
+                .authority(parts.authority.unwrap())
+                .path_and_query("/html")
+                .build()
+                .unwrap();
+            let req = Request::builder()
+                .method(Method::GET)
+                .header("accept-encoding", "gzip")
+                .uri(uri)
+                .body(Body::empty())
+                .unwrap();
+            let resp = client.request(req).await.unwrap();
+            let (parts, body) = resp.into_parts();
+            let body_bytes = hyper::body::to_bytes(body).await.unwrap();
+            let body = decode_reader(&body_bytes.to_vec()).unwrap();
+
+            assert_eq!(200, parts.status.as_u16());
+            assert!(body.contains("<html>"));
+        }
+    }
+}