Skip to content

Commit 5d3c3a4

Browse files
IamPritamAcharyaIamPritamAcharya
committed
fix(lambda-http): handle fallible response body errors
1 parent 831d481 commit 5d3c3a4

2 files changed

Lines changed: 232 additions & 24 deletions

File tree

‎lambda-http/src/lib.rs‎

Lines changed: 153 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ pub use crate::{
8888
};
8989
use crate::{
9090
request::{LambdaRequest, RequestOrigin},
91-
response::LambdaResponse,
91+
response::{BodyConversionError, LambdaResponse},
9292
};
9393

9494
// 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
112112
/// Type alias for `http::Request`s with a fixed [`Body`](enum.Body.html) type
113113
pub type Request = http::Request<Body>;
114114

115-
/// Future that will convert an [`IntoResponse`] into an actual [`LambdaResponse`]
116-
///
117-
/// This is used by the `Adapter` wrapper and is completely internal to the `lambda_http::run` function.
115+
/// Future used by [`Adapter`] to convert an [`IntoResponse`] into a [`LambdaResponse`].
118116
#[non_exhaustive]
119117
#[doc(hidden)]
120118
pub enum TransformResponse<'a, R, E> {
@@ -146,9 +144,109 @@ where
146144
}
147145
}
148146

147+
// The public Adapter must preserve its handler's error type. The runtime helpers
148+
// can use Diagnostic as an internal common error channel for conversion failures.
149+
enum RuntimeTransformResponse<'a, R, E> {
150+
Request(RequestOrigin, RequestFuture<'a, R, E>),
151+
Response(RequestOrigin, ResponseFuture),
152+
}
153+
154+
impl<R, E> Future for RuntimeTransformResponse<'_, R, E>
155+
where
156+
R: IntoResponse,
157+
E: Into<Diagnostic>,
158+
{
159+
type Output = Result<LambdaResponse, Diagnostic>;
160+
161+
fn poll(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
162+
match *self {
163+
RuntimeTransformResponse::Request(ref mut origin, ref mut request) => match request.as_mut().poll(cx) {
164+
Poll::Ready(Ok(resp)) => {
165+
*self = RuntimeTransformResponse::Response(origin.clone(), resp.into_response());
166+
self.poll(cx)
167+
}
168+
Poll::Ready(Err(err)) => Poll::Ready(Err(err.into())),
169+
Poll::Pending => Poll::Pending,
170+
},
171+
RuntimeTransformResponse::Response(ref mut origin, ref mut response) => match response.as_mut().poll(cx) {
172+
Poll::Ready(mut resp) => {
173+
if let Some(error) = resp.extensions_mut().remove::<BodyConversionError>() {
174+
return Poll::Ready(Err(Diagnostic {
175+
error_type: error.error_type.to_owned(),
176+
error_message: error.error_message,
177+
}));
178+
}
179+
180+
Poll::Ready(Ok(LambdaResponse::from_response(origin, resp)))
181+
}
182+
Poll::Pending => Poll::Pending,
183+
},
184+
}
185+
}
186+
}
187+
188+
struct RuntimeAdapter<'a, R, S> {
189+
service: S,
190+
_phantom_data: PhantomData<&'a R>,
191+
}
192+
193+
impl<'a, R, S> Clone for RuntimeAdapter<'a, R, S>
194+
where
195+
S: Clone,
196+
{
197+
fn clone(&self) -> Self {
198+
Self {
199+
service: self.service.clone(),
200+
_phantom_data: PhantomData,
201+
}
202+
}
203+
}
204+
205+
impl<'a, R, S, E> From<S> for RuntimeAdapter<'a, R, S>
206+
where
207+
S: Service<Request, Response = R, Error = E>,
208+
S::Future: Send + 'a,
209+
R: IntoResponse,
210+
E: Into<Diagnostic>,
211+
{
212+
fn from(service: S) -> Self {
213+
Self {
214+
service,
215+
_phantom_data: PhantomData,
216+
}
217+
}
218+
}
219+
220+
impl<'a, R, S, E> Service<LambdaEvent<LambdaRequest>> for RuntimeAdapter<'a, R, S>
221+
where
222+
S: Service<Request, Response = R, Error = E>,
223+
S::Future: Send + 'a,
224+
R: IntoResponse,
225+
E: Into<Diagnostic>,
226+
{
227+
type Response = LambdaResponse;
228+
type Error = Diagnostic;
229+
type Future = RuntimeTransformResponse<'a, R, E>;
230+
231+
fn poll_ready(&mut self, cx: &mut core::task::Context<'_>) -> core::task::Poll<Result<(), Self::Error>> {
232+
self.service.poll_ready(cx).map_err(Into::into)
233+
}
234+
235+
fn call(&mut self, req: LambdaEvent<LambdaRequest>) -> Self::Future {
236+
let LambdaEvent { payload, context } = req;
237+
let request_origin = payload.request_origin();
238+
let mut event: Request = payload.into();
239+
update_xray_trace_id_header(event.headers_mut(), &context);
240+
let fut = Box::pin(self.service.call(event.with_lambda_context(context)));
241+
242+
RuntimeTransformResponse::Request(request_origin, fut)
243+
}
244+
}
245+
149246
/// Wraps a `Service<Request>` in a `Service<LambdaEvent<Request>>`
150247
///
151-
/// This is completely internal to the `lambda_http::run` function.
248+
/// This adapter preserves the wrapped service's error type. Response body conversion
249+
/// failures are returned as deterministic HTTP 500 responses.
152250
#[non_exhaustive]
153251
#[doc(hidden)]
154252
pub struct Adapter<'a, R, S> {
@@ -232,7 +330,7 @@ where
232330
R: IntoResponse,
233331
E: std::fmt::Debug + Into<Diagnostic>,
234332
{
235-
lambda_runtime::run(Adapter::from(handler)).await
333+
lambda_runtime::run(RuntimeAdapter::from(handler)).await
236334
}
237335

