Repository: algesten/ureq
Pull request #288 (2021-01-04 01:22:20)

Fix 307/308 redirects (again, for 2.0)
I guess the unit test got backported but never added to `main`.

Fixes the issue in https://github.com/deadlinks/cargo-deadlinks/pull/134.

Also add some improvements to registration of test handlers, and add env_logger support to tests.


Patch:
diff --git a/src/unit.rs b/src/unit.rs
index f0d75a0..22e6ce3 100644
--- a/src/unit.rs
+++ b/src/unit.rs
@@ -170,8 +170,6 @@ pub(crate) fn connect(
     body: SizedReader,
     previous: Option<Arc<Response>>,
 ) -> Result<Response, Error> {
-    //
-
     let host = unit
         .url
         .host_str()
@@ -273,7 +271,9 @@ pub(crate) fn connect(
                 307 | 308 if ["GET", "HEAD", "OPTIONS", "TRACE"].contains(&method.as_str()) => {
                     let empty = Payload::Empty.into_read();
                     debug!("redirect {} {} -> {}", resp.status(), url, new_url);
-                    return connect(unit, use_pooled, empty, Some(Arc::new(resp)));
+                    // recreate the unit to get a new hostname and cookies for the new host.
+                    let new_unit = Unit::new(&unit.agent, &unit.method, &new_url, &unit.headers, &empty);
+                    return connect(new_unit, use_pooled, empty, Some(Arc::new(resp)));
                 }
                 _ => (),
             };


Test patch:
diff --git a/src/test/mod.rs b/src/test/mod.rs
index 15a918d..15f2bc0 100644
--- a/src/test/mod.rs
+++ b/src/test/mod.rs
@@ -24,8 +24,14 @@ pub(crate) fn set_handler<H>(path: &str, handler: H)
 where
     H: Fn(&Unit) -> Result<Stream, Error> + Send + 'static,
 {
-    let mut handlers = TEST_HANDLERS.lock().unwrap();
-    handlers.insert(path.to_string(), Box::new(handler));
+    let path = path.to_string();
+    let handler = Box::new(handler);
+    // See `resolve_handler` below for why poisoning isn't necessary.
+    let mut handlers = match TEST_HANDLERS.lock() {
+        Ok(h) => h,
+        Err(poison) => poison.into_inner(),
+    };
+    handlers.insert(path, handler);
 }
 
 #[allow(clippy::write_with_newline)]
@@ -46,8 +52,17 @@ pub(crate) fn make_response(
 }
 
 pub(crate) fn resolve_handler(unit: &Unit) -> Result<Stream, Error> {
-    let mut handlers = TEST_HANDLERS.lock().unwrap();
     let path = unit.url.path();
-    let handler = handlers.remove(path).unwrap();
+    // The only way this can panic is if
+    // 1. `remove(path).unwrap()` panics, in which case the HANDLERS haven't been modified.
+    // 2. `make_hash` for `handlers.insert` panics (in `set_handler`), in which case the HANDLERS haven't been modified.
+    // In all cases, another test will fail as a result, so it's ok to continue other tests in parallel.
+    let mut handlers = match TEST_HANDLERS.lock() {
+        Ok(h) => h,
+        Err(poison) => poison.into_inner(),
+    };
+    let handler = handlers.remove(path)
+        .unwrap_or_else(|| panic!("call make_response(\"{}\") before fetching it in tests (or if you did make it, avoid fetching it more than once)", path));
+    drop(handlers);
     handler(unit)
 }
diff --git a/src/test/redirect.rs b/src/test/redirect.rs
index 7a88a79..d360702 100644
--- a/src/test/redirect.rs
+++ b/src/test/redirect.rs
@@ -127,3 +127,17 @@ fn redirect_post() {
     assert!(resp.has("x-foo"));
     assert_eq!(resp.header("x-foo").unwrap(), "bar");
 }
+
+#[test]
+fn redirect_308() {
+    test::set_handler("/redirect_get3", |_| {
+        test::make_response(308, "Go here", vec!["Location: /valid_response"], vec![])
+    });
+    test::set_handler("/valid_response", |unit| {
+        assert_eq!(unit.method, "GET");
+        test::make_response(200, "OK", vec![], vec![])
+    });
+    let resp = get("test://host/redirect_get3").call().unwrap();
+    assert_eq!(resp.status(), 200);
+    assert_eq!(resp.get_url(), "test://host/valid_response");
+}
diff --git a/src/testserver.rs b/src/testserver.rs
index 3cdef51..d1f3d37 100644
--- a/src/testserver.rs
+++ b/src/testserver.rs
@@ -13,6 +13,9 @@ use crate::{Agent, AgentBuilder};
 // An agent to be installed by default for tests and doctests, such
 // that all hostnames resolve to a TestServer on localhost.
 pub(crate) fn test_agent() -> Agent {
+    #[cfg(test)]
+    let _ = env_logger::try_init();
+
     let testserver = TestServer::new(|mut stream: TcpStream| -> io::Result<()> {
         let headers = read_request(&stream);
         if headers.0.is_empty() {