Files
axum/src/handler.rs
T

306 lines
7.9 KiB
Rust
Raw Normal View History

2021-06-06 11:37:08 +02:00
use crate::{
2021-06-06 15:19:54 +02:00
body::{Body, BoxBody},
2021-06-06 11:37:08 +02:00
extract::FromRequest,
response::IntoResponse,
2021-06-06 15:19:54 +02:00
routing::{BoxResponseBody, EmptyRouter, MethodFilter},
service::HandleError,
2021-06-06 11:37:08 +02:00
};
2021-05-30 13:24:03 +02:00
use async_trait::async_trait;
2021-06-01 17:17:10 +02:00
use bytes::Bytes;
2021-05-30 13:24:03 +02:00
use futures_util::future;
use http::{Request, Response};
use std::{
2021-05-31 22:54:21 +02:00
convert::Infallible,
2021-05-30 13:24:03 +02:00
future::Future,
marker::PhantomData,
task::{Context, Poll},
};
2021-06-06 15:19:54 +02:00
use tower::{util::Oneshot, BoxError, Layer, Service, ServiceExt};
2021-06-06 11:37:08 +02:00
2021-06-06 22:41:52 +02:00
pub fn get<H, T>(handler: H) -> OnMethod<IntoService<H, T>, EmptyRouter>
2021-06-06 11:37:08 +02:00
where
2021-06-06 22:41:52 +02:00
H: Handler<T>,
2021-06-06 11:37:08 +02:00
{
on(MethodFilter::Get, handler)
}
2021-06-06 22:41:52 +02:00
pub fn post<H, T>(handler: H) -> OnMethod<IntoService<H, T>, EmptyRouter>
2021-06-06 11:37:08 +02:00
where
2021-06-06 22:41:52 +02:00
H: Handler<T>,
2021-06-06 11:37:08 +02:00
{
on(MethodFilter::Post, handler)
}
2021-06-06 22:41:52 +02:00
pub fn on<H, T>(method: MethodFilter, handler: H) -> OnMethod<IntoService<H, T>, EmptyRouter>
2021-06-06 11:37:08 +02:00
where
2021-06-06 22:41:52 +02:00
H: Handler<T>,
2021-06-06 11:37:08 +02:00
{
2021-06-06 15:19:54 +02:00
OnMethod {
method,
svc: handler.into_service(),
fallback: EmptyRouter,
}
2021-06-06 11:37:08 +02:00
}
2021-05-30 13:24:03 +02:00
mod sealed {
2021-06-06 11:37:08 +02:00
#![allow(unreachable_pub)]
2021-05-30 13:24:03 +02:00
pub trait HiddentTrait {}
pub struct Hidden;
impl HiddentTrait for Hidden {}
}
#[async_trait]
2021-06-06 22:41:52 +02:00
pub trait Handler<In>: Sized {
2021-05-30 13:24:03 +02:00
// This seals the trait. We cannot use the regular "sealed super trait" approach
// due to coherence.
#[doc(hidden)]
type Sealed: sealed::HiddentTrait;
2021-06-06 22:41:52 +02:00
async fn call(self, req: Request<Body>) -> Response<BoxBody>;
2021-05-30 13:24:03 +02:00
fn layer<L>(self, layer: L) -> Layered<L::Service, In>
where
2021-06-06 22:41:52 +02:00
L: Layer<IntoService<Self, In>>,
2021-05-30 13:24:03 +02:00
{
2021-06-06 11:37:08 +02:00
Layered::new(layer.layer(IntoService::new(self)))
}
2021-06-06 22:41:52 +02:00
fn into_service(self) -> IntoService<Self, In> {
2021-06-06 11:37:08 +02:00
IntoService::new(self)
2021-05-30 13:24:03 +02:00
}
}
#[async_trait]
2021-06-06 22:41:52 +02:00
impl<F, Fut, Res> Handler<()> for F
2021-05-30 13:24:03 +02:00
where
2021-06-01 08:32:58 +02:00
F: FnOnce(Request<Body>) -> Fut + Send + Sync,
2021-05-31 20:42:57 +02:00
Fut: Future<Output = Res> + Send,
2021-06-06 22:41:52 +02:00
Res: IntoResponse,
2021-05-30 13:24:03 +02:00
{
type Sealed = sealed::Hidden;
2021-06-06 22:41:52 +02:00
async fn call(self, req: Request<Body>) -> Response<BoxBody> {
self(req).await.into_response().map(BoxBody::new)
2021-05-30 13:24:03 +02:00
}
}
macro_rules! impl_handler {
2021-05-30 16:53:27 +02:00
() => {};
2021-05-30 13:24:03 +02:00
( $head:ident, $($tail:ident),* $(,)? ) => {
#[async_trait]
#[allow(non_snake_case)]
2021-06-06 22:41:52 +02:00
impl<F, Fut, Res, $head, $($tail,)*> Handler<($head, $($tail,)*)> for F
2021-05-30 13:24:03 +02:00
where
2021-06-01 08:32:58 +02:00
F: FnOnce(Request<Body>, $head, $($tail,)*) -> Fut + Send + Sync,
2021-05-31 20:42:57 +02:00
Fut: Future<Output = Res> + Send,
2021-06-06 22:41:52 +02:00
Res: IntoResponse,
$head: FromRequest + Send,
$( $tail: FromRequest + Send, )*
2021-05-30 13:24:03 +02:00
{
type Sealed = sealed::Hidden;
2021-06-06 22:41:52 +02:00
async fn call(self, mut req: Request<Body>) -> Response<BoxBody> {
2021-05-31 22:54:21 +02:00
let $head = match $head::from_request(&mut req).await {
Ok(value) => value,
2021-06-06 22:41:52 +02:00
Err(rejection) => return rejection.into_response().map(BoxBody::new),
2021-05-31 22:54:21 +02:00
};
2021-05-30 13:24:03 +02:00
$(
2021-05-31 22:54:21 +02:00
let $tail = match $tail::from_request(&mut req).await {
Ok(value) => value,
2021-06-06 22:41:52 +02:00
Err(rejection) => return rejection.into_response().map(BoxBody::new),
2021-05-31 22:54:21 +02:00
};
2021-05-30 13:24:03 +02:00
)*
2021-05-31 22:54:21 +02:00
2021-05-31 20:42:57 +02:00
let res = self(req, $head, $($tail,)*).await;
2021-05-31 22:54:21 +02:00
2021-06-06 22:41:52 +02:00
res.into_response().map(BoxBody::new)
2021-05-30 13:24:03 +02:00
}
}
impl_handler!($($tail,)*);
};
}
impl_handler!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16);
pub struct Layered<S, T> {
svc: S,
_input: PhantomData<fn() -> T>,
}
impl<S, T> Clone for Layered<S, T>
where
S: Clone,
{
fn clone(&self) -> Self {
Self::new(self.svc.clone())
}
}
#[async_trait]
2021-06-06 22:41:52 +02:00
impl<S, T, B> Handler<T> for Layered<S, T>
2021-05-30 13:24:03 +02:00
where
S: Service<Request<Body>, Response = Response<B>> + Send,
2021-06-06 22:41:52 +02:00
// S::Response: IntoResponse,
S::Error: IntoResponse,
2021-05-30 13:24:03 +02:00
S::Future: Send,
2021-06-06 22:41:52 +02:00
B: http_body::Body<Data = Bytes> + Send + Sync + 'static,
B::Error: Into<BoxError> + Send + Sync + 'static,
2021-05-30 13:24:03 +02:00
{
type Sealed = sealed::Hidden;
2021-06-06 22:41:52 +02:00
async fn call(self, req: Request<Body>) -> Response<BoxBody> {
2021-05-31 22:54:21 +02:00
match self
.svc
2021-05-30 13:24:03 +02:00
.oneshot(req)
.await
2021-05-31 22:54:21 +02:00
.map_err(IntoResponse::into_response)
{
2021-06-06 22:41:52 +02:00
Ok(res) => res.map(BoxBody::new),
Err(res) => res.map(BoxBody::new),
2021-05-31 22:54:21 +02:00
}
2021-05-30 13:24:03 +02:00
}
}
impl<S, T> Layered<S, T> {
pub(crate) fn new(svc: S) -> Self {
Self {
svc,
_input: PhantomData,
}
}
2021-06-01 17:17:10 +02:00
2021-06-04 01:00:48 +02:00
pub fn handle_error<F, B, Res>(self, f: F) -> Layered<HandleError<S, F>, T>
2021-06-01 17:17:10 +02:00
where
S: Service<Request<Body>, Response = Response<B>>,
F: FnOnce(S::Error) -> Res,
2021-06-06 22:41:52 +02:00
Res: IntoResponse,
2021-06-01 17:17:10 +02:00
{
let svc = HandleError::new(self.svc, f);
Layered::new(svc)
}
2021-05-30 13:24:03 +02:00
}
2021-06-06 22:41:52 +02:00
pub struct IntoService<H, T> {
2021-05-30 13:24:03 +02:00
handler: H,
2021-06-06 22:41:52 +02:00
_marker: PhantomData<fn() -> T>,
2021-05-30 13:24:03 +02:00
}
2021-06-06 22:41:52 +02:00
impl<H, T> IntoService<H, T> {
2021-06-06 11:37:08 +02:00
fn new(handler: H) -> Self {
2021-05-30 13:24:03 +02:00
Self {
handler,
2021-06-06 11:37:08 +02:00
_marker: PhantomData,
2021-05-30 13:24:03 +02:00
}
}
}
2021-06-06 22:41:52 +02:00
impl<H, T> Clone for IntoService<H, T>
2021-05-30 13:24:03 +02:00
where
H: Clone,
{
fn clone(&self) -> Self {
Self {
handler: self.handler.clone(),
2021-06-06 11:37:08 +02:00
_marker: PhantomData,
2021-05-30 13:24:03 +02:00
}
}
}
2021-06-06 22:41:52 +02:00
impl<H, T> Service<Request<Body>> for IntoService<H, T>
2021-05-30 13:24:03 +02:00
where
2021-06-06 22:41:52 +02:00
H: Handler<T> + Clone + Send + 'static,
2021-05-30 13:24:03 +02:00
{
2021-06-06 22:41:52 +02:00
type Response = Response<BoxBody>;
2021-05-31 22:54:21 +02:00
type Error = Infallible;
2021-05-30 13:24:03 +02:00
type Future = future::BoxFuture<'static, Result<Self::Response, Self::Error>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
2021-06-06 11:37:08 +02:00
// `IntoService` can only be constructed from async functions which are always ready, or from
2021-05-30 13:24:03 +02:00
// `Layered` which bufferes in `<Layered as Handler>::call` and is therefore also always
// ready.
Poll::Ready(Ok(()))
}
fn call(&mut self, req: Request<Body>) -> Self::Future {
let handler = self.handler.clone();
2021-06-06 22:41:52 +02:00
Box::pin(async move {
let res = Handler::call(handler, req).await;
Ok(res)
})
2021-05-30 13:24:03 +02:00
}
}
2021-06-06 15:19:54 +02:00
#[derive(Clone)]
pub struct OnMethod<S, F> {
pub(crate) method: MethodFilter,
pub(crate) svc: S,
pub(crate) fallback: F,
}
impl<S, F> OnMethod<S, F> {
2021-06-06 22:41:52 +02:00
pub fn get<H, T>(self, handler: H) -> OnMethod<IntoService<H, T>, Self>
2021-06-06 15:19:54 +02:00
where
2021-06-06 22:41:52 +02:00
H: Handler<T>,
2021-06-06 15:19:54 +02:00
{
self.on(MethodFilter::Get, handler)
}
2021-06-06 22:41:52 +02:00
pub fn post<H, T>(self, handler: H) -> OnMethod<IntoService<H, T>, Self>
2021-06-06 15:19:54 +02:00
where
2021-06-06 22:41:52 +02:00
H: Handler<T>,
2021-06-06 15:19:54 +02:00
{
self.on(MethodFilter::Post, handler)
}
2021-06-06 22:41:52 +02:00
pub fn on<H, T>(self, method: MethodFilter, handler: H) -> OnMethod<IntoService<H, T>, Self>
2021-06-06 15:19:54 +02:00
where
2021-06-06 22:41:52 +02:00
H: Handler<T>,
2021-06-06 15:19:54 +02:00
{
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<S, F, SB, FB> Service<Request<Body>> for OnMethod<S, F>
where
S: Service<Request<Body>, Response = Response<SB>, Error = Infallible> + Clone,
SB: http_body::Body<Data = Bytes> + Send + Sync + 'static,
SB::Error: Into<BoxError>,
F: Service<Request<Body>, Response = Response<FB>, Error = Infallible> + Clone,
FB: http_body::Body<Data = Bytes> + Send + Sync + 'static,
FB::Error: Into<BoxError>,
{
type Response = Response<BoxBody>;
type Error = Infallible;
#[allow(clippy::type_complexity)]
type Future = future::Either<
BoxResponseBody<Oneshot<S, Request<Body>>>,
BoxResponseBody<Oneshot<F, Request<Body>>>,
>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, req: Request<Body>) -> 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))
}
}
}