238336
/// Starts the Lambda Rust runtime and begins polling for events on the [Lambda
@@ -265,7 +363,7 @@ where
265363
R: IntoResponse + Send + Sync + 'static,
266364
E: std::fmt::Debug + Into<Diagnostic> + Send + 'static,
267365
{
268-
lambda_runtime::run_concurrent(Adapter::from(handler)).await
366+
lambda_runtime::run_concurrent(RuntimeAdapter::from(handler)).await
269367
}
270368

271369
/// Returns a configured [`Runtime`](lambda_runtime::Runtime) wrapping the given
@@ -323,7 +421,7 @@ where
323421
R: IntoResponse + Send + Sync + 'static,
324422
E: std::fmt::Debug + Into<Diagnostic> + Send + 'static,
325423
{
326-
lambda_runtime::Runtime::new(Adapter::from(handler))
424+
lambda_runtime::Runtime::new(RuntimeAdapter::from(handler))
327425
}
328426

329427
/// Returns a configured [`Runtime`](lambda_runtime::Runtime) wrapping the given
@@ -355,7 +453,7 @@ where
355453
R: IntoResponse + Send + Sync + 'static,
356454
E: std::fmt::Debug + Into<Diagnostic> + Send + 'static,
357455
{
358-
lambda_runtime::Runtime::new(Adapter::from(handler))
456+
lambda_runtime::Runtime::new(RuntimeAdapter::from(handler))
359457
}
360458

361459
// 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)
369467

