feat(lambda-http): accept http_body::Body in responses (#491)

hugobast · nmoutschen · calavera · web-flow · commit 8e6ae59874d5 · 2022-07-14T07:00:02.000+08:00
* feat!(lambda-http): accept http_body::Body in responses

Co-authored-by: Nicolas Moutschen &lt;nicolas.moutschen@gmail.com&gt;
Co-authored-by: David Calavera &lt;david.calavera@gmail.com&gt;
diff --git a/Makefile b/Makefile
@@ -6,6 +6,10 @@ INTEG_EXTENSIONS := extension-fn extension-trait logs-trait
 # Using musl to run extensions on both AL1 and AL2
 INTEG_ARCH := x86_64-unknown-linux-musl
 
+define uppercase
+$(shell sed -r 's/(^|-)(\w)/\U\2/g' <<< $(1))
+endef
+
 pr-check:
 	cargo +1.54.0 check --all
 	cargo +stable fmt --all -- --check
@@ -15,7 +19,7 @@ pr-check:
 
 integration-tests:
 # Build Integration functions
-	cross build --release --target $(INTEG_ARCH) -p lambda_integration_tests
+	cargo zigbuild --release --target $(INTEG_ARCH) -p lambda_integration_tests
 	rm -rf ./build
 	mkdir -p ./build
 	${MAKE} ${MAKEOPTS} $(foreach function,${INTEG_FUNCTIONS_BUILD}, build-integration-function-${function})
@@ -37,7 +41,7 @@ build-integration-function-%:
 
 build-integration-extension-%:
 	mkdir -p ./build/$*/extensions
-	cp -v ./target/$(INTEG_ARCH)/release/$* ./build/$*/extensions/$*
+	cp -v ./target/$(INTEG_ARCH)/release/$* ./build/$*/extensions/$(call uppercase,$*)
 
 invoke-integration-function-%:
 	aws lambda invoke --function-name $$(aws cloudformation describe-stacks --stack-name $(INTEG_STACK_NAME) \
@@ -56,4 +60,3 @@ invoke-integration-api-%:
 	curl -X POST -d '{"command": "hello"}' $(API_URL)/trait/post
 	curl -X POST -d '{"command": "hello"}' $(API_URL)/al2/post
 	curl -X POST -d '{"command": "hello"}' $(API_URL)/al2-trait/post
-	
diff --git a/examples/http-basic-lambda/src/main.rs b/examples/http-basic-lambda/src/main.rs
@@ -1,18 +1,18 @@
-use lambda_http::{run, service_fn, Error, IntoResponse, Request, Response};
+use lambda_http::{run, service_fn, Body, Error, Request, Response};
 
 /// This is the main body for the function.
 /// Write your code inside it.
 /// There are some code examples in the Runtime repository:
 /// - https://github.com/awslabs/aws-lambda-rust-runtime/tree/main/examples
-async fn function_handler(_event: Request) -> Result<impl IntoResponse, Error> {
+async fn function_handler(_event: Request) -> Result<Response<Body>, Error> {
     // Extract some useful information from the request
 
     // Return something that implements IntoResponse.
     // It will be serialized to the right response event automatically by the runtime
     let resp = Response::builder()
         .status(200)
         .header("content-type", "text/html")
-        .body("Hello AWS Lambda HTTP request")
+        .body("Hello AWS Lambda HTTP request".into())
         .map_err(Box::new)?;
     Ok(resp)
 }
diff --git a/examples/http-cors/src/main.rs b/examples/http-cors/src/main.rs
@@ -29,7 +29,7 @@ async fn main() -> Result<(), Error> {
 
 async fn func(event: Request) -> Result<Response<Body>, Error> {
     Ok(match event.query_string_parameters().first("first_name") {
-        Some(first_name) => format!("Hello, {}!", first_name).into_response(),
+        Some(first_name) => format!("Hello, {}!", first_name).into_response().await,
         _ => Response::builder()
             .status(400)
             .body("Empty first name".into())
diff --git a/examples/http-query-parameters/src/main.rs b/examples/http-query-parameters/src/main.rs
@@ -7,7 +7,7 @@ use lambda_http::{run, service_fn, Error, IntoResponse, Request, RequestExt, Res
 async fn function_handler(event: Request) -> Result<impl IntoResponse, Error> {
     // Extract some useful information from the request
     Ok(match event.query_string_parameters().first("first_name") {
-        Some(first_name) => format!("Hello, {}!", first_name).into_response(),
+        Some(first_name) => format!("Hello, {}!", first_name).into_response().await,
         _ => Response::builder()
             .status(400)
             .body("Empty first name".into())
diff --git a/examples/http-raw-path/src/main.rs b/examples/http-raw-path/src/main.rs
@@ -15,7 +15,9 @@ async fn main() -> Result<(), Error> {
 }
 
 async fn func(event: Request) -> Result<impl IntoResponse, Error> {
-    let res = format!("The raw path for this request is: {}", event.raw_http_path()).into_response();
+    let res = format!("The raw path for this request is: {}", event.raw_http_path())
+        .into_response()
+        .await;
 
     Ok(res)
 }
diff --git a/examples/http-shared-resource/src/main.rs b/examples/http-shared-resource/src/main.rs
@@ -29,9 +29,12 @@ async fn main() -> Result<(), Error> {
     // Define a closure here that makes use of the shared client.
     let handler_func_closure = move |event: Request| async move {
         Result::<Response<Body>, Error>::Ok(match event.query_string_parameters().first("first_name") {
-            Some(first_name) => shared_client_ref
-                .response(event.lambda_context().request_id, first_name)
-                .into_response(),
+            Some(first_name) => {
+                shared_client_ref
+                    .response(event.lambda_context().request_id, first_name)
+                    .into_response()
+                    .await
+            }
             _ => Response::builder()
                 .status(400)
                 .body("Empty first name".into())
diff --git a/lambda-http/Cargo.toml b/lambda-http/Cargo.toml
@@ -24,11 +24,14 @@ base64 = "0.13.0"
 bytes = "1"
 http = "0.2"
 http-body = "0.4"
+hyper = "0.14"
 lambda_runtime = { path = "../lambda-runtime", version = "0.5" }
 serde = { version = "^1", features = ["derive"] }
 serde_json = "^1"
 serde_urlencoded = "0.7.0"
 query_map = { version = "0.5", features = ["url-query"] }
+mime = "0.3.16"
+encoding_rs = "0.8.31"
 
 [dependencies.aws_lambda_events]
 version = "^0.6.3"
diff --git a/lambda-http/src/lib.rs b/lambda-http/src/lib.rs
@@ -67,6 +67,8 @@ extern crate maplit;
 pub use http::{self, Response};
 use lambda_runtime::LambdaEvent;
 pub use lambda_runtime::{self, service_fn, tower, Context, Error, Service};
+use request::RequestFuture;
+use response::ResponseFuture;
 
 pub mod ext;
 pub mod request;
@@ -91,9 +93,9 @@ pub type Request = http::Request<Body>;
 ///
 /// This is used by the `Adapter` wrapper and is completely internal to the `lambda_http::run` function.
 #[doc(hidden)]
-pub struct TransformResponse<'a, R, E> {
-    request_origin: RequestOrigin,
-    fut: Pin<Box<dyn Future<Output = Result<R, E>> + 'a>>,
+pub enum TransformResponse<'a, R, E> {
+    Request(RequestOrigin, RequestFuture<'a, R, E>),
+    Response(RequestOrigin, ResponseFuture),
 }
 
 impl<'a, R, E> Future for TransformResponse<'a, R, E>
@@ -103,11 +105,19 @@ where
     type Output = Result<LambdaResponse, E>;
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut TaskContext) -> Poll<Self::Output> {
-        match self.fut.as_mut().poll(cx) {
-            Poll::Ready(result) => Poll::Ready(
-                result.map(|resp| LambdaResponse::from_response(&self.request_origin, resp.into_response())),
-            ),
-            Poll::Pending => Poll::Pending,
+        match *self {
+            TransformResponse::Request(ref mut origin, ref mut request) => match request.as_mut().poll(cx) {
+                Poll::Ready(Ok(resp)) => {
+                    *self = TransformResponse::Response(origin.clone(), resp.into_response());
+                    self.poll(cx)
+                }
+                Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
+                Poll::Pending => Poll::Pending,
+            },
+            TransformResponse::Response(ref mut origin, ref mut response) => match response.as_mut().poll(cx) {
+                Poll::Ready(resp) => Poll::Ready(Ok(LambdaResponse::from_response(origin, resp))),
+                Poll::Pending => Poll::Pending,
+            },
         }
     }
 }
@@ -153,7 +163,8 @@ where
         let request_origin = req.payload.request_origin();
         let event: Request = req.payload.into();
         let fut = Box::pin(self.service.call(event.with_lambda_context(req.context)));
-        TransformResponse { request_origin, fut }
+
+        TransformResponse::Request(request_origin, fut)
     }
 }
 
diff --git a/lambda-http/src/request.rs b/lambda-http/src/request.rs
@@ -17,6 +17,8 @@ use http::header::HeaderName;
 use query_map::QueryMap;
 use serde::Deserialize;
 use serde_json::error::Error as JsonError;
+use std::future::Future;
+use std::pin::Pin;
 use std::{io::Read, mem};
 
 /// Internal representation of an Lambda http event from
@@ -56,9 +58,12 @@ impl LambdaRequest {
     }
 }
 
+/// RequestFuture type
+pub type RequestFuture<'a, R, E> = Pin<Box<dyn Future<Output = Result<R, E>> + 'a>>;
+
 /// Represents the origin from which the lambda was requested from.
 #[doc(hidden)]
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum RequestOrigin {
     /// API Gateway request origin
     #[cfg(feature = "apigw_rest")]
diff --git a/lambda-http/src/response.rs b/lambda-http/src/response.rs
diff --git a/lambda-integration-tests/src/bin/http-fn.rs b/lambda-integration-tests/src/bin/http-fn.rs
diff --git a/lambda-integration-tests/src/bin/http-trait.rs b/lambda-integration-tests/src/bin/http-trait.rs

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,9 @@ async fn main() -> Result<(), Error> {`
`15`	`15`	`}`
`16`	`16`
`17`	`17`	`async fn func(event: Request) -> Result<impl IntoResponse, Error> {`
`18`		`- let res = format!("The raw path for this request is: {}", event.raw_http_path()).into_response();`
	`18`	`+ let res = format!("The raw path for this request is: {}", event.raw_http_path())`
	`19`	`+ .into_response()`
	`20`	`+ .await;`
`19`	`21`
`20`	`22`	`Ok(res)`
`21`	`23`	`}`