feat(redirect): allow custom redirects to access response headers (#916)

Author: 0x676e67

examples/request_with_redirect.rs

@@ -2,13 +2,20 @@ use wreq::redirect::Policy;
 
 #[tokio::main]
 async fn main() -> wreq::Result<()> {
-    tracing_subscriber::fmt()
-        .with_max_level(tracing::Level::TRACE)
-        .init();
-
     // Use the API you're already familiar with
-    let resp = wreq::get("http://google.com/")
-        .redirect(Policy::default())
+    let resp = wreq::get("https://google.com/")
+        .redirect(Policy::custom(|attempt| {
+            // we can inspect the redirect attempt
+            println!(
+                "Redirecting (status: {}) to {:?} and headers: {:#?}",
+                attempt.status(),
+                attempt.uri(),
+                attempt.headers()
+            );
+
+            // we can follow redirects as normal
+            attempt.follow()
+        }))
         .send()
         .await?;
     println!("{}", resp.text().await?);

src/client/layer/redirect/future.rs

@@ -121,9 +121,11 @@ where
 
                 let attempt = Attempt {
                     status: res.status(),
+                    headers: res.headers(),
                     location: &location,
                     previous: uri,
                 };
+
                 match policy.redirect(&attempt)? {
                     Action::Follow => {
                         *uri = location;

src/client/layer/redirect/policy.rs

@@ -1,6 +1,6 @@
 //! Tools for customizing the behavior of a [`FollowRedirect`][super::FollowRedirect] middleware.
 
-use http::{Extensions, Request, StatusCode, Uri};
+use http::{Extensions, HeaderMap, Request, StatusCode, Uri};
 
 /// Trait for the policy on handling redirection responses.
 pub trait Policy<B, E> {
@@ -83,9 +83,10 @@ where
 
 /// A type that holds information on a redirection attempt.
 pub struct Attempt<'a> {
-    pub(crate) status: StatusCode,
-    pub(crate) location: &'a Uri,
-    pub(crate) previous: &'a Uri,
+    pub(super) status: StatusCode,
+    pub(super) headers: &'a HeaderMap,
+    pub(super) location: &'a Uri,
+    pub(super) previous: &'a Uri,
 }
 
 impl<'a> Attempt<'a> {
@@ -95,6 +96,12 @@ impl<'a> Attempt<'a> {
         self.status
     }
 
+    /// Returns the headers of the redirection response.
+    #[inline(always)]
+    pub fn headers(&self) -> &'a HeaderMap {
+        self.headers
+    }
+
     /// Returns the destination URI of the redirection.
     #[inline(always)]
     pub fn location(&self) -> &'a Uri {

src/redirect.rs

@@ -6,7 +6,8 @@
 
 use std::{error::Error as StdError, fmt, sync::Arc};
 
-use http::{Extensions, HeaderMap, HeaderValue, StatusCode, Uri};
+use bytes::Bytes;
+use http::{Extensions, HeaderMap, HeaderValue, StatusCode, Uri, uri::Authority};
 
 use crate::{
     client::{
@@ -28,7 +29,7 @@ use crate::{
 ///   redirect hops in a chain.
 /// - `none` can be used to disable all redirect behavior.
 /// - `custom` can be used to create a customized policy.
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 pub struct Policy {
     inner: PolicyKind,
 }
@@ -38,6 +39,7 @@ pub struct Policy {
 #[derive(Debug)]
 pub struct Attempt<'a> {
     status: StatusCode,
+    headers: &'a HeaderMap,
     next: &'a Uri,
     previous: &'a [Uri],
 }
@@ -146,9 +148,16 @@ impl Policy {
         }
     }
 
-    pub(crate) fn check(&self, status: StatusCode, next: &Uri, previous: &[Uri]) -> ActionKind {
+    fn check(
+        &self,
+        status: StatusCode,
+        headers: &HeaderMap,
+        next: &Uri,
+        previous: &[Uri],
+    ) -> ActionKind {
         self.redirect(Attempt {
             status,
+            headers,
             next,
             previous,
         })
@@ -169,6 +178,11 @@ impl<'a> Attempt<'a> {
         self.status
     }
 
+    /// Get the headers of redirect.
+    pub fn headers(&self) -> &HeaderMap {
+        self.headers
+    }
+
     /// Get the next URI to redirect to.
     pub fn uri(&self) -> &Uri {
         self.next
@@ -212,12 +226,6 @@ enum PolicyKind {
     None,
 }
 
-impl fmt::Debug for Policy {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        f.debug_tuple("Policy").field(&self.inner).finish()
-    }
-}
-
 impl fmt::Debug for PolicyKind {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match *self {
@@ -235,21 +243,6 @@ pub(crate) enum ActionKind {
     Error(BoxError),
 }
 
-fn remove_sensitive_headers(headers: &mut HeaderMap, next: &Uri, previous: &[Uri]) {
-    if let Some(previous) = previous.last() {
-        let cross_host = next.host() != previous.host()
-            || next.port() != previous.port()
-            || next.scheme() != previous.scheme();
-        if cross_host {
-            headers.remove(AUTHORIZATION);
-            headers.remove(COOKIE);
-            headers.remove("cookie2");
-            headers.remove(PROXY_AUTHORIZATION);
-            headers.remove(WWW_AUTHENTICATE);
-        }
-    }
-}
-
 #[derive(Debug)]
 struct TooManyRedirects;
 
@@ -279,37 +272,17 @@ impl FollowRedirectPolicy {
         }
     }
 
-    pub(crate) fn with_referer(mut self, referer: bool) -> Self {
+    pub(crate) const fn with_referer(mut self, referer: bool) -> Self {
         self.referer = referer;
         self
     }
 
-    pub(crate) fn with_https_only(mut self, https_only: bool) -> Self {
+    pub(crate) const fn with_https_only(mut self, https_only: bool) -> Self {
         self.https_only = https_only;
         self
     }
 }
 
-fn make_referer(next: &Uri, previous: &Uri) -> Option<HeaderValue> {
-    if next.is_http() && previous.is_https() {
-        return None;
-    }
-
-    let mut parts = previous.clone().into_parts();
-    if let Some(authority) = &mut parts.authority {
-        let host_port = authority.host();
-        let port = authority.port();
-        let new_authority = match port {
-            Some(port) => format!("{}:{}", host_port, port),
-            None => host_port.to_string(),
-        };
-        parts.authority = Some(new_authority.parse().ok()?);
-    }
-
-    let referer = Uri::from_parts(parts).ok()?;
-    referer.to_string().parse().ok()
-}
-
 impl policy::Policy<Body, BoxError> for FollowRedirectPolicy {
     fn redirect(&mut self, attempt: &policy::Attempt<'_>) -> Result<policy::Action, BoxError> {
         // Parse the next URI from the attempt.
@@ -326,7 +299,7 @@ impl policy::Policy<Body, BoxError> for FollowRedirectPolicy {
             .expect("FollowRedirectPolicy should always have a policy set");
 
         // Check if the next URI is already in the list of URLs.
-        match policy.check(attempt.status(), next_uri, &self.uris) {
+        match policy.check(attempt.status(), attempt.headers(), next_uri, &self.uris) {
             ActionKind::Follow => {
                 // Validate the next URI's scheme.
                 if !next_uri.is_http() && !next_uri.is_https() {
@@ -378,6 +351,43 @@ impl policy::Policy<Body, BoxError> for FollowRedirectPolicy {
     }
 }
 
+fn make_referer(next: &Uri, previous: &Uri) -> Option<HeaderValue> {
+    if next.is_http() && previous.is_https() {
+        return None;
+    }
+
+    let referer = {
+        let mut parts = previous.clone().into_parts();
+        if let Some(authority) = &mut parts.authority {
+            let host = authority.host();
+            parts.authority = match authority.port() {
+                Some(port) => {
+                    Authority::from_maybe_shared(Bytes::from(format!("{host}:{port}"))).ok()
+                }
+                None => host.parse().ok(),
+            };
+        }
+        Uri::from_parts(parts).ok()?
+    };
+
+    HeaderValue::from_maybe_shared(Bytes::from(referer.to_string())).ok()
+}
+
+fn remove_sensitive_headers(headers: &mut HeaderMap, next: &Uri, previous: &[Uri]) {
+    if let Some(previous) = previous.last() {
+        let cross_host = next.host() != previous.host()
+            || next.port() != previous.port()
+            || next.scheme() != previous.scheme();
+        if cross_host {
+            headers.remove(AUTHORIZATION);
+            headers.remove(COOKIE);
+            headers.remove("cookie2");
+            headers.remove(PROXY_AUTHORIZATION);
+            headers.remove(WWW_AUTHENTICATE);
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -390,14 +400,14 @@ mod tests {
             .map(|i| Uri::try_from(&format!("http://a.b/c/{i}")).unwrap())
             .collect::<Vec<_>>();
 
-        match policy.check(StatusCode::FOUND, &next, &previous) {
+        match policy.check(StatusCode::FOUND, &HeaderMap::new(), &next, &previous) {
             ActionKind::Follow => (),
             other => panic!("unexpected {other:?}"),
         }
 
         previous.push(Uri::try_from("http://a.b.d/e/33").unwrap());
 
-        match policy.check(StatusCode::FOUND, &next, &previous) {
+        match policy.check(StatusCode::FOUND, &HeaderMap::new(), &next, &previous) {
             ActionKind::Error(err) if err.is::<TooManyRedirects>() => (),
             other => panic!("unexpected {other:?}"),
         }
@@ -409,7 +419,7 @@ mod tests {
         let next = Uri::try_from("http://x.y/z").unwrap();
         let previous = vec![Uri::try_from("http://a.b/c").unwrap()];
 
-        match policy.check(StatusCode::FOUND, &next, &previous) {
+        match policy.check(StatusCode::FOUND, &HeaderMap::new(), &next, &previous) {
             ActionKind::Error(err) if err.is::<TooManyRedirects>() => (),
             other => panic!("unexpected {other:?}"),
         }
@@ -426,13 +436,13 @@ mod tests {
         });
 
         let next = Uri::try_from("http://bar/baz").unwrap();
-        match policy.check(StatusCode::FOUND, &next, &[]) {
+        match policy.check(StatusCode::FOUND, &HeaderMap::new(), &next, &[]) {
             ActionKind::Follow => (),
             other => panic!("unexpected {other:?}"),
         }
 
         let next = Uri::try_from("http://foo/baz").unwrap();
-        match policy.check(StatusCode::FOUND, &next, &[]) {
+        match policy.check(StatusCode::FOUND, &HeaderMap::new(), &next, &[]) {
             ActionKind::Stop => (),
             other => panic!("unexpected {other:?}"),
         }