From 2f53545637a784ce945e5b1fd9a0fcd4bb50c6e2 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 30 Aug 2025 19:37:23 +0800 Subject: [PATCH 1/2] feat(redirect): allow custom redirects to access response headers --- examples/request_with_redirect.rs | 19 +++-- src/client/layer/redirect/future.rs | 2 + src/client/layer/redirect/policy.rs | 15 +++- src/redirect.rs | 115 +++++++++++++++------------- 4 files changed, 89 insertions(+), 62 deletions(-) diff --git a/examples/request_with_redirect.rs b/examples/request_with_redirect.rs index 1ce1b20b4..e1cedfe1d 100644 --- a/examples/request_with_redirect.rs +++ b/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?); diff --git a/src/client/layer/redirect/future.rs b/src/client/layer/redirect/future.rs index d4eac22d5..f19844016 100644 --- a/src/client/layer/redirect/future.rs +++ b/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; diff --git a/src/client/layer/redirect/policy.rs b/src/client/layer/redirect/policy.rs index 3556b5b0a..c077df00c 100644 --- a/src/client/layer/redirect/policy.rs +++ b/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 { @@ -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 { diff --git a/src/redirect.rs b/src/redirect.rs index aa06cd3ac..8d619561c 100644 --- a/src/redirect.rs +++ b/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 { - 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 for FollowRedirectPolicy { fn redirect(&mut self, attempt: &policy::Attempt<'_>) -> Result { // Parse the next URI from the attempt. @@ -326,7 +299,7 @@ impl policy::Policy 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,44 @@ impl policy::Policy for FollowRedirectPolicy { } } +fn make_referer(next: &Uri, previous: &Uri) -> Option { + 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 = authority.host(); + match authority.port() { + Some(port) => { + parts.authority = + Authority::from_maybe_shared(Bytes::from(format!("{host}:{port}"))).ok() + } + None => { + parts.authority = Some(host.parse().ok()?); + } + }; + } + + let referer = 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 +401,14 @@ mod tests { .map(|i| Uri::try_from(&format!("http://a.b/c/{i}")).unwrap()) .collect::>(); - 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::() => (), other => panic!("unexpected {other:?}"), } @@ -409,7 +420,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::() => (), other => panic!("unexpected {other:?}"), } @@ -426,13 +437,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:?}"), } From 7604ddeaaec0cb81b6f76a95acb3704a93c2d82a Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 30 Aug 2025 19:55:12 +0800 Subject: [PATCH 2/2] fmt --- src/redirect.rs | 25 ++++++++++++------------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/src/redirect.rs b/src/redirect.rs index 8d619561c..132b4acf2 100644 --- a/src/redirect.rs +++ b/src/redirect.rs @@ -356,21 +356,20 @@ fn make_referer(next: &Uri, previous: &Uri) -> Option { return None; } - let mut parts = previous.clone().into_parts(); - if let Some(authority) = &mut parts.authority { - let host = authority.host(); - match authority.port() { - Some(port) => { - parts.authority = + 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 => { - parts.authority = Some(host.parse().ok()?); - } - }; - } + } + None => host.parse().ok(), + }; + } + Uri::from_parts(parts).ok()? + }; - let referer = Uri::from_parts(parts).ok()?; HeaderValue::from_maybe_shared(Bytes::from(referer.to_string())).ok() }