tonic/transport/server/
recover_error.rs

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

tonic/transport/server/
recover_error.rs
