tonic/service/
recover_error.rs

1/*
2 *
3 * Copyright 2025 gRPC authors.
4 *
5 * Permission is hereby granted, free of charge, to any person obtaining a copy
6 * of this software and associated documentation files (the "Software"), to
7 * deal in the Software without restriction, including without limitation the
8 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
9 * sell copies of the Software, and to permit persons to whom the Software is
10 * furnished to do so, subject to the following conditions:
11 *
12 * The above copyright notice and this permission notice shall be included in
13 * all copies or substantial portions of the Software.
14 *
15 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
20 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
21 * IN THE SOFTWARE.
22 *
23 */
24
25//! Middleware which recovers from error.
26
27use std::{
28    fmt,
29    future::Future,
30    pin::Pin,
31    task::{Context, Poll, ready},
32};
33
34use http::Response;
35use pin_project::pin_project;
36use tower_layer::Layer;
37use tower_service::Service;
38
39use crate::Status;
40
41/// Layer which applies the [`RecoverError`] middleware.
42#[derive(Debug, Default, Clone)]
43pub struct RecoverErrorLayer {
44    _priv: (),
45}
46
47impl RecoverErrorLayer {
48    /// Create a new `RecoverErrorLayer`.
49    pub fn new() -> Self {
50        Self { _priv: () }
51    }
52}
53
54impl<S> Layer<S> for RecoverErrorLayer {
55    type Service = RecoverError<S>;
56
57    fn layer(&self, inner: S) -> Self::Service {
58        RecoverError::new(inner)
59    }
60}
61
62/// Middleware that attempts to recover from service errors by turning them into a response built
63/// from the `Status`.
64#[derive(Debug, Clone)]
65pub struct RecoverError<S> {
66    inner: S,
67}
68
69impl<S> RecoverError<S> {
70    /// Create a new `RecoverError` middleware.
71    pub fn new(inner: S) -> Self {
72        Self { inner }
73    }
74}
75
76impl<S, Req, ResBody> Service<Req> for RecoverError<S>
77where
78    S: Service<Req, Response = Response<ResBody>>,
79    S::Error: Into<crate::BoxError>,
80{
81    type Response = Response<ResponseBody<ResBody>>;
82    type Error = crate::BoxError;
83    type Future = ResponseFuture<S::Future>;
84
85    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
86        self.inner.poll_ready(cx).map_err(Into::into)
87    }
88
89    fn call(&mut self, req: Req) -> Self::Future {
90        ResponseFuture {
91            inner: self.inner.call(req),
92        }
93    }
94}
95
96/// Response future for [`RecoverError`].
97#[pin_project]
98pub struct ResponseFuture<F> {
99    #[pin]
100    inner: F,
101}
102
103impl<F> fmt::Debug for ResponseFuture<F> {
104    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
105        f.debug_struct("ResponseFuture").finish()
106    }
107}
108
109impl<F, E, ResBody> Future for ResponseFuture<F>
110where
111    F: Future<Output = Result<Response<ResBody>, E>>,
112    E: Into<crate::BoxError>,
113{
114    type Output = Result<Response<ResponseBody<ResBody>>, crate::BoxError>;
115
116    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
117        match ready!(self.project().inner.poll(cx)) {
118            Ok(response) => {
119                let response = response.map(ResponseBody::full);
120                Poll::Ready(Ok(response))
121            }
122            Err(err) => match Status::try_from_error(err.into()) {
123                Ok(status) => {
124                    let (parts, ()) = status.into_http::<()>().into_parts();
125                    let res = Response::from_parts(parts, ResponseBody::empty());
126                    Poll::Ready(Ok(res))
127                }
128                Err(err) => Poll::Ready(Err(err)),
129            },
130        }
131    }
132}
133
134/// Response body for [`RecoverError`].
135#[pin_project]
136pub struct ResponseBody<B> {
137    #[pin]
138    inner: Option<B>,
139}
140
141impl<B> fmt::Debug for ResponseBody<B> {
142    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
143        f.debug_struct("ResponseBody").finish()
144    }
145}
146
147impl<B> ResponseBody<B> {
148    fn full(inner: B) -> Self {
149        Self { inner: Some(inner) }
150    }
151
152    const fn empty() -> Self {
153        Self { inner: None }
154    }
155}
156
157impl<B> http_body::Body for ResponseBody<B>
158where
159    B: http_body::Body,
160{
161    type Data = B::Data;
162    type Error = B::Error;
163
164    fn poll_frame(
165        self: Pin<&mut Self>,
166        cx: &mut Context<'_>,
167    ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
168        match self.project().inner.as_pin_mut() {
169            Some(b) => b.poll_frame(cx),
170            None => Poll::Ready(None),
171        }
172    }
173
174    fn is_end_stream(&self) -> bool {
175        match &self.inner {
176            Some(b) => b.is_end_stream(),
177            None => true,
178        }
179    }
180
181    fn size_hint(&self) -> http_body::SizeHint {
182        match &self.inner {
183            Some(body) => body.size_hint(),
184            None => http_body::SizeHint::with_exact(0),
185        }
186    }
187}
tonic/service/recover_error.rs

tonic/service/
recover_error.rs
