Skip to content

Commit

Permalink
Optionally record redirect history
Browse files Browse the repository at this point in the history
  • Loading branch information
harmless-tech committed Jan 17, 2025
1 parent d609a07 commit e7758fd
Show file tree
Hide file tree
Showing 4 changed files with 111 additions and 1 deletion.
26 changes: 26 additions & 0 deletions src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,7 @@ pub struct Config {
max_redirects: u32,
max_redirects_will_error: bool,
redirect_auth_headers: RedirectAuthHeaders,
save_redirect_history: bool,
user_agent: AutoHeaderValue,
accept: AutoHeaderValue,
accept_encoding: AutoHeaderValue,
Expand Down Expand Up @@ -277,6 +278,17 @@ impl Config {
self.redirect_auth_headers
}

/// If we should record a history of every redirect location,
/// including the request and final locations.
///
/// Comes at the cost of allocating/retaining the Uri for
/// every redirect loop.
///
/// Defaults to false
pub fn save_redirect_history(&self) -> bool {
self.save_redirect_history
}

/// Value to use for the `User-Agent` header.
///
/// This can be overridden by setting a `user-agent` header on the request
Expand Down Expand Up @@ -482,6 +494,18 @@ impl<Scope: private::ConfigScope> ConfigBuilder<Scope> {
self
}

/// If we should record a history of every redirect location,
/// including the request and final locations.
///
/// Comes at the cost of allocating/retaining the Uri for
/// every redirect loop.
///
/// Defaults to false
pub fn save_redirect_history(mut self, v: bool) -> Self {
self.config().save_redirect_history = v;
self
}

/// Value to use for the `User-Agent` header.
///
/// This can be overridden by setting a `user-agent` header on the request
Expand Down Expand Up @@ -809,6 +833,7 @@ impl Default for Config {
max_redirects: 10,
max_redirects_will_error: true,
redirect_auth_headers: RedirectAuthHeaders::Never,
save_redirect_history: false,
user_agent: AutoHeaderValue::default(),
accept: AutoHeaderValue::default(),
accept_encoding: AutoHeaderValue::default(),
Expand Down Expand Up @@ -884,6 +909,7 @@ impl fmt::Debug for Config {
.field("no_delay", &self.no_delay)
.field("max_redirects", &self.max_redirects)
.field("redirect_auth_headers", &self.redirect_auth_headers)
.field("save_redirect_history", &self.save_redirect_history)
.field("user_agent", &self.user_agent)
.field("timeouts", &self.timeouts)
.field("max_response_header_size", &self.max_response_header_size)
Expand Down
59 changes: 59 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -847,6 +847,65 @@ pub(crate) mod test {
assert_eq!(response_uri.path(), "/get")
}

#[test]
fn redirect_history_none() {
init_test_log();
let res = get("http://httpbin.org/redirect-to?url=%2Fget")
.call()
.unwrap();
let redirect_history = res.get_redirect_history();
assert_eq!(redirect_history, None)
}

#[test]
fn redirect_history_some() {
init_test_log();
let agent: Agent = Config::builder()
.max_redirects(3)
.max_redirects_will_error(false)
.save_redirect_history(true)
.build()
.into();
let res = agent
.get("http://httpbin.org/redirect-to?url=%2Fget")
.call()
.unwrap();
let redirect_history = res.get_redirect_history();
assert_eq!(
redirect_history,
Some(
vec![
"http://httpbin.org/redirect-to?url=%2Fget".parse().unwrap(),
"http://httpbin.org/get".parse().unwrap()
]
.as_ref()
)
);
let res = agent
.get(
"http://httpbin.org/redirect-to?url=%2Fredirect-to%3F\
url%3D%2Fredirect-to%3Furl%3D%252Fredirect-to%253Furl%253D",
)
.call()
.unwrap();
let redirect_history = res.get_redirect_history();
assert_eq!(
redirect_history,
Some(vec![
"http://httpbin.org/redirect-to?url=%2Fredirect-to%3Furl%3D%2Fredirect-to%3Furl%3D%252Fredirect-to%253Furl%253D".parse().unwrap(),
"http://httpbin.org/redirect-to?url=/redirect-to?url=%2Fredirect-to%3Furl%3D".parse().unwrap(),
"http://httpbin.org/redirect-to?url=/redirect-to?url=".parse().unwrap(),
"http://httpbin.org/redirect-to?url=".parse().unwrap(),
].as_ref())
);
let res = agent.get("https://www.google.com/").call().unwrap();
let redirect_history = res.get_redirect_history();
assert_eq!(
redirect_history,
Some(vec!["https://www.google.com/".parse().unwrap()].as_ref())
);
}

#[test]
fn connect_https_invalid_name() {
let result = get("https://example.com{REQUEST_URI}/").call();
Expand Down
14 changes: 14 additions & 0 deletions src/response.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,20 @@ use crate::http;
#[derive(Debug, Clone)]
pub(crate) struct ResponseUri(pub http::Uri);

#[derive(Debug, Clone)]
pub(crate) struct RedirectHistory(pub Vec<Uri>);

/// Extension trait for `http::Response<Body>` objects
///
/// Allows the user to access the `Uri` in http::Response
pub trait ResponseExt {
/// The Uri we ended up at. This can differ from the request uri when we have followed redirects.
fn get_uri(&self) -> &Uri;

/// The full history of uris, including the request and final uri.
///
/// Returns None when `save_redirect_history` is false.
fn get_redirect_history(&self) -> Option<&[Uri]>;
}

impl ResponseExt for http::Response<Body> {
Expand All @@ -22,4 +30,10 @@ impl ResponseExt for http::Response<Body> {
.expect("uri to have been set")
.0
}

fn get_redirect_history(&self) -> Option<&[Uri]> {
self.extensions()
.get::<RedirectHistory>()
.map(|r| r.0.as_ref())
}
}
13 changes: 12 additions & 1 deletion src/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use crate::body::ResponseInfo;
use crate::config::{Config, RequestLevelConfig, DEFAULT_USER_AGENT};
use crate::http;
use crate::pool::Connection;
use crate::response::ResponseUri;
use crate::response::{RedirectHistory, ResponseUri};
use crate::timings::{CallTimings, CurrentTime};
use crate::transport::time::{Duration, Instant};
use crate::transport::ConnectionDetails;
Expand All @@ -41,6 +41,9 @@ pub(crate) fn run(
.map(Arc::new)
.unwrap_or_else(|| agent.config.clone());

let mut redirect_history: Option<Vec<Uri>> =
config.save_redirect_history().then_some(Vec::new());

let timeouts = config.timeouts();

let mut timings = CallTimings::new(timeouts, CurrentTime::default());
Expand All @@ -67,6 +70,7 @@ pub(crate) fn run(
flow,
&mut body,
redirect_count,
&mut redirect_history,
&mut timings,
)? {
// Follow redirect
Expand Down Expand Up @@ -112,6 +116,7 @@ fn flow_run(
mut flow: Flow<Prepare>,
body: &mut SendBody,
redirect_count: u32,
redirect_history: &mut Option<Vec<Uri>>,
timings: &mut CallTimings,
) -> Result<FlowResult, Error> {
let uri = flow.uri().clone();
Expand Down Expand Up @@ -166,6 +171,12 @@ fn flow_run(
jar.store_response_cookies(iter, &uri);
}

if let Some(history) = redirect_history.as_mut() {
history.push(uri.clone());
response
.extensions_mut()
.insert(RedirectHistory(history.clone()));
}
response.extensions_mut().insert(ResponseUri(uri));

let ret = match response_result {
Expand Down

0 comments on commit e7758fd

Please sign in to comment.