use crate::{ body::{Body, BoxBody}, response::IntoResponse, routing::{EmptyRouter, MethodFilter, OnMethod}, }; use bytes::Bytes; use futures_util::ready; use http::{Request, Response}; use pin_project::pin_project; use std::{ convert::Infallible, fmt, future::Future, pin::Pin, task::{Context, Poll}, }; use tower::{util::Oneshot, BoxError, Service, ServiceExt as _}; pub fn get(svc: S) -> OnMethod { on(MethodFilter::Get, svc) } pub fn post(svc: S) -> OnMethod { on(MethodFilter::Post, svc) } pub fn on(method: MethodFilter, svc: S) -> OnMethod { OnMethod { method, svc, fallback: EmptyRouter, } } #[derive(Clone)] pub struct HandleError { inner: S, f: F, } impl HandleError { pub(crate) fn new(inner: S, f: F) -> Self { Self { inner, f } } } impl fmt::Debug for HandleError where S: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("HandleError") .field("inner", &self.inner) .field("f", &format_args!("{}", std::any::type_name::())) .finish() } } impl Service> for HandleError where S: Service, Response = Response> + Clone, F: FnOnce(S::Error) -> Res + Clone, Res: IntoResponse, B: http_body::Body + Send + Sync + 'static, B::Error: Into + Send + Sync + 'static, { type Response = Response; type Error = Infallible; type Future = HandleErrorFuture>, F>; fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn call(&mut self, req: Request) -> Self::Future { HandleErrorFuture { f: Some(self.f.clone()), inner: self.inner.clone().oneshot(req), } } } #[pin_project] pub struct HandleErrorFuture { #[pin] inner: Fut, f: Option, } impl Future for HandleErrorFuture where Fut: Future, E>>, F: FnOnce(E) -> Res, Res: IntoResponse, B: http_body::Body + Send + Sync + 'static, B::Error: Into + Send + Sync + 'static, { type Output = Result, Infallible>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); match ready!(this.inner.poll(cx)) { Ok(res) => Ok(res.map(BoxBody::new)).into(), Err(err) => { let f = this.f.take().unwrap(); let res = f(err).into_response(); Ok(res.map(BoxBody::new)).into() } } } }