tonic/transport/server/service/
recover_error.rs

1use crate::Status;
2use http::Response;
3use http_body::Frame;
4use pin_project::pin_project;
5use std::{
6    future::Future,
7    pin::Pin,
8    task::{ready, Context, Poll},
9};
10use tower::Service;
11
12/// Middleware that attempts to recover from service errors by turning them into a response built
13/// from the `Status`.
14#[derive(Debug, Clone)]
15pub(crate) struct RecoverError<S> {
16    inner: S,
17}
18
19impl<S> RecoverError<S> {
20    pub(crate) fn new(inner: S) -> Self {
21        Self { inner }
22    }
23}
24
25impl<S, R, ResBody> Service<R> for RecoverError<S>
26where
27    S: Service<R, Response = Response<ResBody>>,
28    S::Error: Into<crate::Error>,
29{
30    type Response = Response<MaybeEmptyBody<ResBody>>;
31    type Error = crate::Error;
32    type Future = ResponseFuture<S::Future>;
33
34    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
35        self.inner.poll_ready(cx).map_err(Into::into)
36    }
37
38    fn call(&mut self, req: R) -> Self::Future {
39        ResponseFuture {
40            inner: self.inner.call(req),
41        }
42    }
43}
44
45#[pin_project]
46pub(crate) struct ResponseFuture<F> {
47    #[pin]
48    inner: F,
49}
50
51impl<F, E, ResBody> Future for ResponseFuture<F>
52where
53    F: Future<Output = Result<Response<ResBody>, E>>,
54    E: Into<crate::Error>,
55{
56    type Output = Result<Response<MaybeEmptyBody<ResBody>>, crate::Error>;
57
58    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
59        let result: Result<Response<_>, crate::Error> =
60            ready!(self.project().inner.poll(cx)).map_err(Into::into);
61
62        match result {
63            Ok(response) => {
64                let response = response.map(MaybeEmptyBody::full);
65                Poll::Ready(Ok(response))
66            }
67            Err(err) => match Status::try_from_error(err) {
68                Ok(status) => {
69                    let mut res = Response::new(MaybeEmptyBody::empty());
70                    status.add_header(res.headers_mut()).unwrap();
71                    Poll::Ready(Ok(res))
72                }
73                Err(err) => Poll::Ready(Err(err)),
74            },
75        }
76    }
77}
78
79#[pin_project]
80pub(crate) struct MaybeEmptyBody<B> {
81    #[pin]
82    inner: Option<B>,
83}
84
85impl<B> MaybeEmptyBody<B> {
86    fn full(inner: B) -> Self {
87        Self { inner: Some(inner) }
88    }
89
90    fn empty() -> Self {
91        Self { inner: None }
92    }
93}
94
95impl<B> http_body::Body for MaybeEmptyBody<B>
96where
97    B: http_body::Body + Send,
98{
99    type Data = B::Data;
100    type Error = B::Error;
101
102    fn poll_frame(
103        self: Pin<&mut Self>,
104        cx: &mut Context<'_>,
105    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
106        match self.project().inner.as_pin_mut() {
107            Some(b) => b.poll_frame(cx),
108            None => Poll::Ready(None),
109        }
110    }
111
112    fn is_end_stream(&self) -> bool {
113        match &self.inner {
114            Some(b) => b.is_end_stream(),
115            None => true,
116        }
117    }
118
119    fn size_hint(&self) -> http_body::SizeHint {
120        match &self.inner {
121            Some(body) => body.size_hint(),
122            None => http_body::SizeHint::with_exact(0),
123        }
124    }
125}
tonic/transport/server/service/recover_error.rs

tonic/transport/server/service/
recover_error.rs
