diff --git a/lambda-http/src/lib.rs b/lambda-http/src/lib.rs index caea8618..d76c0b59 100644 --- a/lambda-http/src/lib.rs +++ b/lambda-http/src/lib.rs @@ -88,7 +88,7 @@ pub use crate::{ }; use crate::{ request::{LambdaRequest, RequestOrigin}, - response::LambdaResponse, + response::{BodyConversionError, LambdaResponse}, }; // Reexported in its entirety, regardless of what feature flags are enabled @@ -112,9 +112,7 @@ pub use streaming::{run_with_streaming_response_concurrent, streaming_runtime_co /// Type alias for `http::Request`s with a fixed [`Body`](enum.Body.html) type pub type Request = http::Request; -/// Future that will convert an [`IntoResponse`] into an actual [`LambdaResponse`] -/// -/// This is used by the `Adapter` wrapper and is completely internal to the `lambda_http::run` function. +/// Future used by [`Adapter`] to convert an [`IntoResponse`] into a [`LambdaResponse`]. #[non_exhaustive] #[doc(hidden)] pub enum TransformResponse<'a, R, E> { @@ -146,9 +144,109 @@ where } } +// The public Adapter must preserve its handler's error type. The runtime helpers +// can use Diagnostic as an internal common error channel for conversion failures. +enum RuntimeTransformResponse<'a, R, E> { + Request(RequestOrigin, RequestFuture<'a, R, E>), + Response(RequestOrigin, ResponseFuture), +} + +impl Future for RuntimeTransformResponse<'_, R, E> +where + R: IntoResponse, + E: Into, +{ + type Output = Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll { + match *self { + RuntimeTransformResponse::Request(ref mut origin, ref mut request) => match request.as_mut().poll(cx) { + Poll::Ready(Ok(resp)) => { + *self = RuntimeTransformResponse::Response(origin.clone(), resp.into_response()); + self.poll(cx) + } + Poll::Ready(Err(err)) => Poll::Ready(Err(err.into())), + Poll::Pending => Poll::Pending, + }, + RuntimeTransformResponse::Response(ref mut origin, ref mut response) => match response.as_mut().poll(cx) { + Poll::Ready(mut resp) => { + if let Some(error) = resp.extensions_mut().remove::() { + return Poll::Ready(Err(Diagnostic { + error_type: error.error_type.to_owned(), + error_message: error.error_message, + })); + } + + Poll::Ready(Ok(LambdaResponse::from_response(origin, resp))) + } + Poll::Pending => Poll::Pending, + }, + } + } +} + +struct RuntimeAdapter<'a, R, S> { + service: S, + _phantom_data: PhantomData<&'a R>, +} + +impl<'a, R, S> Clone for RuntimeAdapter<'a, R, S> +where + S: Clone, +{ + fn clone(&self) -> Self { + Self { + service: self.service.clone(), + _phantom_data: PhantomData, + } + } +} + +impl<'a, R, S, E> From for RuntimeAdapter<'a, R, S> +where + S: Service, + S::Future: Send + 'a, + R: IntoResponse, + E: Into, +{ + fn from(service: S) -> Self { + Self { + service, + _phantom_data: PhantomData, + } + } +} + +impl<'a, R, S, E> Service> for RuntimeAdapter<'a, R, S> +where + S: Service, + S::Future: Send + 'a, + R: IntoResponse, + E: Into, +{ + type Response = LambdaResponse; + type Error = Diagnostic; + type Future = RuntimeTransformResponse<'a, R, E>; + + fn poll_ready(&mut self, cx: &mut core::task::Context<'_>) -> core::task::Poll> { + self.service.poll_ready(cx).map_err(Into::into) + } + + fn call(&mut self, req: LambdaEvent) -> Self::Future { + let LambdaEvent { payload, context } = req; + let request_origin = payload.request_origin(); + let mut event: Request = payload.into(); + update_xray_trace_id_header(event.headers_mut(), &context); + let fut = Box::pin(self.service.call(event.with_lambda_context(context))); + + RuntimeTransformResponse::Request(request_origin, fut) + } +} + /// Wraps a `Service` in a `Service>` /// -/// This is completely internal to the `lambda_http::run` function. +/// This adapter preserves the wrapped service's error type. Response body conversion +/// failures are returned as deterministic HTTP 500 responses. #[non_exhaustive] #[doc(hidden)] pub struct Adapter<'a, R, S> { @@ -232,7 +330,7 @@ where R: IntoResponse, E: std::fmt::Debug + Into, { - lambda_runtime::run(Adapter::from(handler)).await + lambda_runtime::run(RuntimeAdapter::from(handler)).await } /// Starts the Lambda Rust runtime and begins polling for events on the [Lambda @@ -265,7 +363,7 @@ where R: IntoResponse + Send + Sync + 'static, E: std::fmt::Debug + Into + Send + 'static, { - lambda_runtime::run_concurrent(Adapter::from(handler)).await + lambda_runtime::run_concurrent(RuntimeAdapter::from(handler)).await } /// Returns a configured [`Runtime`](lambda_runtime::Runtime) wrapping the given @@ -323,7 +421,7 @@ where R: IntoResponse + Send + Sync + 'static, E: std::fmt::Debug + Into + Send + 'static, { - lambda_runtime::Runtime::new(Adapter::from(handler)) + lambda_runtime::Runtime::new(RuntimeAdapter::from(handler)) } /// Returns a configured [`Runtime`](lambda_runtime::Runtime) wrapping the given @@ -355,7 +453,7 @@ where R: IntoResponse + Send + Sync + 'static, E: std::fmt::Debug + Into + Send + 'static, { - lambda_runtime::Runtime::new(Adapter::from(handler)) + lambda_runtime::Runtime::new(RuntimeAdapter::from(handler)) } // In concurrent mode we must use the per-request context. @@ -369,17 +467,35 @@ fn update_xray_trace_id_header(headers: &mut http::HeaderMap, context: &Context) #[cfg(test)] mod test_adapter { - use std::task::{Context, Poll}; + use bytes::Bytes; + use futures_util::stream; + use http_body::Frame; + use http_body_util::StreamBody; + use std::{ + io::{self, ErrorKind}, + task::{Context, Poll}, + }; use crate::{ + aws_lambda_events::apigw::ApiGatewayV2httpRequest, http::{Response, StatusCode}, lambda_runtime::LambdaEvent, request::LambdaRequest, response::LambdaResponse, tower::{util::BoxService, Service, ServiceBuilder, ServiceExt}, - Adapter, Body, Request, + Adapter, Body, Request, RuntimeAdapter, }; + fn fallible_body() -> impl http_body::Body + Unpin { + StreamBody::new(stream::iter([ + Ok(Frame::data(Bytes::from_static(b"partial response"))), + Err(io::Error::new( + ErrorKind::UnexpectedEof, + "simulated truncated response body", + )), + ])) + } + // A middleware that logs requests before forwarding them to another service struct LogService { inner: S, @@ -422,6 +538,32 @@ mod test_adapter { .boxed(); } + #[tokio::test] + async fn runtime_adapter_propagates_body_errors() { + for content_type in ["text/plain; charset=utf-8", "application/octet-stream"] { + let handler = crate::service_fn(move |_event: Request| async move { + Ok::<_, std::convert::Infallible>( + Response::builder() + .header(http::header::CONTENT_TYPE, content_type) + .body(fallible_body()) + .expect("unable to build http::Response"), + ) + }); + let event = LambdaEvent::new( + LambdaRequest::ApiGatewayV2(ApiGatewayV2httpRequest::default()), + crate::Context::default(), + ); + + let error = RuntimeAdapter::from(handler) + .oneshot(event) + .await + .expect_err("body collection error should be propagated"); + + assert_eq!(error.error_type, std::any::type_name::()); + assert!(error.error_message.contains("simulated truncated response body")); + } + } + async fn http_handler(_req: Request) -> Result<&'static str, std::convert::Infallible> { Ok("hello") } diff --git a/lambda-http/src/response.rs b/lambda-http/src/response.rs index da387322..674be274 100644 --- a/lambda-http/src/response.rs +++ b/lambda-http/src/response.rs @@ -20,6 +20,7 @@ use http_body_util::BodyExt; use mime::{Mime, CHARSET}; use serde::Serialize; use std::{ + any::type_name, borrow::Cow, fmt, future::{ready, Future}, @@ -216,7 +217,12 @@ where let (parts, body) = self.into_parts(); let headers = parts.headers.clone(); - let fut = async { Response::from_parts(parts, body.convert(headers).await) }; + let fut = async { + match body.convert(headers).await { + Ok(body) => Response::from_parts(parts, body), + Err(error) => error.into_response(), + } + }; Box::pin(fut) } @@ -328,7 +334,29 @@ impl IntoResponse for (StatusCode, serde_json::Value) { pub type ResponseFuture = Pin> + Send>>; -pub trait ConvertBody { +#[derive(Clone, Debug)] +pub(crate) struct BodyConversionError { + pub(crate) error_type: &'static str, + pub(crate) error_message: String, +} + +impl BodyConversionError { + fn new(error: E) -> Self { + Self { + error_type: type_name::(), + error_message: format!("unable to read bytes from response body: {error:?}"), + } + } + + fn into_response(self) -> Response { + let mut response = Response::new(Body::Empty); + *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR; + response.extensions_mut().insert(self); + response + } +} + +pub(crate) trait ConvertBody { fn convert(self, parts: HeaderMap) -> BodyFuture; } @@ -381,13 +409,8 @@ where B::Error: fmt::Debug, { Box::pin(async move { - Body::from( - body.collect() - .await - .expect("unable to read bytes from body") - .to_bytes() - .to_vec(), - ) + let bytes = body.collect().await.map_err(BodyConversionError::new)?.to_bytes(); + Ok(Body::from(bytes.to_vec())) }) } @@ -409,30 +432,45 @@ where // assumes utf-8 Box::pin(async move { - let bytes = body.collect().await.expect("unable to read bytes from body").to_bytes(); + let bytes = body.collect().await.map_err(BodyConversionError::new)?.to_bytes(); let (content, _, _) = encoding.decode(&bytes); - match content { + Ok(match content { Cow::Borrowed(content) => Body::from(content), Cow::Owned(content) => Body::from(content), - } + }) }) } -pub type BodyFuture = Pin + Send>>; +pub(crate) type BodyFuture = Pin> + Send>>; #[cfg(test)] mod tests { use super::{Body, IntoResponse, LambdaResponse, RequestOrigin, X_LAMBDA_HTTP_CONTENT_ENCODING}; + use bytes::Bytes; + use futures_util::stream; use http::{ header::{CONTENT_ENCODING, CONTENT_TYPE}, Response, StatusCode, }; + use http_body::Frame; + use http_body_util::StreamBody; use lambda_runtime_api_client::body::Body as HyperBody; use serde_json::{self, json}; + use std::io::{self, ErrorKind}; const SVG_LOGO: &str = include_str!("../tests/data/svg_logo.svg"); + fn fallible_body() -> impl http_body::Body + Unpin { + StreamBody::new(stream::iter([ + Ok(Frame::data(Bytes::from_static(b"partial response"))), + Err(io::Error::new( + ErrorKind::UnexpectedEof, + "simulated truncated response body", + )), + ])) + } + #[tokio::test] async fn json_into_response() { let response = json!({ "hello": "lambda"}).into_response().await; @@ -467,6 +505,34 @@ mod tests { } } + #[tokio::test] + async fn fallible_text_body_returns_internal_server_error() { + let response = Response::builder() + .header(CONTENT_TYPE, "text/plain; charset=utf-8") + .body(fallible_body()) + .expect("unable to build http::Response") + .into_response() + .await; + + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + assert!(response.headers().is_empty()); + assert!(matches!(response.body(), Body::Empty)); + } + + #[tokio::test] + async fn fallible_binary_body_returns_internal_server_error() { + let response = Response::builder() + .header(CONTENT_TYPE, "application/octet-stream") + .body(fallible_body()) + .expect("unable to build http::Response") + .into_response() + .await; + + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + assert!(response.headers().is_empty()); + assert!(matches!(response.body(), Body::Empty)); + } + #[tokio::test] async fn json_with_status_code_into_response() { let response = (StatusCode::CREATED, json!({ "hello": "lambda"})).into_response().await;