Repository: algesten/ureq
Pull request #959 (2025-01-28 01:22:46)

Use std::sync::OnceLock instead of once_cell::sync::Lazy
This is a revival of #725 but for ureq 3.0 adhering to the "lightweight" spirit of this library.
MSRV is 1.71.1 and `OnceLock` was added in 1.70.0 therefore no bump is required.

I'm not exactly sure if the changes to public structs in `tls/` are OK.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 57fcd13..3dac2b2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,6 @@
 # Unreleased
 
+  * Fix bug parsing partial redirects (#958)
   * Expose typestate variables (#956)
 
 # 3.0.1
diff --git a/Cargo.lock b/Cargo.lock
index 42c6314..d61123b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1418,7 +1418,6 @@ dependencies = [
  "flate2",
  "log",
  "native-tls",
- "once_cell",
  "percent-encoding",
  "rustls",
  "rustls-pemfile",
@@ -1436,9 +1435,9 @@ dependencies = [
 
 [[package]]
 name = "ureq-proto"
-version = "0.2.5"
+version = "0.2.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "42506c1d9fc65bcf2508cb8f9b5e7d21f88e71567fbb7595ca5e5bd5218ca7db"
+checksum = "b930c6bd0de6a824c1bc380378a00f8f437f5da0a441740510a3dc4d02bd677e"
 dependencies = [
  "base64",
  "http",
diff --git a/Cargo.toml b/Cargo.toml
index cd9d03c..5d8e9f5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -56,10 +56,9 @@ features = ["rustls", "platform-verifier", "native-tls", "socks-proxy", "cookies
 
 [dependencies]
 base64 = "0.22.1"
-ureq-proto = "0.2.5"
+ureq-proto = "0.2.6"
 # ureq-proto = { path = "../ureq-proto" }
 log = "0.4.22"
-once_cell = "1.19.0"
 utf-8 = "0.7.6"
 percent-encoding = "2.3.1"
 
diff --git a/src/lib.rs b/src/lib.rs
index 04a809c..b531a73 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -596,11 +596,10 @@ mk_method!(trace, TRACE, WithoutBody);
 
 #[cfg(test)]
 pub(crate) mod test {
-    use std::io;
+    use std::{io, sync::OnceLock};
 
     use assert_no_alloc::AllocDisabler;
     use config::{Config, ConfigBuilder};
-    use once_cell::sync::Lazy;
     use typestate::AgentScope;
 
     use super::*;
@@ -610,8 +609,8 @@ pub(crate) mod test {
     static A: AllocDisabler = AllocDisabler;
 
     pub fn init_test_log() {
-        static INIT_LOG: Lazy<()> = Lazy::new(env_logger::init);
-        *INIT_LOG
+        static INIT_LOG: OnceLock<()> = OnceLock::new();
+        INIT_LOG.get_or_init(|| env_logger::init());
     }
 
     #[test]
@@ -973,6 +972,28 @@ pub(crate) mod test {
         let _ = agent.get("http://cookie.test/cookie-test").call().unwrap();
     }
 
+    #[test]
+    #[cfg(all(feature = "_test", not(feature = "cookies")))]
+    fn partial_redirect_when_following() {
+        init_test_log();
+        // this should work because we follow the redirect and go to /get
+        get("http://my-host.com/partial-redirect").call().unwrap();
+    }
+
+    #[test]
+    #[cfg(feature = "_test")]
+    fn partial_redirect_when_not_following() {
+        init_test_log();
+        // this should fail because we are not following redirects, and the
+        // response is partial before the server is hanging up
+        get("http://my-host.com/partial-redirect")
+            .config()
+            .max_redirects(0)
+            .build()
+            .call()
+            .unwrap_err();
+    }
+
     #[test]
     #[cfg(feature = "_test")]
     fn http_connect_proxy() {
diff --git a/src/run.rs b/src/run.rs
index ad88c1c..439acc0 100644
--- a/src/run.rs
+++ b/src/run.rs
@@ -1,9 +1,8 @@
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 use std::{io, mem};
 
 use http::uri::Scheme;
 use http::{header, HeaderValue, Request, Response, Uri};
-use once_cell::sync::Lazy;
 use ureq_proto::client::flow::state::{Await100, RecvBody, RecvResponse, Redirect, SendRequest};
 use ureq_proto::client::flow::state::{Prepare, SendBody as SendBodyState};
 use ureq_proto::client::flow::{Await100Result, RecvBodyResult};
@@ -153,7 +152,15 @@ fn flow_run(
         SendRequestResult::RecvResponse(flow) => flow,
     };
 
-    let (mut response, response_result) = recv_response(flow, &mut connection, config, timings)?;
+    let is_following_redirects = redirect_count < config.max_redirects();
+
+    let (mut response, response_result) = recv_response(
+        flow,
+        &mut connection,
+        config,
+        timings,
+        is_following_redirects,
+    )?;
 
     info!("{:?}", DebugResponse(&response));
 
@@ -190,7 +197,7 @@ fn flow_run(
             };
 
             if response.status().is_redirection() {
-                if redirect_count < config.max_redirects() {
+                if is_following_redirects {
                     let flow = handler.consume_redirect_body()?;
 
                     FlowResult::Redirect(flow, handler.timings)
@@ -206,7 +213,7 @@ fn flow_run(
         RecvResponseResult::Redirect(flow) => {
             cleanup(connection, flow.must_close_connection(), timings.now());
 
-            if redirect_count < config.max_redirects() {
+            if is_following_redirects {
                 FlowResult::Redirect(flow, mem::take(timings))
             } else if config.max_redirects_do_error() {
                 return Err(Error::TooManyRedirects);
@@ -267,7 +274,9 @@ fn add_headers(
     }
 
     {
-        static ACCEPTS: Lazy<String> = Lazy::new(|| {
+        static ACCEPTS: OnceLock<String> = OnceLock::new();
+
+        let accepts = ACCEPTS.get_or_init(|| {
             #[allow(unused_mut)]
             let mut value = String::with_capacity(10);
             #[cfg(feature = "gzip")]
@@ -278,8 +287,9 @@ fn add_headers(
             value.push_str("br");
             value
         });
+
         if !has_header_accept_enc {
-            if let Some(v) = config.accept_encoding().as_str(&ACCEPTS) {
+            if let Some(v) = config.accept_encoding().as_str(accepts) {
                 // unwrap is ok because above ACCEPTS will produce a valid value,
                 // or the value is user provided in which case it must be valid.
                 let value = HeaderValue::from_str(v).unwrap();
@@ -488,6 +498,7 @@ fn recv_response(
     connection: &mut Connection,
     config: &Config,
     timings: &mut CallTimings,
+    is_following_redirects: bool,
 ) -> Result<(Response<()>, RecvResponseResult<()>), Error> {
     let response = loop {
         let timeout = timings.next_timeout(Timeout::RecvResponse);
@@ -495,7 +506,21 @@ fn recv_response(
 
         let input = connection.buffers().input();
 
-        let (amount, maybe_response) = flow.try_response(input)?;
+        // If cookies are disabled, we can allow ourselves to follow
+        // the redirect as soon as we discover the `Location` header.
+        // There are plenty of broken servers out there that don't send
+        // the finishing \r\n on redirects. With cookies off, we can
+        // handle that situation.
+        //
+        // If cookies are enabled, we risk mising a `Set-Cookie` header
+        // with such a strategy.
+        let cookies_enabled = cfg!(feature = "cookies");
+
+        // If we are not following redirects, do not allow partial returned
+        // 302 responses.
+        let allow_partial_redirect = !cookies_enabled && is_following_redirects;
+
+        let (amount, maybe_response) = flow.try_response(input, allow_partial_redirect)?;
 
         if input.len() > config.max_response_header_size() {
             return Err(Error::LargeResponseHeader(
diff --git a/src/tls/native_tls.rs b/src/tls/native_tls.rs
index f014e9d..61561ea 100644
--- a/src/tls/native_tls.rs
+++ b/src/tls/native_tls.rs
@@ -1,7 +1,7 @@
 use std::convert::TryFrom;
 use std::fmt;
 use std::io::{Read, Write};
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 
 use crate::tls::{RootCerts, TlsProvider};
 use crate::{transport::*, Error};
@@ -9,7 +9,6 @@ use der::pem::LineEnding;
 use der::Document;
 use native_tls::{Certificate, HandshakeError, Identity, TlsConnector};
 use native_tls::{TlsConnectorBuilder, TlsStream};
-use once_cell::sync::OnceCell;
 
 use super::TlsConfig;
 
@@ -18,7 +17,7 @@ use super::TlsConfig;
 /// Requires feature flag **native-tls**.
 #[derive(Default)]
 pub struct NativeTlsConnector {
-    connector: OnceCell<Arc<TlsConnector>>,
+    connector: OnceLock<Arc<TlsConnector>>,
 }
 
 impl<In: Transport> Connector<In> for NativeTlsConnector {
diff --git a/src/tls/rustls.rs b/src/tls/rustls.rs
index e18d89b..0fd01ed 100644
--- a/src/tls/rustls.rs
+++ b/src/tls/rustls.rs
@@ -1,9 +1,8 @@
 use std::convert::TryInto;
 use std::fmt;
 use std::io::{Read, Write};
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 
-use once_cell::sync::OnceCell;
 use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
 use rustls::crypto::CryptoProvider;
 use rustls::{ClientConfig, ClientConnection, RootCertStore, StreamOwned, ALL_VERSIONS};
@@ -23,7 +22,7 @@ use super::TlsConfig;
 /// Requires feature flag **rustls**.
 #[derive(Default)]
 pub struct RustlsConnector {
-    config: OnceCell<Arc<ClientConfig>>,
+    config: OnceLock<Arc<ClientConfig>>,
 }
 
 impl<In: Transport> Connector<In> for RustlsConnector {


Test patch:
diff --git a/src/unversioned/transport/test.rs b/src/unversioned/transport/test.rs
index b1be72e..04df9db 100644
--- a/src/unversioned/transport/test.rs
+++ b/src/unversioned/transport/test.rs
@@ -314,6 +314,21 @@ fn setup_default_handlers(handlers: &mut Vec<TestHandler>) {
         handlers,
     );
 
+    maybe_add(
+        TestHandler::new("/partial-redirect", |_uri, _req, w| {
+            write!(
+                w,
+                "HTTP/1.1 302 OK\r\n\
+                Location: /get\r\n\
+                set-cookie: AEC=AVYB7cpadYFS8ZgaioQ17NnxHl1QcSQ_2aH2WEIg1KGDXD5kjk2HhpGVhfk; \
+                    expires=Mon, 14-Apr-2025 17:23:39 GMT; path=/; domain=.google.com; \
+                    Secure; HttpOnly; SameSite=lax\r\n\
+                " // deliberately omit final \r\n
+            )
+        }),
+        handlers,
+    );
+
     maybe_add(
         TestHandler::new("/cookie-test", |_uri, req, w| {
             let mut all: Vec<_> = req