Files
axum/src/extract/mod.rs
T

319 lines
9.0 KiB
Rust
Raw Normal View History

2021-06-07 16:28:40 +02:00
//! Types and traits for extracting data from requests.
2021-06-01 14:52:18 +02:00
use crate::{body::Body, response::IntoResponse};
2021-05-31 12:55:39 +02:00
use async_trait::async_trait;
2021-05-30 14:33:20 +02:00
use bytes::Bytes;
2021-06-06 11:37:08 +02:00
use http::{header, Request, Response};
use rejection::{
2021-06-06 22:41:52 +02:00
BodyAlreadyTaken, FailedToBufferBody, InvalidJsonBody, InvalidUrlParam, InvalidUtf8,
LengthRequired, MissingExtension, MissingJsonContentType, MissingRouteParams, PayloadTooLarge,
2021-06-06 11:37:08 +02:00
QueryStringMissing,
};
2021-05-30 13:24:03 +02:00
use serde::de::DeserializeOwned;
2021-05-31 22:54:21 +02:00
use std::{collections::HashMap, convert::Infallible, str::FromStr};
2021-05-30 13:24:03 +02:00
2021-06-06 11:37:08 +02:00
pub mod rejection;
2021-05-31 12:55:39 +02:00
#[async_trait]
2021-06-06 22:41:52 +02:00
pub trait FromRequest: Sized {
type Rejection: IntoResponse;
2021-05-31 14:04:05 +02:00
2021-05-31 22:54:21 +02:00
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection>;
2021-05-31 14:04:05 +02:00
}
2021-05-31 12:55:39 +02:00
#[async_trait]
2021-06-06 22:41:52 +02:00
impl<T> FromRequest for Option<T>
2021-05-30 13:24:03 +02:00
where
2021-06-06 22:41:52 +02:00
T: FromRequest,
2021-05-30 13:24:03 +02:00
{
2021-05-31 22:54:21 +02:00
type Rejection = Infallible;
async fn from_request(req: &mut Request<Body>) -> Result<Option<T>, Self::Rejection> {
2021-05-31 12:55:39 +02:00
Ok(T::from_request(req).await.ok())
2021-05-30 13:24:03 +02:00
}
}
2021-06-06 15:19:54 +02:00
#[derive(Debug, Clone, Copy, Default)]
2021-06-01 14:52:18 +02:00
pub struct Query<T>(pub T);
2021-05-30 13:24:03 +02:00
2021-05-31 12:55:39 +02:00
#[async_trait]
2021-06-06 22:41:52 +02:00
impl<T> FromRequest for Query<T>
2021-05-30 13:24:03 +02:00
where
2021-05-31 12:55:39 +02:00
T: DeserializeOwned,
2021-05-30 13:24:03 +02:00
{
2021-05-31 22:54:21 +02:00
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)?;
2021-05-31 12:55:39 +02:00
Ok(Query(value))
2021-05-30 13:24:03 +02:00
}
}
2021-06-06 15:19:54 +02:00
#[derive(Debug, Clone, Copy, Default)]
2021-06-01 14:52:18 +02:00
pub struct Json<T>(pub T);
2021-05-30 13:24:03 +02:00
2021-05-31 12:55:39 +02:00
#[async_trait]
2021-06-06 22:41:52 +02:00
impl<T> FromRequest for Json<T>
2021-05-30 13:24:03 +02:00
where
T: DeserializeOwned,
{
2021-06-01 14:52:18 +02:00
type Rejection = Response<Body>;
2021-05-31 22:54:21 +02:00
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
2021-06-06 11:37:08 +02:00
if has_content_type(req, "application/json") {
2021-06-01 14:52:18 +02:00
let body = take_body(req).map_err(IntoResponse::into_response)?;
2021-05-31 12:22:16 +02:00
2021-05-31 12:55:39 +02:00
let bytes = hyper::body::to_bytes(body)
.await
2021-05-31 22:54:21 +02:00
.map_err(InvalidJsonBody::from_err)
2021-06-01 14:52:18 +02:00
.map_err(IntoResponse::into_response)?;
2021-05-31 22:54:21 +02:00
let value = serde_json::from_slice(&bytes)
.map_err(InvalidJsonBody::from_err)
2021-06-01 14:52:18 +02:00
.map_err(IntoResponse::into_response)?;
2021-05-31 22:54:21 +02:00
2021-05-31 12:55:39 +02:00
Ok(Json(value))
2021-05-31 12:22:16 +02:00
} else {
Err(MissingJsonContentType.into_response())
2021-05-31 12:22:16 +02:00
}
}
}
2021-05-30 13:24:03 +02:00
2021-05-31 12:22:16 +02:00
fn has_content_type<B>(req: &Request<B>, expected_content_type: &str) -> bool {
let content_type = if let Some(content_type) = req.headers().get(header::CONTENT_TYPE) {
content_type
} else {
return false;
};
2021-05-30 13:24:03 +02:00
2021-05-31 12:22:16 +02:00
let content_type = if let Ok(content_type) = content_type.to_str() {
content_type
} else {
return false;
};
content_type.starts_with(expected_content_type)
2021-05-30 13:24:03 +02:00
}
#[derive(Debug, Clone, Copy)]
2021-06-01 14:52:18 +02:00
pub struct Extension<T>(pub T);
2021-05-30 13:24:03 +02:00
2021-05-31 12:55:39 +02:00
#[async_trait]
2021-06-06 22:41:52 +02:00
impl<T> FromRequest for Extension<T>
2021-05-30 13:24:03 +02:00
where
T: Clone + Send + Sync + 'static,
{
2021-05-31 22:54:21 +02:00
type Rejection = MissingExtension;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
2021-05-31 12:55:39 +02:00
let value = req
.extensions()
.get::<T>()
.ok_or(MissingExtension)
2021-05-31 12:55:39 +02:00
.map(|x| x.clone())?;
Ok(Extension(value))
2021-05-30 13:24:03 +02:00
}
}
2021-05-31 12:55:39 +02:00
#[async_trait]
2021-06-06 22:41:52 +02:00
impl FromRequest for Bytes {
2021-06-01 14:52:18 +02:00
type Rejection = Response<Body>;
2021-05-31 22:54:21 +02:00
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
2021-06-01 14:52:18 +02:00
let body = take_body(req).map_err(IntoResponse::into_response)?;
2021-05-30 14:33:20 +02:00
2021-05-31 12:55:39 +02:00
let bytes = hyper::body::to_bytes(body)
.await
2021-05-31 22:54:21 +02:00
.map_err(FailedToBufferBody::from_err)
2021-06-01 14:52:18 +02:00
.map_err(IntoResponse::into_response)?;
2021-05-31 12:55:39 +02:00
Ok(bytes)
2021-05-30 14:33:20 +02:00
}
}
2021-05-31 12:55:39 +02:00
#[async_trait]
2021-06-06 22:41:52 +02:00
impl FromRequest for String {
2021-06-01 14:52:18 +02:00
type Rejection = Response<Body>;
2021-05-31 22:54:21 +02:00
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
2021-06-01 14:52:18 +02:00
let body = take_body(req).map_err(IntoResponse::into_response)?;
2021-05-31 12:22:16 +02:00
2021-05-31 12:55:39 +02:00
let bytes = hyper::body::to_bytes(body)
.await
2021-05-31 22:54:21 +02:00
.map_err(FailedToBufferBody::from_err)
2021-06-01 14:52:18 +02:00
.map_err(IntoResponse::into_response)?
2021-05-31 12:55:39 +02:00
.to_vec();
2021-05-31 22:54:21 +02:00
let string = String::from_utf8(bytes)
.map_err(InvalidUtf8::from_err)
2021-06-01 14:52:18 +02:00
.map_err(IntoResponse::into_response)?;
2021-05-31 12:55:39 +02:00
Ok(string)
2021-05-31 12:22:16 +02:00
}
}
2021-05-31 12:55:39 +02:00
#[async_trait]
2021-06-06 22:41:52 +02:00
impl FromRequest for Body {
2021-05-31 22:54:21 +02:00
type Rejection = BodyAlreadyTaken;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
take_body(req)
2021-05-31 12:22:16 +02:00
}
}
2021-05-30 13:24:03 +02:00
#[derive(Debug, Clone)]
2021-06-01 14:52:18 +02:00
pub struct BytesMaxLength<const N: u64>(pub Bytes);
2021-05-30 13:24:03 +02:00
2021-05-31 12:55:39 +02:00
#[async_trait]
2021-06-06 22:41:52 +02:00
impl<const N: u64> FromRequest for BytesMaxLength<N> {
2021-06-01 14:52:18 +02:00
type Rejection = Response<Body>;
2021-05-31 22:54:21 +02:00
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
2021-05-30 14:33:20 +02:00
let content_length = req.headers().get(http::header::CONTENT_LENGTH).cloned();
2021-06-01 14:52:18 +02:00
let body = take_body(req).map_err(|reject| reject.into_response())?;
2021-05-30 13:24:03 +02:00
2021-05-31 12:55:39 +02:00
let content_length =
content_length.and_then(|value| value.to_str().ok()?.parse::<u64>().ok());
2021-05-30 14:33:20 +02:00
2021-05-31 12:55:39 +02:00
if let Some(length) = content_length {
if length > N {
return Err(PayloadTooLarge.into_response());
2021-05-31 12:55:39 +02:00
}
} else {
return Err(LengthRequired.into_response());
2021-05-31 12:55:39 +02:00
};
2021-05-30 14:33:20 +02:00
2021-05-31 12:55:39 +02:00
let bytes = hyper::body::to_bytes(body)
.await
2021-06-01 14:52:18 +02:00
.map_err(|e| FailedToBufferBody::from_err(e).into_response())?;
2021-05-30 14:33:20 +02:00
2021-05-31 12:55:39 +02:00
Ok(BytesMaxLength(bytes))
2021-05-30 13:24:03 +02:00
}
}
2021-05-30 15:44:26 +02:00
2021-06-03 21:36:39 +02:00
#[derive(Debug)]
2021-05-30 16:37:27 +02:00
pub struct UrlParamsMap(HashMap<String, String>);
2021-05-30 15:44:26 +02:00
2021-05-30 16:37:27 +02:00
impl UrlParamsMap {
2021-06-01 00:34:09 +02:00
pub fn get(&self, key: &str) -> Option<&str> {
self.0.get(key).map(|s| &**s)
2021-05-30 15:44:26 +02:00
}
2021-05-30 16:37:27 +02:00
2021-06-01 00:34:09 +02:00
pub fn get_typed<T>(&self, key: &str) -> Option<T>
2021-05-30 16:37:27 +02:00
where
2021-05-30 16:53:27 +02:00
T: FromStr,
2021-05-30 16:37:27 +02:00
{
2021-06-01 00:34:09 +02:00
self.get(key)?.parse().ok()
2021-05-30 16:37:27 +02:00
}
2021-05-30 15:44:26 +02:00
}
2021-05-31 12:55:39 +02:00
#[async_trait]
2021-06-06 22:41:52 +02:00
impl FromRequest for UrlParamsMap {
2021-05-31 22:54:21 +02:00
type Rejection = MissingRouteParams;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
2021-05-30 15:44:26 +02:00
if let Some(params) = req
.extensions_mut()
.get_mut::<Option<crate::routing::UrlParams>>()
{
let params = params.take().expect("params already taken").0;
2021-05-31 12:55:39 +02:00
Ok(Self(params.into_iter().collect()))
2021-05-30 15:44:26 +02:00
} else {
Err(MissingRouteParams)
2021-05-30 15:44:26 +02:00
}
}
}
2021-05-30 16:53:27 +02:00
2021-06-01 14:52:18 +02:00
pub struct UrlParams<T>(pub T);
2021-05-31 22:54:21 +02:00
2021-05-30 16:53:27 +02:00
macro_rules! impl_parse_url {
() => {};
( $head:ident, $($tail:ident),* $(,)? ) => {
2021-05-31 12:55:39 +02:00
#[async_trait]
2021-06-06 22:41:52 +02:00
impl<$head, $($tail,)*> FromRequest for UrlParams<($head, $($tail,)*)>
2021-05-30 16:53:27 +02:00
where
$head: FromStr + Send,
$( $tail: FromStr + Send, )*
{
2021-06-01 14:52:18 +02:00
type Rejection = Response<Body>;
2021-05-31 22:54:21 +02:00
2021-05-30 16:53:27 +02:00
#[allow(non_snake_case)]
2021-05-31 22:54:21 +02:00
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
2021-05-30 16:53:27 +02:00
let params = if let Some(params) = req
.extensions_mut()
.get_mut::<Option<crate::routing::UrlParams>>()
{
params.take().expect("params already taken").0
} else {
return Err(MissingRouteParams.into_response())
2021-05-30 16:53:27 +02:00
};
if let [(_, $head), $((_, $tail),)*] = &*params {
let $head = if let Ok(x) = $head.parse::<$head>() {
x
} else {
2021-06-01 14:52:18 +02:00
return Err(InvalidUrlParam::new::<$head>().into_response());
2021-05-30 16:53:27 +02:00
};
$(
let $tail = if let Ok(x) = $tail.parse::<$tail>() {
x
} else {
2021-06-01 14:52:18 +02:00
return Err(InvalidUrlParam::new::<$tail>().into_response());
2021-05-30 16:53:27 +02:00
};
)*
2021-05-31 12:55:39 +02:00
Ok(UrlParams(($head, $($tail,)*)))
2021-05-30 16:53:27 +02:00
} else {
return Err(MissingRouteParams.into_response())
2021-05-30 16:53:27 +02:00
}
}
}
impl_parse_url!($($tail,)*);
};
}
2021-06-07 16:28:40 +02:00
impl_parse_url!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16);
2021-05-31 22:54:21 +02:00
fn take_body(req: &mut Request<Body>) -> Result<Body, BodyAlreadyTaken> {
struct BodyAlreadyTakenExt;
if req.extensions_mut().insert(BodyAlreadyTakenExt).is_some() {
Err(BodyAlreadyTaken)
2021-05-31 22:54:21 +02:00
} else {
let body = std::mem::take(req.body_mut());
Ok(body)
}
}
2021-06-07 16:28:40 +02:00
macro_rules! impl_from_request_tuple {
() => {};
( $head:ident, $($tail:ident),* $(,)? ) => {
#[allow(non_snake_case)]
#[async_trait]
impl<R, $head, $($tail,)*> FromRequest for ($head, $($tail,)*)
where
R: IntoResponse,
$head: FromRequest<Rejection = R> + Send,
$( $tail: FromRequest<Rejection = R> + Send, )*
{
type Rejection = R;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
let $head = FromRequest::from_request(req).await?;
$( let $tail = FromRequest::from_request(req).await?; )*
Ok(($head, $($tail,)*))
}
}
impl_from_request_tuple!($($tail,)*);
};
}
impl_from_request_tuple!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16);