ureq 3.5.1

mod read;

use std::io::{Read, Write};
use std::net::{Shutdown, SocketAddr, TcpStream};
use std::sync::{Arc, Weak};
use std::{fmt, io, time};

use crate::config::Config;
use crate::unversioned::transport::time::Instant;
use crate::util::IoResultExt;
use crate::{Error, Timeout};

use self::read::read_retrying_interrupts;

use super::ResolvedSocketAddrs;
use super::chain::Either;

use super::time::Duration;
use super::{Buffers, ConnectionDetails, Connector, LazyBuffers, NextTimeout};
use super::{ShutdownHandle, Transport};

#[derive(Default)]
/// Connector for regular TCP sockets.
pub struct TcpConnector(());

impl<In: Transport> Connector<In> for TcpConnector {
    type Out = Either<In, TcpTransport>;

    fn connect(
        &self,
        details: &ConnectionDetails,
        chained: Option<In>,
    ) -> Result<Option<Self::Out>, Error> {
        if chained.is_some() {
            // The chained connection overrides whatever we were to open here.
            // In the DefaultConnector chain this would be a SOCKS proxy connection.
            trace!("Skip");
            return Ok(chained.map(Either::A));
        }

        let config = &details.config;
        let stream = try_connect(
            &details.addrs,
            details.now,
            details.timeout,
            details.current_time.clone(),
            config,
        )?;

        let buffers = LazyBuffers::new(config.input_buffer_size(), config.output_buffer_size());
        let transport = TcpTransport::new(stream, buffers);
        (details.register_shutdown)(transport.shutdown_handle())?;

        Ok(Some(Either::B(transport)))
    }
}

impl ShutdownHandle for TcpStream {
    fn shutdown(&self) {
        // Shutdown can wake a read with EOF rather than an error. We considered
        // distinguishing this for close-delimited bodies (e.g. Connection: close
        // without Content-Length or chunked framing), which can then appear
        // complete despite missing data. That framing already cannot distinguish
        // a complete body from a premature connection close. With Content-Length,
        // missing bytes still produce an error.
        //
        // We deliberately accept the close-delimited ambiguity: guaranteeing a
        // cancellation error would require extra synchronized state, such as an
        // atomic flag checked by the I/O path. That edge case does not justify
        // the machinery. Sharing the socket lets us interrupt I/O directly,
        // without a mutex around it or a separate cancellation flag.
        debug!("Shutdown: {:?}", self.peer_addr().ok());
        // An error means the socket is already closed or shut down.
        let _ = TcpStream::shutdown(self, Shutdown::Both);
    }
}

pub(super) fn try_connect(
    addrs: &ResolvedSocketAddrs,
    start: Instant,
    timeout: NextTimeout,
    current_time: Arc<dyn Fn() -> Instant + Send + Sync + 'static>,
    config: &Config,
) -> Result<TcpStream, Error> {
    try_connect_with(addrs, start, timeout, current_time, |addr, per_addr| {
        try_connect_single(addr, per_addr, config)
    })
}

