From b6e67eefd725fedb9ce996f625390c66678eac4c Mon Sep 17 00:00:00 2001 From: David Pedersen Date: Tue, 15 Jun 2021 21:27:21 +0200 Subject: [PATCH] Add support for extracting typed headers (#18) Uses the `headers` crate. --- Cargo.toml | 1 + src/extract/mod.rs | 41 ++++++++++++++++++++++++++++++++++++++++ src/extract/rejection.rs | 20 ++++++++++++++++++++ src/lib.rs | 1 + src/tests.rs | 29 +++++++++++++++++++++++++++- 5 files changed, 91 insertions(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index cce7af88..ced7cc15 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,6 +38,7 @@ tower-http = { version = "0.1", features = ["add-extension", "map-response-body" tokio-tungstenite = { optional = true, version = "0.14" } sha-1 = { optional = true, version = "0.9.6" } base64 = { optional = true, version = "0.13" } +headers = { optional = true, version = "0.3" } [dev-dependencies] askama = "0.10.5" diff --git a/src/extract/mod.rs b/src/extract/mod.rs index 1245f409..3b13392a 100644 --- a/src/extract/mod.rs +++ b/src/extract/mod.rs @@ -766,3 +766,44 @@ macro_rules! impl_from_request_tuple { } impl_from_request_tuple!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16); + +/// Extractor that extracts a typed header value from [`headers`]. +/// +/// # Example +/// +/// ```rust,no_run +/// use awebframework::{extract::TypedHeader, prelude::*}; +/// use headers::UserAgent; +/// +/// async fn users_teams_show( +/// TypedHeader(user_agent): TypedHeader, +/// ) { +/// // ... +/// } +/// +/// let app = route("/users/:user_id/team/:team_id", get(users_teams_show)); +/// ``` +#[cfg(feature = "headers")] +#[cfg_attr(docsrs, doc(cfg(feature = "headers")))] +#[derive(Debug, Clone, Copy)] +pub struct TypedHeader(pub T); + +#[cfg(feature = "headers")] +#[cfg_attr(docsrs, doc(cfg(feature = "headers")))] +#[async_trait] +impl FromRequest for TypedHeader +where + T: headers::Header, +{ + type Rejection = rejection::TypedHeaderRejection; + + async fn from_request(req: &mut Request) -> Result { + let header_values = req.headers().get_all(T::name()); + T::decode(&mut header_values.iter()) + .map(Self) + .map_err(|err| rejection::TypedHeaderRejection { + err, + name: T::name(), + }) + } +} diff --git a/src/extract/rejection.rs b/src/extract/rejection.rs index df4e6b34..edd42913 100644 --- a/src/extract/rejection.rs +++ b/src/extract/rejection.rs @@ -1,5 +1,6 @@ //! Rejection response types. +use http::StatusCode; use tower::BoxError; use super::IntoResponse; @@ -341,3 +342,22 @@ where } } } + +/// Rejection used for [`TypedHeader`](super::TypedHeader). +#[cfg(feature = "headers")] +#[cfg_attr(docsrs, doc(cfg(feature = "headers")))] +#[derive(Debug)] +pub struct TypedHeaderRejection { + pub(super) name: &'static http::header::HeaderName, + pub(super) err: headers::Error, +} + +#[cfg(feature = "headers")] +#[cfg_attr(docsrs, doc(cfg(feature = "headers")))] +impl IntoResponse for TypedHeaderRejection { + fn into_response(self) -> http::Response { + let mut res = format!("{} ({})", self.err, self.name).into_response(); + *res.status_mut() = StatusCode::BAD_REQUEST; + res + } +} diff --git a/src/lib.rs b/src/lib.rs index 52861a15..93c98699 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -551,6 +551,7 @@ //! - `ws`: Enables WebSockets support. //! - `hyper-h1`: Enables hyper's `http1` feature. On by default. //! - `hyper-h2`: Enables hyper's `http2` feature. +//! - `headers`: Enables extracing typed headers via [`extract::TypedHeader`]. //! //! [tower]: https://crates.io/crates/tower //! [tower-http]: https://crates.io/crates/tower-http diff --git a/src/tests.rs b/src/tests.rs index c0877a24..50a48ffd 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -1,4 +1,4 @@ -use crate::{handler::on, prelude::*, routing::MethodFilter, service}; +use crate::{handler::on, prelude::*, response::IntoResponse, routing::MethodFilter, service}; use bytes::Bytes; use http::{Request, Response, StatusCode}; use hyper::{Body, Server}; @@ -580,6 +580,33 @@ async fn disjunction() { assert_eq!(res.text().await.unwrap(), "v0: games#show (123)"); } +#[tokio::test] +async fn typed_header() { + use extract::TypedHeader; + async fn handle(TypedHeader(user_agent): TypedHeader) -> impl IntoResponse { + user_agent.to_string() + } + + let app = route("/", get(handle)); + + let addr = run_in_background(app).await; + + let client = reqwest::Client::new(); + + let res = client + .get(format!("http://{}", addr)) + .header("user-agent", "foobar") + .send() + .await + .unwrap(); + let body = res.text().await.unwrap(); + assert_eq!(body, "foobar"); + + let res = client.get(format!("http://{}", addr)).send().await.unwrap(); + let body = res.text().await.unwrap(); + assert_eq!(body, "invalid HTTP header (user-agent)"); +} + /// Run a `tower::Service` in the background and get a URI for it. async fn run_in_background(svc: S) -> SocketAddr where