use crate::{ body::{Body, BoxBody}, extract::FromRequest, response::IntoResponse, routing::{BoxResponseBody, EmptyRouter, MethodFilter}, service::HandleError, }; use async_trait::async_trait; use bytes::Bytes; use futures_util::future; use http::{Request, Response}; use std::{ convert::Infallible, future::Future, marker::PhantomData, task::{Context, Poll}, }; use tower::{util::Oneshot, BoxError, Layer, Service, ServiceExt}; pub fn get(handler: H) -> OnMethod, EmptyRouter> where H: Handler, { on(MethodFilter::Get, handler) } pub fn post(handler: H) -> OnMethod, EmptyRouter> where H: Handler, { on(MethodFilter::Post, handler) } pub fn on(method: MethodFilter, handler: H) -> OnMethod, EmptyRouter> where H: Handler, { OnMethod { method, svc: handler.into_service(), fallback: EmptyRouter, } } mod sealed { #![allow(unreachable_pub)] pub trait HiddentTrait {} pub struct Hidden; impl HiddentTrait for Hidden {} } #[async_trait] pub trait Handler: Sized { // This seals the trait. We cannot use the regular "sealed super trait" approach // due to coherence. #[doc(hidden)] type Sealed: sealed::HiddentTrait; async fn call(self, req: Request) -> Response; fn layer(self, layer: L) -> Layered where L: Layer>, { Layered::new(layer.layer(IntoService::new(self))) } fn into_service(self) -> IntoService { IntoService::new(self) } } #[async_trait] impl Handler<()> for F where F: FnOnce(Request) -> Fut + Send + Sync, Fut: Future + Send, Res: IntoResponse, { type Sealed = sealed::Hidden; async fn call(self, req: Request) -> Response { self(req).await.into_response().map(BoxBody::new) } } macro_rules! impl_handler { () => {}; ( $head:ident, $($tail:ident),* $(,)? ) => { #[async_trait] #[allow(non_snake_case)] impl Handler<($head, $($tail,)*)> for F where F: FnOnce(Request, $head, $($tail,)*) -> Fut + Send + Sync, Fut: Future + Send, Res: IntoResponse, $head: FromRequest + Send, $( $tail: FromRequest + Send, )* { type Sealed = sealed::Hidden; async fn call(self, mut req: Request) -> Response { let $head = match $head::from_request(&mut req).await { Ok(value) => value, Err(rejection) => return rejection.into_response().map(BoxBody::new), }; $( let $tail = match $tail::from_request(&mut req).await { Ok(value) => value, Err(rejection) => return rejection.into_response().map(BoxBody::new), }; )* let res = self(req, $head, $($tail,)*).await; res.into_response().map(BoxBody::new) } } impl_handler!($($tail,)*); }; } impl_handler!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16); pub struct Layered { svc: S, _input: PhantomData T>, } impl Clone for Layered where S: Clone, { fn clone(&self) -> Self { Self::new(self.svc.clone()) } } #[async_trait] impl Handler for Layered where S: Service, Response = Response> + Send, // S::Response: IntoResponse, S::Error: IntoResponse, S::Future: Send, B: http_body::Body + Send + Sync + 'static, B::Error: Into + Send + Sync + 'static, { type Sealed = sealed::Hidden; async fn call(self, req: Request) -> Response { match self .svc .oneshot(req) .await .map_err(IntoResponse::into_response) { Ok(res) => res.map(BoxBody::new), Err(res) => res.map(BoxBody::new), } } } impl Layered { pub(crate) fn new(svc: S) -> Self { Self { svc, _input: PhantomData, } } pub fn handle_error(self, f: F) -> Layered, T> where S: Service, Response = Response>, F: FnOnce(S::Error) -> Res, Res: IntoResponse, { let svc = HandleError::new(self.svc, f); Layered::new(svc) } } pub struct IntoService { handler: H, _marker: PhantomData T>, } impl IntoService { fn new(handler: H) -> Self { Self { handler, _marker: PhantomData, } } } impl Clone for IntoService where H: Clone, { fn clone(&self) -> Self { Self { handler: self.handler.clone(), _marker: PhantomData, } } } impl Service> for IntoService where H: Handler + Clone + Send + 'static, { type Response = Response; type Error = Infallible; type Future = future::BoxFuture<'static, Result>; fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { // `IntoService` can only be constructed from async functions which are always ready, or from // `Layered` which bufferes in `::call` and is therefore also always // ready. Poll::Ready(Ok(())) } fn call(&mut self, req: Request) -> Self::Future { let handler = self.handler.clone(); Box::pin(async move { let res = Handler::call(handler, req).await; Ok(res) }) } } #[derive(Clone)] pub struct OnMethod { pub(crate) method: MethodFilter, pub(crate) svc: S, pub(crate) fallback: F, } impl OnMethod { pub fn get(self, handler: H) -> OnMethod, Self> where H: Handler, { self.on(MethodFilter::Get, handler) } pub fn post(self, handler: H) -> OnMethod, Self> where H: Handler, { self.on(MethodFilter::Post, handler) } pub fn on(self, method: MethodFilter, handler: H) -> OnMethod, Self> where H: Handler, { OnMethod { method, svc: handler.into_service(), fallback: self, } } } // this is identical to `routing::OnMethod`'s implementation. Would be nice to find a way to clean // that up, but not sure its possible. impl Service> for OnMethod where S: Service, Response = Response, Error = Infallible> + Clone, SB: http_body::Body + Send + Sync + 'static, SB::Error: Into, F: Service, Response = Response, Error = Infallible> + Clone, FB: http_body::Body + Send + Sync + 'static, FB::Error: Into, { type Response = Response; type Error = Infallible; #[allow(clippy::type_complexity)] type Future = future::Either< BoxResponseBody>>, BoxResponseBody>>, >; fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn call(&mut self, req: Request) -> Self::Future { if self.method.matches(req.method()) { let response_future = self.svc.clone().oneshot(req); future::Either::Left(BoxResponseBody(response_future)) } else { let response_future = self.fallback.clone().oneshot(req); future::Either::Right(BoxResponseBody(response_future)) } } }