fn try_connect_with<T>(
    addrs: &ResolvedSocketAddrs,
    start: Instant,
    timeout: NextTimeout,
    current_time: Arc<dyn Fn() -> Instant + Send + Sync + 'static>,
    mut connect: impl FnMut(SocketAddr, Option<Duration>) -> Result<T, Error>,
) -> Result<T, Error> {
    // The idea here is to give each attempt a budget of the total time to try.
    // For a host returning multiple addresses, we share the budget between them
    // using a geometric series that sums to exactly the total budget.
    //
    // Background: https://curl.se/mail/lib-2021-01/0037.html
    //
    // Example: Timeout is 10 seconds, and the host returns 4 addresses.
    //
    // Address 0: 5.33 seconds (53.3% of budget)
    // Address 1: 2.67 seconds (26.7% of budget)
    // Address 2: 1.33 seconds (13.3% of budget)
    // Address 3: 0.67 seconds (6.7% of budget)
    // Sum: 10.0 seconds
    //
    // For a single address, it gets the full budget (100%).
    // We cap the lowest to 10ms.
    //
    const MIN_PER_ADDRESS_TIMEOUT: Duration = Duration::from_millis(10);

    let num_addrs = addrs.len();

    // Pre-calculate the total weight for the geometric series.
    // For weights [1, 1/2, 1/4, 1/8, ...], the sum is 2 * (1 - 1/2^n)
    let total_weight = 2.0 * (1.0 - 0.5_f64.powi(num_addrs as i32));

    // Start with weight 1.0 for the first address, then halve for each subsequent.
    let mut weight = 1.0_f64;

    // The most recent per-address failure, so that a host whose every address
    // fails reports what actually happened instead of a synthesized refusal.
    let mut last_err: Option<Error> = None;

    for addr in addrs {
        // Calculate this address's timeout using geometric series.
        let per_addr = timeout.not_zero().map(|t| {
            let secs = t.as_secs_f64() * weight / total_weight;
            let timeout = Duration::from_millis((secs * 1000.0) as u64);
            timeout.max(MIN_PER_ADDRESS_TIMEOUT)
        });

        match connect(*addr, per_addr) {
            // First that connects
            Ok(v) => return Ok(v),
            // Intercept errors that concern only this address to try next addrs
            Err(Error::Io(e)) if is_addr_specific_error(&e) => {
                trace!("{} failed: {}", addr, e);
                last_err = Some(Error::Io(e));
                continue;
            }
            Err(e @ Error::Timeout(_)) => {
                // Check if we hit the overall global timeout for the connect.
                let elapsed = current_time().duration_since(start);
                if elapsed > timeout.after {
                    return Err(Error::Timeout(timeout.reason));
                }

                // We still got time to try the next address.
                last_err = Some(e);
            }
            // Other errors bail
            Err(e) => return Err(e),
        }

        // Halve the weight for the next address
        weight /= 2.0;
    }

    debug!("Failed to connect to any resolved address");
    Err(last_err.unwrap_or_else(|| {
        Error::Io(io::Error::new(
            io::ErrorKind::ConnectionRefused,
            "Connection refused",
        ))
    }))
}

/// Whether a failed connect concerns only the address tried, meaning the next
/// resolved address might still succeed.
///
/// `ConnectionRefused` means this address answered and said no. The
/// unreachable/unavailable kinds mean the local network stack rejected this
/// address at routing level without asking anything: the typical case is a
/// host whose resolver returns IPv6 addresses first but which has no IPv6
/// route, where the AAAA connect fails instantly while the A record would
/// have worked (#1184). Browsers and curl mask that condition by moving on to
/// the next address, which is the behavior matched here.
fn is_addr_specific_error(e: &io::Error) -> bool {
    // On Windows, a VPN or firewall can surface the blocked address family as
    // WSAEACCES, which maps to ErrorKind::PermissionDenied. Match the raw OS
    // error to retry this socket error without retrying every permission error (#1184).
    #[cfg(windows)]
    const WSAEACCES: i32 = 10013;
    #[cfg(windows)]
    if e.raw_os_error() == Some(WSAEACCES) {
        return true;
    }

    matches!(
        e.kind(),
        io::ErrorKind::ConnectionRefused
            | io::ErrorKind::HostUnreachable
            | io::ErrorKind::NetworkUnreachable
            | io::ErrorKind::AddrNotAvailable
    )
}

fn try_connect_single(
    addr: SocketAddr,
    per_addr: Option<Duration>,
    config: &Config,
) -> Result<TcpStream, Error> {
    trace!("Try connect TcpStream to {}", addr);

    let maybe_stream = if let Some(when) = per_addr {
        TcpStream::connect_timeout(&addr, *when)
    } else {
        TcpStream::connect(addr)
    }
    .normalize_would_block();

    let stream = match maybe_stream {
        Ok(v) => v,
        Err(e) if e.kind() == io::ErrorKind::TimedOut => {
            // The parent replaces this reason if the overall deadline has expired.
            return Err(Error::Timeout(Timeout::Connect));
        }
        Err(e) => return Err(e.into()),
    };

    if config.no_delay() {
        stream.set_nodelay(true)?;
    }

    debug!("Connected TcpStream to {}", addr);

    Ok(stream)
}

pub struct TcpTransport {
    /// In an `Arc` so that the [`ShutdownHandle`] can share it (weakly) and shut the
    /// socket down while this transport is in use. All I/O goes via `&TcpStream`.
    stream: Arc<TcpStream>,
    buffers: LazyBuffers,
    timeout_write: Option<Duration>,
    timeout_read: Option<Duration>,
}

