Line data    Source code

       1              : use crate::config::Host;
       2              : use crate::Error;
       3              : use std::future::Future;
       4              : use std::io;
       5              : use std::time::Duration;
       6              : use tokio::net::{self, TcpStream};
       7              : use tokio::time;
       8              : 
       9            0 : pub(crate) async fn connect_socket(
      10            0 :     host: &Host,
      11            0 :     port: u16,
      12            0 :     connect_timeout: Option<Duration>,
      13            0 : ) -> Result<TcpStream, Error> {
      14            0 :     match host {
      15            0 :         Host::Tcp(host) => {
      16            0 :             let addrs = net::lookup_host((&**host, port))
      17            0 :                 .await
      18            0 :                 .map_err(Error::connect)?;
      19              : 
      20            0 :             let mut last_err = None;
      21              : 
      22            0 :             for addr in addrs {
      23            0 :                 let stream =
      24            0 :                     match connect_with_timeout(TcpStream::connect(addr), connect_timeout).await {
      25            0 :                         Ok(stream) => stream,
      26            0 :                         Err(e) => {
      27            0 :                             last_err = Some(e);
      28            0 :                             continue;
      29              :                         }
      30              :                     };
      31              : 
      32            0 :                 stream.set_nodelay(true).map_err(Error::connect)?;
      33              : 
      34            0 :                 return Ok(stream);
      35              :             }
      36              : 
      37            0 :             Err(last_err.unwrap_or_else(|| {
      38            0 :                 Error::connect(io::Error::new(
      39            0 :                     io::ErrorKind::InvalidInput,
      40            0 :                     "could not resolve any addresses",
      41            0 :                 ))
      42            0 :             }))
      43              :         }
      44              :     }
      45            0 : }
      46              : 
      47            0 : async fn connect_with_timeout<F, T>(connect: F, timeout: Option<Duration>) -> Result<T, Error>
      48            0 : where
      49            0 :     F: Future<Output = io::Result<T>>,
      50            0 : {
      51            0 :     match timeout {
      52            0 :         Some(timeout) => match time::timeout(timeout, connect).await {
      53            0 :             Ok(Ok(socket)) => Ok(socket),
      54            0 :             Ok(Err(e)) => Err(Error::connect(e)),
      55            0 :             Err(_) => Err(Error::connect(io::Error::new(
      56            0 :                 io::ErrorKind::TimedOut,
      57            0 :                 "connection timed out",
      58            0 :             ))),
      59              :         },
      60            0 :         None => match connect.await {
      61            0 :             Ok(socket) => Ok(socket),
      62            0 :             Err(e) => Err(Error::connect(e)),
      63              :         },
      64              :     }
      65            0 : }