From b246f0a9d2e6150b29b5504924edc090d2bf9168 Mon Sep 17 00:00:00 2001
From: Jacob Hoffman-Andrews <github@hoffman-andrews.com>
Date: Sun, 7 Feb 2021 12:29:35 -0800
Subject: [PATCH] Apply deadline across redirects. (#313)

Previously, each redirect could take timeout time, so a series of slow
redirects could run for longer than expected, or indefinitely.
---
 src/request.rs       | 19 +++++++++++++++++--
 src/test/redirect.rs | 27 +++++++++++++++++++++++++++
 src/unit.rs          | 18 +++++++++---------
 3 files changed, 53 insertions(+), 11 deletions(-)

diff --git a/src/request.rs b/src/request.rs
index b28fcdde..32f38cf9 100644
--- a/src/request.rs
+++ b/src/request.rs
@@ -1,5 +1,5 @@
-use std::fmt;
 use std::io::Read;
+use std::{fmt, time};
 
 use url::{form_urlencoded, Url};
 
@@ -116,8 +116,23 @@ impl Request {
         for (name, value) in self.query_params.clone() {
             url.query_pairs_mut().append_pair(&name, &value);
         }
+        let deadline = match self.agent.config.timeout {
+            None => None,
+            Some(timeout) => {
+                let now = time::Instant::now();
+                Some(now.checked_add(timeout).unwrap())
+            }
+        };
+
         let reader = payload.into_read();
-        let unit = Unit::new(&self.agent, &self.method, &url, &self.headers, &reader);
+        let unit = Unit::new(
+            &self.agent,
+            &self.method,
+            &url,
+            &self.headers,
+            &reader,
+            deadline,
+        );
         let response = unit::connect(unit, true, reader).map_err(|e| e.url(url.clone()))?;
 
         if response.status() >= 400 {
diff --git a/src/test/redirect.rs b/src/test/redirect.rs
index 83a5e0ec..b6f1fa00 100644
--- a/src/test/redirect.rs
+++ b/src/test/redirect.rs
@@ -1,6 +1,8 @@
 use std::{
     io::{self, Write},
     net::TcpStream,
+    thread,
+    time::Duration,
 };
 use testserver::{self, TestServer};
 
@@ -158,6 +160,31 @@ fn redirect_308() {
     assert_eq!(resp.get_url(), "test://host/valid_response");
 }
 
+#[test]
+fn redirects_hit_timeout() {
+    // A chain of 2 redirects and an OK, each of which takes 50ms; we set a
+    // timeout of 100ms and should error.
+    test::set_handler("/redirect_sleep1", |_| {
+        thread::sleep(Duration::from_millis(50));
+        test::make_response(302, "Go here", vec!["Location: /redirect_sleep2"], vec![])
+    });
+    test::set_handler("/redirect_sleep2", |_| {
+        thread::sleep(Duration::from_millis(50));
+        test::make_response(302, "Go here", vec!["Location: /ok"], vec![])
+    });
+    test::set_handler("/ok", |_| {
+        thread::sleep(Duration::from_millis(50));
+        test::make_response(200, "Go here", vec![], vec![])
+    });
+    let req = crate::builder().timeout(Duration::from_millis(100)).build();
+    let result = req.get("test://host/redirect_sleep1").call();
+    assert!(
+        matches!(result, Err(Error::Transport(_))),
+        "expected Transport error, got {:?}",
+        result
+    );
+}
+
 #[test]
 fn too_many_redirects() {
     for i in 0..10_000 {
diff --git a/src/unit.rs b/src/unit.rs
index 2c6be0a9..65dc49bb 100644
--- a/src/unit.rs
+++ b/src/unit.rs
@@ -38,6 +38,7 @@ impl Unit {
         url: &Url,
         headers: &Vec<Header>,
         body: &SizedReader,
+        deadline: Option<time::Instant>,
     ) -> Self {
         //
 
@@ -99,14 +100,6 @@ impl Unit {
             .cloned()
             .collect();
 
-        let deadline = match agent.config.timeout {
-            None => None,
-            Some(timeout) => {
-                let now = time::Instant::now();
-                Some(now.checked_add(timeout).unwrap())
-            }
-        };
-
         Unit {
             agent: agent.clone(),
             method: method.to_string(),
@@ -213,7 +206,14 @@ pub(crate) fn connect(
         history.push(unit.url.to_string());
         body = Payload::Empty.into_read();
         // recreate the unit to get a new hostname and cookies for the new host.
-        unit = Unit::new(&unit.agent, &new_method, &new_url, &unit.headers, &body);
+        unit = Unit::new(
+            &unit.agent,
+            &new_method,
+            &new_url,
+            &unit.headers,
+            &body,
+            unit.deadline,
+        );
     };
     resp.history = history;
     Ok(resp)