impl TcpTransport {
    pub fn new(stream: TcpStream, buffers: LazyBuffers) -> TcpTransport {
        TcpTransport {
            stream: Arc::new(stream),
            buffers,
            timeout_read: None,
            timeout_write: None,
        }
    }

    /// A handle to shut down the socket while this transport is in use.
    pub(crate) fn shutdown_handle(&self) -> Weak<dyn ShutdownHandle> {
        Arc::downgrade(&self.stream) as Weak<dyn ShutdownHandle>
    }
}

// The goal here is to only cause a syscall to set the timeout if it's necessary.
fn maybe_update_timeout(
    timeout: NextTimeout,
    previous: &mut Option<Duration>,
    stream: &TcpStream,
    f: impl Fn(&TcpStream, Option<time::Duration>) -> io::Result<()>,
) -> io::Result<()> {
    let maybe_timeout = timeout.not_zero();

    if maybe_timeout != *previous {
        update_timeout(previous, maybe_timeout, || {
            (f)(stream, maybe_timeout.map(|t| *t))
        })?;
    }

    Ok(())
}

fn update_timeout(
    previous: &mut Option<Duration>,
    timeout: Option<Duration>,
    set_timeout: impl FnOnce() -> io::Result<()>,
) -> io::Result<()> {
    match set_timeout() {
        Ok(()) => *previous = timeout,
        // macOS can reject timeout changes after the peer closes with unread
        // request data. The response may still be readable: let I/O report the
        // outcome instead. Leave the cache unchanged since the setter failed.
        Err(e) if cfg!(target_os = "macos") && e.kind() == io::ErrorKind::InvalidInput => {}
        Err(e) => return Err(e),
    }
    Ok(())
}

impl Transport for TcpTransport {
    fn buffers(&mut self) -> &mut dyn Buffers {
        &mut self.buffers
    }

    fn transmit_output(&mut self, amount: usize, timeout: NextTimeout) -> Result<(), Error> {
        let timeout = timeout.for_write().check()?;
        maybe_update_timeout(
            timeout,
            &mut self.timeout_write,
            &self.stream,
            TcpStream::set_write_timeout,
        )?;

        let output = &self.buffers.output()[..amount];
        match (&*self.stream).write_all(output).normalize_would_block() {
            Ok(v) => Ok(v),
            Err(e) if e.kind() == io::ErrorKind::TimedOut => Err(Error::Timeout(timeout.reason)),
            Err(e) => Err(e.into()),
        }?;

        Ok(())
    }

    fn await_input(&mut self, timeout: NextTimeout) -> Result<bool, Error> {
        let timeout = timeout.for_read().check()?;
        // Proceed to fill the buffers from the TcpStream
        maybe_update_timeout(
            timeout,
            &mut self.timeout_read,
            &self.stream,
            TcpStream::set_read_timeout,
        )?;

        let input = self.buffers.input_append_buf();
        let started = time::Instant::now();
        let mut stream = &*self.stream;
        let result = read_retrying_interrupts(
            &mut stream,
            input,
            timeout.not_zero().map(|t| *t),
            || started.elapsed(),
            |stream, left| {
                // Keep the cache in step so maybe_update_timeout() later restores the full value.
                update_timeout(&mut self.timeout_read, Some(left.into()), || {
                    stream.set_read_timeout(Some(left))
                })
            },
        );
        let amount = match result.normalize_would_block() {
            Ok(v) => Ok(v),
            Err(e) if e.kind() == io::ErrorKind::TimedOut => Err(Error::Timeout(timeout.reason)),
            Err(e) => Err(e.into()),
        }?;
        self.buffers.input_appended(amount);

        Ok(amount > 0)
    }

    fn is_open(&mut self) -> bool {
        probe_tcp_stream(&self.stream).unwrap_or(false)
    }
}

