tokio-rs/axum
1
0
Fork 0
mirror of https://github.com/tokio-rs/axum.git synced 2026-08-25 00:00:23 +02:00
Code Issues Packages Projects Releases Wiki Activity

This changes everything

Browse Source
This commit is contained in:
David Pedersen
2021-05-31 22:54:21 +02:00
parent 19cbece1dc
commit 18f613ff98
5 changed files with 424 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/extract.rs
+265 -70
View File
@@ -1,36 +1,85 @@
use crate::{body::Body, Error};
use crate::{
body::Body,
response::{BoxIntoResponse, IntoResponse},
Error,
};
use async_trait::async_trait;
use bytes::Bytes;
use http::{header, Request, StatusCode};
use http::{header, Request};
use serde::de::DeserializeOwned;
use std::{collections::HashMap, str::FromStr};
use std::{collections::HashMap, convert::Infallible, str::FromStr};
#[async_trait]
pub trait FromRequest: Sized {
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error>;
}
pub trait FromRequest<B>: Sized {
type Rejection: IntoResponse<B>;
fn take_body(req: &mut Request<Body>) -> Body {
struct BodyAlreadyTaken;
if req.extensions_mut().insert(BodyAlreadyTaken).is_some() {
panic!("Cannot have two request body on extractors")
} else {
let body = std::mem::take(req.body_mut());
body
}
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection>;
}
#[async_trait]
impl<T> FromRequest for Option<T>
impl<T, B> FromRequest<B> for Option<T>
where
T: FromRequest,
T: FromRequest<B>,
{
async fn from_request(req: &mut Request<Body>) -> Result<Option<T>, Error> {
type Rejection = Infallible;
async fn from_request(req: &mut Request<Body>) -> Result<Option<T>, Self::Rejection> {
Ok(T::from_request(req).await.ok())
}
}
macro_rules! define_rejection {
(
#[status = $status:ident]
#[body = $body:expr]
pub struct $name:ident (());
) => {
#[derive(Debug)]
pub struct $name(());
impl IntoResponse<Body> for $name {
fn into_response(self) -> http::Response<Body> {
let mut res = http::Response::new(Body::from($body));
*res.status_mut() = http::StatusCode::$status;
res
}
}
};
(
#[status = $status:ident]
#[body = $body:expr]
pub struct $name:ident (BoxError);
) => {
#[derive(Debug)]
pub struct $name(tower::BoxError);
impl $name {
fn from_err<E>(err: E) -> Self
where
E: Into<tower::BoxError>,
{
Self(err.into())
}
}
impl IntoResponse<Body> for $name {
fn into_response(self) -> http::Response<Body> {
let mut res =
http::Response::new(Body::from(format!(concat!($body, ": {}"), self.0)));
*res.status_mut() = http::StatusCode::$status;
res
}
}
};
}
define_rejection! {
#[status = BAD_REQUEST]
#[body = "Query string was invalid or missing"]
pub struct QueryStringMissing(());
}
#[derive(Debug, Clone, Copy)]
pub struct Query<T>(T);
@@ -41,13 +90,15 @@ impl<T> Query<T> {
}
#[async_trait]
impl<T> FromRequest for Query<T>
impl<T> FromRequest<Body> for Query<T>
where
T: DeserializeOwned,
{
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> {
let query = req.uri().query().ok_or(Error::QueryStringMissing)?;
let value = serde_urlencoded::from_str(query).map_err(Error::DeserializeQueryString)?;
type Rejection = QueryStringMissing;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
let query = req.uri().query().ok_or(QueryStringMissing(()))?;
let value = serde_urlencoded::from_str(query).map_err(|_| QueryStringMissing(()))?;
Ok(Query(value))
}
}
@@ -61,22 +112,41 @@ impl<T> Json<T> {
}
}
define_rejection! {
#[status = BAD_REQUEST]
#[body = "Failed to parse the response body as JSON"]
pub struct InvalidJsonBody(BoxError);
}
define_rejection! {
#[status = BAD_REQUEST]
#[body = "Expected request with `Content-Type: application/json`"]
pub struct MissingJsonContentType(());
}
#[async_trait]
impl<T> FromRequest for Json<T>
impl<T> FromRequest<Body> for Json<T>
where
T: DeserializeOwned,
{
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> {
type Rejection = BoxIntoResponse<Body>;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
if has_content_type(&req, "application/json") {
let body = take_body(req);
let body = take_body(req).map_err(IntoResponse::boxed)?;
let bytes = hyper::body::to_bytes(body)
.await
.map_err(Error::ConsumeRequestBody)?;
let value = serde_json::from_slice(&bytes).map_err(Error::DeserializeRequestBody)?;
.map_err(InvalidJsonBody::from_err)
.map_err(IntoResponse::boxed)?;
let value = serde_json::from_slice(&bytes)
.map_err(InvalidJsonBody::from_err)
.map_err(IntoResponse::boxed)?;
Ok(Json(value))
} else {
Err(Error::Status(StatusCode::BAD_REQUEST))
Err(MissingJsonContentType(()).boxed())
}
}
}
@@ -97,6 +167,12 @@ fn has_content_type<B>(req: &Request<B>, expected_content_type: &str) -> bool {
content_type.starts_with(expected_content_type)
}
define_rejection! {
#[status = INTERNAL_SERVER_ERROR]
#[body = "Missing request extension"]
pub struct MissingExtension(());
}
#[derive(Debug, Clone, Copy)]
pub struct Extension<T>(T);
@@ -107,60 +183,93 @@ impl<T> Extension<T> {
}
#[async_trait]
impl<T> FromRequest for Extension<T>
impl<T> FromRequest<Body> for Extension<T>
where
T: Clone + Send + Sync + 'static,
{
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> {
type Rejection = MissingExtension;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
let value = req
.extensions()
.get::<T>()
.ok_or_else(|| Error::MissingExtension {
type_name: std::any::type_name::<T>(),
})
.ok_or(MissingExtension(()))
.map(|x| x.clone())?;
Ok(Extension(value))
}
}
define_rejection! {
#[status = BAD_REQUEST]
#[body = "Failed to buffer the request body"]
pub struct FailedToBufferBody(BoxError);
}
#[async_trait]
impl FromRequest for Bytes {
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> {
let body = take_body(req);
impl FromRequest<Body> for Bytes {
type Rejection = BoxIntoResponse<Body>;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
let body = take_body(req).map_err(IntoResponse::boxed)?;
let bytes = hyper::body::to_bytes(body)
.await
.map_err(Error::ConsumeRequestBody)?;
.map_err(FailedToBufferBody::from_err)
.map_err(IntoResponse::boxed)?;
Ok(bytes)
}
}
define_rejection! {
#[status = BAD_REQUEST]
#[body = "Response body didn't contain valid UTF-8"]
pub struct InvalidUtf8(BoxError);
}
#[async_trait]
impl FromRequest for String {
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> {
let body = take_body(req);
impl FromRequest<Body> for String {
type Rejection = BoxIntoResponse<Body>;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
let body = take_body(req).map_err(IntoResponse::boxed)?;
let bytes = hyper::body::to_bytes(body)
.await
.map_err(Error::ConsumeRequestBody)?
.map_err(FailedToBufferBody::from_err)
.map_err(IntoResponse::boxed)?
.to_vec();
let string = String::from_utf8(bytes).map_err(|_| Error::InvalidUtf8)?;
let string = String::from_utf8(bytes)
.map_err(InvalidUtf8::from_err)
.map_err(IntoResponse::boxed)?;
Ok(string)
}
}
#[async_trait]
impl FromRequest for Body {
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> {
let body = take_body(req);
Ok(body)
impl FromRequest<Body> for Body {
type Rejection = BodyAlreadyTaken;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
take_body(req)
}
}
define_rejection! {
#[status = PAYLOAD_TOO_LARGE]
#[body = "Request payload is too large"]
pub struct PayloadTooLarge(());
}
define_rejection! {
#[status = LENGTH_REQUIRED]
#[body = "Content length header is required"]
pub struct LengthRequired(());
}
#[derive(Debug, Clone)]
pub struct BytesMaxLength<const N: u64>(Bytes);
@@ -171,30 +280,38 @@ impl<const N: u64> BytesMaxLength<N> {
}
#[async_trait]
impl<const N: u64> FromRequest for BytesMaxLength<N> {
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> {
impl<const N: u64> FromRequest<Body> for BytesMaxLength<N> {
type Rejection = BoxIntoResponse<Body>;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
let content_length = req.headers().get(http::header::CONTENT_LENGTH).cloned();
let body = take_body(req);
let body = take_body(req).map_err(|reject| reject.boxed())?;
let content_length =
content_length.and_then(|value| value.to_str().ok()?.parse::<u64>().ok());
if let Some(length) = content_length {
if length > N {
return Err(Error::PayloadTooLarge);
return Err(PayloadTooLarge(()).boxed());
}
} else {
return Err(Error::LengthRequired);
return Err(LengthRequired(()).boxed());
};
let bytes = hyper::body::to_bytes(body)
.await
.map_err(Error::ConsumeRequestBody)?;
.map_err(|e| FailedToBufferBody::from_err(e).boxed())?;
Ok(BytesMaxLength(bytes))
}
}
define_rejection! {
#[status = INTERNAL_SERVER_ERROR]
#[body = "No url params found for matched route. This is a bug in tower-web. Please open an issue"]
pub struct MissingRouteParams(());
}
pub struct UrlParamsMap(HashMap<String, String>);
impl UrlParamsMap {
@@ -217,8 +334,10 @@ impl UrlParamsMap {
}
#[async_trait]
impl FromRequest for UrlParamsMap {
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> {
impl FromRequest<Body> for UrlParamsMap {
type Rejection = MissingRouteParams;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
if let Some(params) = req
.extensions_mut()
.get_mut::<Option<crate::routing::UrlParams>>()
@@ -226,62 +345,78 @@ impl FromRequest for UrlParamsMap {
let params = params.take().expect("params already taken").0;
Ok(Self(params.into_iter().collect()))
} else {
panic!("no url params found for matched route. This is a bug in tower-web")
Err(MissingRouteParams(()))
}
}
}
pub struct UrlParams<T>(T);
#[derive(Debug)]
pub struct InvalidUrlParam {
type_name: &'static str,
}
impl<T> UrlParams<T> {
pub fn into_inner(self) -> T {
self.0
impl InvalidUrlParam {
fn new<T>() -> Self {
InvalidUrlParam {
type_name: std::any::type_name::<T>(),
}
}
}
impl IntoResponse<Body> for InvalidUrlParam {
fn into_response(self) -> http::Response<Body> {
let mut res = http::Response::new(Body::from(format!(
"Invalid URL param. Expected something of type `{}`",
self.type_name
)));
*res.status_mut() = http::StatusCode::BAD_REQUEST;
res
}
}
pub struct UrlParams<T>(T);
macro_rules! impl_parse_url {
() => {};
( $head:ident, $($tail:ident),* $(,)? ) => {
#[async_trait]
impl<$head, $($tail,)*> FromRequest for UrlParams<($head, $($tail,)*)>
impl<$head, $($tail,)*> FromRequest<Body> for UrlParams<($head, $($tail,)*)>
where
$head: FromStr + Send,
$( $tail: FromStr + Send, )*
{
type Rejection = BoxIntoResponse<Body>;
#[allow(non_snake_case)]
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> {
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
let params = if let Some(params) = req
.extensions_mut()
.get_mut::<Option<crate::routing::UrlParams>>()
{
params.take().expect("params already taken").0
} else {
panic!("no url params found for matched route. This is a bug in tower-web")
return Err(MissingRouteParams(()).boxed())
};
if let [(_, $head), $((_, $tail),)*] = &*params {
let $head = if let Ok(x) = $head.parse::<$head>() {
x
} else {
return Err(Error::InvalidUrlParam {
type_name: std::any::type_name::<$head>(),
});
return Err(InvalidUrlParam::new::<$head>().boxed());
};
$(
let $tail = if let Ok(x) = $tail.parse::<$tail>() {
x
} else {
return Err(Error::InvalidUrlParam {
type_name: std::any::type_name::<$tail>(),
});
return Err(InvalidUrlParam::new::<$tail>().boxed());
};
)*
Ok(UrlParams(($head, $($tail,)*)))
} else {
panic!("wrong number of url params found for matched route. This is a bug in tower-web")
return Err(MissingRouteParams(()).boxed())
}
}
}
@@ -290,4 +425,64 @@ macro_rules! impl_parse_url {
};
}
impl_parse_url!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16);
impl_parse_url!(T1, T2, T3, T4, T5, T6);
impl<T1> UrlParams<(T1,)> {
pub fn into_inner(self) -> T1 {
(self.0).0
}
}
impl<T1, T2> UrlParams<(T1, T2)> {
pub fn into_inner(self) -> (T1, T2) {
((self.0).0, (self.0).1)
}
}
impl<T1, T2, T3> UrlParams<(T1, T2, T3)> {
pub fn into_inner(self) -> (T1, T2, T3) {
((self.0).0, (self.0).1, (self.0).2)
}
}
impl<T1, T2, T3, T4> UrlParams<(T1, T2, T3, T4)> {
pub fn into_inner(self) -> (T1, T2, T3, T4) {
((self.0).0, (self.0).1, (self.0).2, (self.0).3)
}
}
impl<T1, T2, T3, T4, T5> UrlParams<(T1, T2, T3, T4, T5)> {
pub fn into_inner(self) -> (T1, T2, T3, T4, T5) {
((self.0).0, (self.0).1, (self.0).2, (self.0).3, (self.0).4)
}
}
impl<T1, T2, T3, T4, T5, T6> UrlParams<(T1, T2, T3, T4, T5, T6)> {
pub fn into_inner(self) -> (T1, T2, T3, T4, T5, T6) {
(
(self.0).0,
(self.0).1,
(self.0).2,
(self.0).3,
(self.0).4,
(self.0).5,
)
}
}
define_rejection! {
#[status = INTERNAL_SERVER_ERROR]
#[body = "Cannot have two request body extractors for a single handler"]
pub struct BodyAlreadyTaken(());
}
fn take_body(req: &mut Request<Body>) -> Result<Body, BodyAlreadyTaken> {
struct BodyAlreadyTakenExt;
if req.extensions_mut().insert(BodyAlreadyTakenExt).is_some() {
Err(BodyAlreadyTaken(()))
} else {
let body = std::mem::take(req.body_mut());
Ok(body)
}
}