From 4151aab3b442bd217b40532e46a21470ca010b1f Mon Sep 17 00:00:00 2001
From: Martin Algesten <martin@algesten.se>
Date: Tue, 18 Dec 2018 13:34:05 +0100
Subject: [PATCH] redirects 0

---
 src/body.rs        |  4 ++--
 src/request.rs     |  8 ++++++--
 src/test/simple.rs | 38 ++++++++++++++++++++++++++++++++++++++
 src/unit.rs        | 13 +++++++------
 4 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/src/body.rs b/src/body.rs
index 043282e..5795e8c 100644
--- a/src/body.rs
+++ b/src/body.rs
@@ -2,12 +2,12 @@ use crate::stream::Stream;
 use chunked_transfer;
 use std::io::{copy, empty, Cursor, Read, Result as IoResult};
 
+#[cfg(feature = "charset")]
+use crate::response::DEFAULT_CHARACTER_SET;
 #[cfg(feature = "charset")]
 use encoding::label::encoding_from_whatwg_label;
 #[cfg(feature = "charset")]
 use encoding::EncoderTrap;
-#[cfg(feature = "charset")]
-use crate::response::DEFAULT_CHARACTER_SET;
 
 #[cfg(feature = "json")]
 use super::SerdeValue;
diff --git a/src/request.rs b/src/request.rs
index 7bd2f3d..8b4002c 100644
--- a/src/request.rs
+++ b/src/request.rs
@@ -100,7 +100,7 @@ impl Request {
             .and_then(|url| {
                 let reader = payload.into_read();
                 let unit = Unit::new(&self, &url, true, &reader);
-                connect(&self, unit, &self.method, true, self.redirects, reader)
+                connect(&self, unit, &self.method, true, 0, reader)
             })
             .unwrap_or_else(|e| e.into())
     }
@@ -348,7 +348,11 @@ impl Request {
 
     /// How many redirects to follow.
     ///
-    /// Defaults to `5`.
+    /// Defaults to `5`. Set to `0` to avoid redirects and instead
+    /// get a response object with the 3xx status code.
+    ///
+    /// If the redirect count hits this limit (and it's > 0), a synthetic 500 error
+    /// response is produced.
     ///
     /// ```
     /// let r = ureq::get("/my_page")
diff --git a/src/test/simple.rs b/src/test/simple.rs
index 000f0f3..2f1c0ca 100644
--- a/src/test/simple.rs
+++ b/src/test/simple.rs
@@ -138,3 +138,41 @@ fn non_ascii_header() {
     assert_eq!(resp.status(), 500);
     assert_eq!(resp.status_text(), "Bad Header");
 }
+
+#[test]
+fn redirect_on() {
+    test::set_handler("/redirect_on1", |_| {
+        test::make_response(302, "Go here", vec!["Location: /redirect_on2"], vec![])
+    });
+    test::set_handler("/redirect_on2", |_| {
+        test::make_response(200, "OK", vec!["x-foo: bar"], vec![])
+    });
+    let resp = get("test://host/redirect_on1").call();
+    assert_eq!(resp.status(), 200);
+    assert!(resp.has("x-foo"));
+    assert_eq!(resp.header("x-foo").unwrap(), "bar");
+}
+
+#[test]
+fn redirect_many() {
+    test::set_handler("/redirect_many1", |_| {
+        test::make_response(302, "Go here", vec!["Location: /redirect_many2"], vec![])
+    });
+    test::set_handler("/redirect_many2", |_| {
+        test::make_response(302, "Go here", vec!["Location: /redirect_many3"], vec![])
+    });
+    let resp = get("test://host/redirect_many1").redirects(1).call();
+    assert_eq!(resp.status(), 500);
+    assert_eq!(resp.status_text(), "Too Many Redirects");
+}
+
+#[test]
+fn redirect_off() {
+    test::set_handler("/redirect_off", |_| {
+        test::make_response(302, "Go here", vec!["Location: somewhere.else"], vec![])
+    });
+    let resp = get("test://host/redirect_off").redirects(0).call();
+    assert_eq!(resp.status(), 302);
+    assert!(resp.has("Location"));
+    assert_eq!(resp.header("Location").unwrap(), "somewhere.else");
+}
diff --git a/src/unit.rs b/src/unit.rs
index d3ad240..536d0e8 100644
--- a/src/unit.rs
+++ b/src/unit.rs
@@ -111,7 +111,7 @@ pub fn connect(
     unit: Unit,
     method: &str,
     use_pooled: bool,
-    redirects: u32,
+    redirect_count: u32,
     body: SizedReader,
 ) -> Result<Response, Error> {
     //
@@ -125,7 +125,7 @@ pub fn connect(
         if is_recycled {
             // we try open a new connection, this time there will be
             // no connection in the pool. don't use it.
-            return connect(req, unit, method, false, redirects, body);
+            return connect(req, unit, method, false, redirect_count, body);
         } else {
             // not a pooled connection, propagate the error.
             return Err(send_result.unwrap_err().into());
@@ -142,8 +142,8 @@ pub fn connect(
     save_cookies(&unit, &resp);
 
     // handle redirects
-    if resp.redirect() {
-        if redirects == 0 {
+    if resp.redirect() && req.redirects > 0 {
+        if redirect_count == req.redirects {
             return Err(Error::TooManyRedirects);
         }
 
@@ -162,7 +162,7 @@ pub fn connect(
                     let empty = Payload::Empty.into_read();
                     // recreate the unit to get a new hostname and cookies for the new host.
                     let new_unit = Unit::new(req, &new_url, false, &empty);
-                    return connect(req, new_unit, "GET", use_pooled, redirects - 1, empty);
+                    return connect(req, new_unit, "GET", use_pooled, redirect_count + 1, empty);
                 }
                 , _ => (),
                 // reinstate this with expect-100
@@ -171,7 +171,8 @@ pub fn connect(
         }
     }
 
-    // since it is not a redirect, give away the incoming stream to the response object
+    // since it is not a redirect, or we're not following redirects,
+    // give away the incoming stream to the response object
     response::set_stream(&mut resp, Some(unit), stream);
 
     // release the response