fn probe_tcp_stream(mut stream: &TcpStream) -> Result<bool, Error> {
    // Temporary do non-blocking IO
    stream.set_nonblocking(true)?;

    let mut buf = [0];
    match stream.read(&mut buf) {
        Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
            // This is the correct condition. There should be no waiting
            // bytes, and therefore reading would block
        }
        // Any bytes read means the server sent some garbage we didn't ask for
        Ok(_) => {
            debug!("Unexpected bytes from server. Closing connection");
            return Ok(false);
        }
        // Errors such as closed connection
        Err(_) => return Ok(false),
    };

    // Reset back to blocking
    stream.set_nonblocking(false)?;

    Ok(true)
}

impl fmt::Debug for TcpConnector {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("TcpConnector").finish()
    }
}

impl fmt::Debug for TcpTransport {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("TcpTransport")
            .field("addr", &self.stream.peer_addr().ok())
            .finish()
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn successful_timeout_update_changes_cache() {
        let mut previous = Some(Duration::from_secs(5));
        update_timeout(&mut previous, Some(Duration::from_secs(2)), || Ok(())).unwrap();
        assert_eq!(previous, Some(Duration::from_secs(2)));
        update_timeout(&mut previous, None, || Ok(())).unwrap();
        assert_eq!(previous, None);
    }

    #[test]
    fn failed_timeout_update_preserves_cache() {
        let mut previous = Some(Duration::from_secs(5));
        let error = update_timeout(&mut previous, Some(Duration::from_secs(2)), || {
            Err(io::ErrorKind::PermissionDenied.into())
        })
        .unwrap_err();
        assert_eq!(error.kind(), io::ErrorKind::PermissionDenied);
        assert_eq!(previous, Some(Duration::from_secs(5)));
    }

    #[test]
    fn invalid_timeout_update_is_tolerated_only_on_macos() {
        for initial in [None, Some(Duration::from_secs(5))] {
            let mut previous = initial;
            let result = update_timeout(&mut previous, Some(Duration::from_secs(2)), || {
                Err(io::ErrorKind::InvalidInput.into())
            });
            if cfg!(target_os = "macos") {
                result.unwrap();
            } else {
                assert_eq!(result.unwrap_err().kind(), io::ErrorKind::InvalidInput);
            }
            assert_eq!(previous, initial);
        }
    }

