tower_http/limit/
future.rs

1use super::body::create_error_response;
2use super::ResponseBody;
3use http::Response;
4use http_body::Body;
5use pin_project_lite::pin_project;
6use std::future::Future;
7use std::pin::Pin;
8use std::task::{ready, Context, Poll};
9
10pin_project! {
11    /// Response future for [`RequestBodyLimit`].
12    ///
13    /// [`RequestBodyLimit`]: super::RequestBodyLimit
14    pub struct ResponseFuture<F> {
15        #[pin]
16        inner: ResponseFutureInner<F>,
17    }
18}
19
20impl<F> ResponseFuture<F> {
21    pub(crate) fn payload_too_large() -> Self {
22        Self {
23            inner: ResponseFutureInner::PayloadTooLarge,
24        }
25    }
26
27    pub(crate) fn new(future: F) -> Self {
28        Self {
29            inner: ResponseFutureInner::Future { future },
30        }
31    }
32}
33
34pin_project! {
35    #[project = ResFutProj]
36    enum ResponseFutureInner<F> {
37        PayloadTooLarge,
38        Future {
39            #[pin]
40            future: F,
41        }
42    }
43}
44
45impl<ResBody, F, E> Future for ResponseFuture<F>
46where
47    ResBody: Body,
48    F: Future<Output = Result<Response<ResBody>, E>>,
49{
50    type Output = Result<Response<ResponseBody<ResBody>>, E>;
51
52    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
53        let res = match self.project().inner.project() {
54            ResFutProj::PayloadTooLarge => create_error_response(),
55            ResFutProj::Future { future } => ready!(future.poll(cx))?.map(ResponseBody::new),
56        };
57
58        Poll::Ready(Ok(res))
59    }
60}