Add Scheme extractor

axum/src/extract/host.rs

@@ -142,7 +142,7 @@ mod tests {
         assert_eq!(value, "192.0.2.60");
 
         // is case insensitive
-        let headers = header_map(&[(FORWARDED, "host=192.0.2.60;proto=http;by=203.0.113.43")]);
+        let headers = header_map(&[(FORWARDED, "HOST=192.0.2.60;proto=http;by=203.0.113.43")]);
         let value = parse_forwarded(&headers).unwrap();
         assert_eq!(value, "192.0.2.60");
 

axum/src/extract/mod.rs

@@ -15,6 +15,7 @@ pub(crate) mod nested_path;
 mod raw_form;
 mod raw_query;
 mod request_parts;
+mod scheme;
 mod state;
 
 #[doc(inline)]
@@ -31,6 +32,7 @@ pub use self::{
     path::{Path, RawPathParams},
     raw_form::RawForm,
     raw_query::RawQuery,
+    scheme::Scheme,
     state::State,
 };
 

axum/src/extract/rejection.rs

@@ -73,6 +73,14 @@ define_rejection! {
     pub struct FailedToResolveHost;
 }
 
+define_rejection! {
+    #[status = BAD_REQUEST]
+    #[body = "No scheme found in request"]
+    /// Rejection type used if the [`Scheme`](super::Scheme) extractor is unable to
+    /// resolve a host.
+    pub struct FailedToResolveScheme;
+}
+
 define_rejection! {
     #[status = BAD_REQUEST]
     #[body = "Failed to deserialize form"]
@@ -188,6 +196,16 @@ composite_rejection! {
     }
 }
 
+composite_rejection! {
+    /// Rejection used for [`Scheme`](super::Scheme).
+    ///
+    /// Contains one variant for each way the [`Scheme`](super::Scheme) extractor
+    /// can fail.
+    pub enum SchemeRejection {
+        FailedToResolveScheme,
+    }
+}
+
 #[cfg(feature = "matched-path")]
 define_rejection! {
     #[status = INTERNAL_SERVER_ERROR]

axum/src/extract/scheme.rs

@@ -0,0 +1,156 @@
+use super::{
+    rejection::{FailedToResolveScheme, SchemeRejection},
+    FromRequestParts,
+};
+use async_trait::async_trait;
+use http::{
+    header::{HeaderMap, FORWARDED},
+    request::Parts,
+};
+const X_FORWARDED_PROTO_HEADER_KEY: &str = "X-Forwarded-Proto";
+
+/// Extractor that resolves the scheme / protocol of the request.
+///
+/// The scheme is resolved through the following, in order:
+/// - `Forwarded` header
+/// - `X-Forwarded-Proto` header
+/// - request target / URI
+///
+/// Note that user agents can set the `X-Forwarded-Proto` header to arbitrary values so make
+/// sure to validate them to avoid security issues.
+#[derive(Debug, Clone)]
+pub struct Scheme(pub String);
+
+#[async_trait]
+impl<S> FromRequestParts<S> for Scheme
+where
+    S: Send + Sync,
+{
+    type Rejection = SchemeRejection;
+
+    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
+        // Within Forwarded header
+        if let Some(scheme) = parse_forwarded(&parts.headers) {
+            return Ok(Scheme(scheme.to_owned()));
+        }
+
+        // X-Forwarded-Proto
+        if let Some(scheme) = parts
+            .headers
+            .get(X_FORWARDED_PROTO_HEADER_KEY)
+            .and_then(|scheme| scheme.to_str().ok())
+        {
+            return Ok(Scheme(scheme.to_owned()));
+        }
+
+        // From parts
+        if let Some(scheme) = parts.uri.scheme_str() {
+            return Ok(Scheme(scheme.to_owned()));
+        }
+
+        Err(SchemeRejection::FailedToResolveScheme(
+            FailedToResolveScheme,
+        ))
+    }
+}
+
+fn parse_forwarded(headers: &HeaderMap) -> Option<&str> {
+    // if there are multiple `Forwarded` `HeaderMap::get` will return the first one
+    let forwarded_values = headers.get(FORWARDED)?.to_str().ok()?;
+
+    // get the first set of values
+    let first_value = forwarded_values.split(',').next()?;
+
+    // find the value of the `proto` field
+    first_value.split(';').find_map(|pair| {
+        let (key, value) = pair.split_once('=')?;
+        key.trim()
+            .eq_ignore_ascii_case("proto")
+            .then(|| value.trim().trim_matches('"'))
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{routing::get, test_helpers::TestClient, Router};
+    use http::{header::HeaderName, Request};
+
+    use axum_core::{body::Body, extract::FromRequest};
+
+    fn test_client() -> TestClient {
+        async fn scheme_as_body(Scheme(scheme): Scheme) -> String {
+            scheme
+        }
+
+        TestClient::new(Router::new().route("/", get(scheme_as_body)))
+    }
+
+    #[crate::test]
+    async fn forwarded_scheme_parsing() {
+        // the basic case
+        let headers = header_map(&[(FORWARDED, "host=192.0.2.60;proto=http;by=203.0.113.43")]);
+        let value = parse_forwarded(&headers).unwrap();
+        assert_eq!(value, "http");
+
+        // is case insensitive
+        let headers = header_map(&[(FORWARDED, "host=192.0.2.60;PROTO=https;by=203.0.113.43")]);
+        let value = parse_forwarded(&headers).unwrap();
+        assert_eq!(value, "https");
+
+        // multiple values in one header
+        let headers = header_map(&[(FORWARDED, "proto=ftp, proto=https")]);
+        let value = parse_forwarded(&headers).unwrap();
+        assert_eq!(value, "ftp");
+
+        // multiple header values
+        let headers = header_map(&[(FORWARDED, "proto=ftp"), (FORWARDED, "proto=https")]);
+        let value = parse_forwarded(&headers).unwrap();
+        assert_eq!(value, "ftp");
+    }
+
+    #[crate::test]
+    async fn x_forwarded_scheme_header() {
+        let original_scheme = "https";
+        let scheme = test_client()
+            .get("/")
+            .header(X_FORWARDED_PROTO_HEADER_KEY, original_scheme)
+            .await
+            .text()
+            .await;
+        assert_eq!(scheme, original_scheme);
+    }
+
+    #[crate::test]
+    async fn from_parts() {
+        let original_scheme = "http";
+        let req = Request::builder()
+            .uri(format!("{original_scheme}://localhost/"))
+            .body(Body::empty())
+            .unwrap();
+        assert_eq!(
+            Scheme::from_request(req, &()).await.unwrap().0,
+            original_scheme
+        );
+    }
+
+    #[crate::test]
+    async fn precedence_forwarded_over_x_forwarded() {
+        let scheme = test_client()
+            .get("/")
+            .header(X_FORWARDED_PROTO_HEADER_KEY, "https")
+            .header(FORWARDED, "proto=ftp")
+            .await
+            .text()
+            .await;
+        assert_eq!(scheme, "ftp");
+    }
+
+    fn header_map(values: &[(HeaderName, &str)]) -> HeaderMap {
+        let mut headers = HeaderMap::new();
+        for (key, value) in values {
+            headers.append(key, value.parse().unwrap());
+        }
+        headers
+    }
+}