    #[cfg(target_os = "macos")]
    #[test]
    fn interrupted_read_survives_rejected_timeout_update() {
        struct InterruptedOnce<'a>(bool, &'a [u8]);
        impl Read for InterruptedOnce<'_> {
            fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
                if std::mem::take(&mut self.0) {
                    return Err(io::ErrorKind::Interrupted.into());
                }
                self.1.read(buf)
            }
        }

        let response = b"HTTP/1.1 302 Found\r\nContent-Length: 0\r\n\r\n";
        let mut reader = InterruptedOnce(true, response);
        let mut buf = [0; 128];
        let mut previous = Some(Duration::from_secs(5));
        let amount = read_retrying_interrupts(
            &mut reader,
            &mut buf,
            Some(time::Duration::from_secs(5)),
            || time::Duration::from_secs(1),
            |_, left| {
                update_timeout(&mut previous, Some(left.into()), || {
                    Err(io::Error::from_raw_os_error(22))
                })
            },
        )
        .unwrap();
        assert_eq!(&buf[..amount], response);
        assert_eq!(previous, Some(Duration::from_secs(5)));
    }

    // Script connect outcomes so routing failures and timeouts are deterministic.
    fn scripted_connect(
        outcomes: Vec<Result<(), Error>>,
        elapsed: Duration,
    ) -> (Result<(), Error>, usize) {
        let mut addrs = ResolvedSocketAddrs::from_fn(|_| SocketAddr::from(([0, 0, 0, 0], 0)));
        for i in 0..outcomes.len() {
            addrs.push(SocketAddr::from(([127, 0, 0, 1], 10000 + i as u16)));
        }
        let start = Instant::now();
        let mut outcomes = outcomes.into_iter();
        let mut attempts = 0;
        let result = try_connect_with(
            &addrs,
            start,
            NextTimeout {
                after: Duration::from_secs(10),
                reason: Timeout::Global,
                ..NextTimeout::default()
            },
            Arc::new(move || start + elapsed),
            |addr, _| {
                assert_eq!(addr, addrs[attempts]);
                attempts += 1;
                outcomes.next().unwrap()
            },
        );
        (result, attempts)
    }

    fn io_error(kind: io::ErrorKind) -> Result<(), Error> {
        Err(io::Error::from(kind).into())
    }

    #[test]
    fn timeout_replaces_earlier_unreachable_error() {
        let (result, attempts) = scripted_connect(
            vec![
                io_error(io::ErrorKind::NetworkUnreachable),
                Err(Error::Timeout(Timeout::Connect)),
            ],
            Duration::from_secs(6),
        );
        assert_eq!(attempts, 2);
        assert!(matches!(result, Err(Error::Timeout(Timeout::Connect))));
    }

    #[test]
    fn addr_specific_failures_and_timeout_fall_back_to_success() {
        for failure in [
            io_error(io::ErrorKind::ConnectionRefused),
            io_error(io::ErrorKind::HostUnreachable),
            io_error(io::ErrorKind::NetworkUnreachable),
            io_error(io::ErrorKind::AddrNotAvailable),
            Err(Error::Timeout(Timeout::Connect)),
        ] {
            let (result, attempts) = scripted_connect(
                vec![failure, Ok(()), io_error(io::ErrorKind::PermissionDenied)],
                Duration::from_secs(1),
            );
            assert!(result.is_ok());
            assert_eq!(attempts, 2);
        }
    }

    #[test]
    fn last_io_error_replaces_timeout_and_preserves_details() {
        let (result, attempts) = scripted_connect(
            vec![
                Err(Error::Timeout(Timeout::Connect)),
                Err(io::Error::new(io::ErrorKind::HostUnreachable, "last address").into()),
            ],
            Duration::from_secs(1),
        );
        assert_eq!(attempts, 2);
        let Err(Error::Io(error)) = result else {
            panic!("expected last I/O error: {result:?}");
        };
        assert_eq!(error.kind(), io::ErrorKind::HostUnreachable);
        assert_eq!(error.to_string(), "last address");
    }

    #[test]
    fn all_timeouts_report_connect_timeout() {
        let (result, attempts) = scripted_connect(
            vec![
                Err(Error::Timeout(Timeout::Connect)),
                Err(Error::Timeout(Timeout::Connect)),
            ],
            Duration::from_secs(6),
        );
        assert_eq!(attempts, 2);
        assert!(matches!(result, Err(Error::Timeout(Timeout::Connect))));
    }

    #[test]
    fn overall_timeout_takes_precedence_and_stops_attempts() {
        let (result, attempts) = scripted_connect(
            vec![
                io_error(io::ErrorKind::NetworkUnreachable),
                Err(Error::Timeout(Timeout::Connect)),
                Ok(()),
            ],
            Duration::from_secs(11),
        );
        assert_eq!(attempts, 2);
        assert!(matches!(result, Err(Error::Timeout(Timeout::Global))));
    }

    #[test]
    fn other_io_errors_stop_attempts() {
        let (result, attempts) = scripted_connect(
            vec![io_error(io::ErrorKind::PermissionDenied), Ok(())],
            Duration::from_secs(1),
        );
        assert_eq!(attempts, 1);
        assert!(matches!(result, Err(Error::Io(e)) if e.kind() == io::ErrorKind::PermissionDenied));
    }

    #[test]
    fn empty_address_list_reports_refusal() {
        let (result, attempts) = scripted_connect(vec![], Duration::from_secs(0));
        assert_eq!(attempts, 0);
        assert!(
            matches!(result, Err(Error::Io(e)) if e.kind() == io::ErrorKind::ConnectionRefused)
        );
    }

    #[test]
    fn addr_specific_errors_try_the_next_addr() {
        for kind in [
            io::ErrorKind::ConnectionRefused,
            io::ErrorKind::HostUnreachable,
            io::ErrorKind::NetworkUnreachable,
            io::ErrorKind::AddrNotAvailable,
        ] {
            assert!(
                is_addr_specific_error(&io::Error::from(kind)),
                "{kind:?} should move on to the next address"
            );
        }
        for kind in [io::ErrorKind::TimedOut, io::ErrorKind::PermissionDenied] {
            assert!(
                !is_addr_specific_error(&io::Error::from(kind)),
                "{kind:?} should bail"
            );
        }
    }

    #[test]
    #[cfg(windows)]
    fn wsaeacces_tries_the_next_addr() {
        assert!(is_addr_specific_error(&io::Error::from_raw_os_error(10013)));
    }
}