370468
#[cfg(test)]
371469
mod test_adapter {
372-
use std::task::{Context, Poll};
470+
use bytes::Bytes;
471+
use futures_util::stream;
472+
use http_body::Frame;
473+
use http_body_util::StreamBody;
474+
use std::{
475+
io::{self, ErrorKind},
476+
task::{Context, Poll},
477+
};
373478

374479
use crate::{
480+
aws_lambda_events::apigw::ApiGatewayV2httpRequest,
375481
http::{Response, StatusCode},
376482
lambda_runtime::LambdaEvent,
377483
request::LambdaRequest,
378484
response::LambdaResponse,
379485
tower::{util::BoxService, Service, ServiceBuilder, ServiceExt},
380-
Adapter, Body, Request,
486+
Adapter, Body, Request, RuntimeAdapter,
381487
};
382488

489+
fn fallible_body() -> impl http_body::Body<Data = Bytes, Error = io::Error> + Unpin {
490+
StreamBody::new(stream::iter([
491+
Ok(Frame::data(Bytes::from_static(b"partial response"))),
492+
Err(io::Error::new(
493+
ErrorKind::UnexpectedEof,
494+
"simulated truncated response body",
495+
)),
496+
]))
497+
}
498+
383499
// A middleware that logs requests before forwarding them to another service
384500
struct LogService<S> {
385501
inner: S,
@@ -422,6 +538,32 @@ mod test_adapter {
422538
.boxed();
423539
}
424540

541+
#[tokio::test]
542+
async fn runtime_adapter_propagates_body_errors() {
543+
for content_type in ["text/plain; charset=utf-8", "application/octet-stream"] {
544+
let handler = crate::service_fn(move |_event: Request| async move {
545+
Ok::<_, std::convert::Infallible>(
546+
Response::builder()
547+
.header(http::header::CONTENT_TYPE, content_type)
548+
.body(fallible_body())
549+
.expect("unable to build http::Response"),
550+
)
551+
});
552+
let event = LambdaEvent::new(
553+
LambdaRequest::ApiGatewayV2(ApiGatewayV2httpRequest::default()),
554+
crate::Context::default(),
555+
);
556+
557+
let error = RuntimeAdapter::from(handler)
558+
.oneshot(event)
559+
.await
560+
.expect_err("body collection error should be propagated");
561+
562+
assert_eq!(error.error_type, std::any::type_name::<io::Error>());
563+
assert!(error.error_message.contains("simulated truncated response body"));
564+
}
565+
}
566+
425567
async fn http_handler(_req: Request) -> Result<&'static str, std::convert::Infallible> {
426568
Ok("hello")
427569
}

‎lambda-http/src/response.rs‎

Lines changed: 79 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ use http_body_util::BodyExt;
2020
use mime::{Mime, CHARSET};
2121
use serde::Serialize;
2222
use std::{
23+
any::type_name,
2324
borrow::Cow,
2425
fmt,
2526
future::{ready, Future},
@@ -216,7 +217,12 @@ where
216217
let (parts, body) = self.into_parts();
217218
let headers = parts.headers.clone();
218219

219-
let fut = async { Response::from_parts(parts, body.convert(headers).await) };
220+
let fut = async {
221+
match body.convert(headers).await {
222+
Ok(body) => Response::from_parts(parts, body),
223+
Err(error) => error.into_response(),
224+
}
225+
};
220226

221227
Box::pin(fut)
222228
}
@@ -328,7 +334,29 @@ impl IntoResponse for (StatusCode, serde_json::Value) {
328334

329335
pub type ResponseFuture = Pin<Box<dyn Future<Output = Response<Body>> + Send>>;
330336

331-
pub trait ConvertBody {
337+
#[derive(Clone, Debug)]
338+
pub(crate) struct BodyConversionError {
339+
pub(crate) error_type: &'static str,
340+
pub(crate) error_message: String,
341+
}
342+
343+
impl BodyConversionError {
344+
fn new<E: fmt::Debug>(error: E) -> Self {
345+
Self {
346+
error_type: type_name::<E>(),
347+
error_message: format!("unable to read bytes from response body: {error:?}"),
348+
}
349+
}
350+
351+
fn into_response(self) -> Response<Body> {
352+
let mut response = Response::new(Body::Empty);
353+
*response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
354+
response.extensions_mut().insert(self);
355+
response
356+
}
357+
}
358+
359+
pub(crate) trait ConvertBody {
332360
fn convert(self, parts: HeaderMap) -> BodyFuture;
333361
}
334362

@@ -381,13 +409,8 @@ where
381409
B::Error: fmt::Debug,
382410
{
383411
Box::pin(async move {
384-
Body::from(
385-
body.collect()
386-
.await
387-
.expect("unable to read bytes from body")
388-
.to_bytes()
389-
.to_vec(),
390-
)
412+
let bytes = body.collect().await.map_err(BodyConversionError::new)?.to_bytes();
413+
Ok(Body::from(bytes.to_vec()))
391414
})
392415
}
393416

@@ -409,30 +432,45 @@ where
409432

410433
// assumes utf-8
411434
Box::pin(async move {
412-
let bytes = body.collect().await.expect("unable to read bytes from body").to_bytes();
435+
let bytes = body.collect().await.map_err(BodyConversionError::new)?.to_bytes();
413436
let (content, _, _) = encoding.decode(&bytes);
414437

415-
match content {
438+
Ok(match content {
416439
Cow::Borrowed(content) => Body::from(content),
417440
Cow::Owned(content) => Body::from(content),
418-
}
441+
})
419442
})
420443
}
421444

422-
pub type BodyFuture = Pin<Box<dyn Future<Output = Body> + Send>>;
445+
pub(crate) type BodyFuture = Pin<Box<dyn Future<Output = Result<Body, BodyConversionError>> + Send>>;
423446

424447
#[cfg(test)]
425448
mod tests {
426449
use super::{Body, IntoResponse, LambdaResponse, RequestOrigin, X_LAMBDA_HTTP_CONTENT_ENCODING};
450+
use bytes::Bytes;
451+
use futures_util::stream;
427452
use http::{
428453
header::{CONTENT_ENCODING, CONTENT_TYPE},
429454
Response, StatusCode,
430455
};
456+
use http_body::Frame;
457+
use http_body_util::StreamBody;
431458
use lambda_runtime_api_client::body::Body as HyperBody;
432459
use serde_json::{self, json};
460+
use std::io::{self, ErrorKind};
433461

434462
const SVG_LOGO: &str = include_str!("../tests/data/svg_logo.svg");
435463

464+
fn fallible_body() -> impl http_body::Body<Data = Bytes, Error = io::Error> + Unpin {
465+
StreamBody::new(stream::iter([
466+
Ok(Frame::data(Bytes::from_static(b"partial response"))),
467+
Err(io::Error::new(
468+
ErrorKind::UnexpectedEof,
469+
"simulated truncated response body",
470+
)),
471+
]))
472+
}
473+
436474
#[tokio::test]
437475
async fn json_into_response() {
438476
let response = json!({ "hello": "lambda"}).into_response().await;
@@ -467,6 +505,34 @@ mod tests {
467505
}
468506
}
469507

508+
#[tokio::test]
509+
async fn fallible_text_body_returns_internal_server_error() {
510+
let response = Response::builder()
511+
.header(CONTENT_TYPE, "text/plain; charset=utf-8")
512+
.body(fallible_body())
513+
.expect("unable to build http::Response")
514+
.into_response()
515+
.await;
516+
517+
assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
518+
assert!(response.headers().is_empty());
519+
assert!(matches!(response.body(), Body::Empty));
520+
}
521+
522+
#[tokio::test]
523+
async fn fallible_binary_body_returns_internal_server_error() {
524+
let response = Response::builder()
525+
.header(CONTENT_TYPE, "application/octet-stream")
526+
.body(fallible_body())
527+
.expect("unable to build http::Response")
528+
.into_response()
529+
.await;
530+
531+
assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
532+
assert!(response.headers().is_empty());
533+
assert!(matches!(response.body(), Body::Empty));
534+
}
535+
470536
#[tokio::test]
471537
async fn json_with_status_code_into_response() {
472538
let response = (StatusCode::CREATED, json!({ "hello": "lambda"})).into_response().await;

0 commit comments

Comments
 (0)