use self::{ body::Body, routing::{AlwaysNotFound, RouteAt}, }; use body::BoxBody; 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::{BoxError, Service}; pub mod body; pub mod extract; pub mod handler; pub mod response; pub mod routing; pub use tower_http::add_extension::{AddExtension, AddExtensionLayer}; pub use async_trait::async_trait; #[cfg(test)] mod tests; pub fn app() -> App { App { service_tree: AlwaysNotFound(()), } } #[derive(Debug, Clone)] pub struct App { service_tree: R, } impl App { pub fn at(self, route_spec: &str) -> RouteAt { self.at_bytes(Bytes::copy_from_slice(route_spec.as_bytes())) } fn at_bytes(self, route_spec: Bytes) -> RouteAt { RouteAt { app: self, route_spec, } } } pub struct IntoService { app: App, } impl Clone for IntoService where R: Clone, { fn clone(&self) -> Self { Self { app: self.app.clone(), } } } impl Service for IntoService where R: Service, Error = Infallible>, B: Default, { type Response = Response; type Error = Infallible; type Future = R::Future; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { match ready!(self.app.service_tree.poll_ready(cx)) { Ok(_) => Poll::Ready(Ok(())), Err(err) => match err {}, } } fn call(&mut self, req: T) -> Self::Future { self.app.service_tree.call(req) } } pub(crate) trait ResultExt { fn unwrap_infallible(self) -> T; } impl ResultExt for Result { fn unwrap_infallible(self) -> T { match self { Ok(value) => value, Err(err) => match err {}, } } } // work around for `BoxError` not implementing `std::error::Error` // // This is currently required since tower-http's Compression middleware's body type's // error only implements error when the inner error type does: // https://github.com/tower-rs/tower-http/blob/master/tower-http/src/lib.rs#L310 // // Fixing that is a breaking change to tower-http so we should wait a bit, but should // totally fix it at some point. #[derive(Debug, thiserror::Error)] #[error("{0}")] pub struct BoxStdError(#[source] pub(crate) tower::BoxError); pub trait ServiceExt: Service, Response = Response> { fn handle_error(self, f: F) -> HandleError where Self: Sized, F: FnOnce(Self::Error) -> Response, B: http_body::Body + Send + Sync + 'static, B::Error: Into + Send + Sync + 'static, NewBody: http_body::Body + Send + Sync + 'static, NewBody::Error: Into + Send + Sync + 'static, { HandleError { inner: self, f, poll_ready_error: None, } } } impl ServiceExt for S where S: Service, Response = Response> {} pub struct HandleError { inner: S, f: F, poll_ready_error: Option, } impl fmt::Debug for HandleError where S: fmt::Debug, E: 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::())) .field("poll_ready_error", &self.poll_ready_error) .finish() } } impl Clone for HandleError where S: Clone, F: Clone, { fn clone(&self) -> Self { Self { inner: self.inner.clone(), f: self.f.clone(), poll_ready_error: None, } } } impl Service> for HandleError where S: Service, Response = Response>, F: FnOnce(S::Error) -> Response + Clone, B: http_body::Body + Send + Sync + 'static, B::Error: Into + Send + Sync + 'static, NewBody: http_body::Body + Send + Sync + 'static, NewBody::Error: Into + Send + Sync + 'static, { type Response = Response; type Error = Infallible; type Future = HandleErrorFuture; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { match ready!(self.inner.poll_ready(cx)) { Ok(_) => Poll::Ready(Ok(())), Err(err) => { self.poll_ready_error = Some(err); Poll::Ready(Ok(())) } } } fn call(&mut self, req: Request) -> Self::Future { if let Some(err) = self.poll_ready_error.take() { return HandleErrorFuture { f: Some(self.f.clone()), kind: Kind::Error(Some(err)), }; } HandleErrorFuture { f: Some(self.f.clone()), kind: Kind::Future(self.inner.call(req)), } } } #[pin_project] pub struct HandleErrorFuture { #[pin] kind: Kind, f: Option, } #[pin_project(project = KindProj)] enum Kind { Future(#[pin] Fut), Error(Option), } impl Future for HandleErrorFuture where Fut: Future, E>>, F: FnOnce(E) -> Response, B: http_body::Body + Send + Sync + 'static, B::Error: Into + Send + Sync + 'static, NewBody: http_body::Body + Send + Sync + 'static, NewBody::Error: Into + Send + Sync + 'static, { type Output = Result, Infallible>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); match this.kind.project() { KindProj::Future(future) => match ready!(future.poll(cx)) { Ok(res) => Ok(res.map(BoxBody::new)).into(), Err(err) => { let f = this.f.take().unwrap(); let res = f(err); Ok(res.map(BoxBody::new)).into() } }, KindProj::Error(err) => { let f = this.f.take().unwrap(); let res = f(err.take().unwrap()); Ok(res.map(BoxBody::new)).into() } } } }