use std::{
io::{Error, ErrorKind, Result},
net::{SocketAddr, TcpStream, ToSocketAddrs},
time::{Duration, Instant},
};
use polling::{Event, Poller};
use socket2::{Domain, SockAddr, Socket};
pub fn connect<A: ToSocketAddrs>(addr: A, timeout: Option<Duration>) -> Result<TcpStream> {
let mut happy = HappyEyeballs::new()?;
let start = Instant::now();
let timeout_left = || -> Result<Option<Duration>> {
let Some(v) = timeout else {
return Ok(None);
};
Ok(Some(v.checked_sub(Instant::now() - start).ok_or_else(
|| Error::new(ErrorKind::TimedOut, "Connection timed out."),
)?))
};
for a in prepare_addresses(addr)? {
match happy.add(a.into(), Domain::for_address(a)) {
AddOutcome::Connected(tcp) => return Ok(tcp),
AddOutcome::Error => continue,
AddOutcome::InProgress => (),
}
if let Some(sock) = happy.poll_once(timeout_left()?)? {
return Ok(sock);
}
}
while !happy.is_empty() {
if let Some(sock) = happy.poll_once(timeout_left()?)? {
return Ok(sock);
}
}
Err(happy
.error
.take()
.unwrap_or_else(|| Error::new(ErrorKind::InvalidInput, "could not resolve to any address")))
}
fn prepare_addresses<A>(addr: A) -> Result<Vec<SocketAddr>>
where
A: ToSocketAddrs,
{
let (addrs_v4, addrs_v6): (Vec<_>, Vec<_>) = addr.to_socket_addrs()?.partition(|a| match a {
SocketAddr::V4(_) => true,
SocketAddr::V6(_) => false,
});
let mut addrs = Vec::with_capacity(addrs_v4.len() + addrs_v6.len());
let (mut left, mut right) = (addrs_v6.into_iter(), addrs_v4.into_iter());
while let Some(a) = left.next() {
addrs.push(a);
std::mem::swap(&mut left, &mut right);
}
addrs.extend(right);
Ok(addrs)
}
struct HappyEyeballs {
poller: Poller,
error: Option<Error>,
attempts: Vec<Option<(SockAddr, Socket)>>,
attempts_in_progress: usize,
}
impl HappyEyeballs {
fn new() -> Result<Self> {
Ok(Self {
poller: Poller::new()?,
error: None,
attempts: Vec::new(),
attempts_in_progress: 0,
})
}
fn is_empty(&self) -> bool {
self.attempts_in_progress == 0
}
fn set_error(&mut self, error: Error) {
self.error.get_or_insert(error);
}
fn add(&mut self, saddr: SockAddr, domain: Domain) -> AddOutcome {
let set_nonblocking =
Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP)).and_then(
|sock| {
sock.set_nonblocking(true)?;
Ok(sock)
},
);
let sock = match set_nonblocking {
Ok(sock) => sock,
Err(e) => {
self.set_error(e);
return AddOutcome::Error;
}
};
match sock.connect(&saddr) {
Ok(()) => match sock.set_nonblocking(false) {
Ok(()) => return AddOutcome::Connected(sock.into()),
Err(e) => self.set_error(e),
},
Err(e) if e.raw_os_error() == Some(libc::EINPROGRESS) => {
let interest = Event::writable(self.attempts.len());
match unsafe { self.poller.add(&sock, interest) } {
Ok(()) => {
self.attempts.push(Some((saddr, sock)));
self.attempts_in_progress += 1;
return AddOutcome::InProgress;
}
Err(e) => self.set_error(e),
}
}
Err(e) => self.set_error(e),
}
AddOutcome::Error
}
fn poll_once(&mut self, timeout: Option<Duration>) -> Result<Option<TcpStream>> {
let mut events = polling::Events::new();
self.poller.wait(&mut events, timeout)?;
for evt in events.iter() {
assert!(evt.writable);
let (sock_addr, sock) = self.attempts[evt.key].take().expect("attempt exists");
self.attempts_in_progress -= 1;
self.poller.delete(&sock).expect("socket is in poll set");
match nix::sys::socket::getsockopt(&sock, nix::sys::socket::sockopt::SocketError) {
Err(err) => self.set_error(err.into()),
Ok(0) => {
if let Some(tcp) = self.socket_into_blocking_tcp_stream(sock) {
return Ok(Some(tcp));
}
return Ok(None);
}
Ok(_) => {}
}
match sock.connect(&sock_addr) {
Err(err) => self.set_error(err),
Ok(()) => {
if let Some(tcp) = self.socket_into_blocking_tcp_stream(sock) {
return Ok(Some(tcp));
}
}
}
}
Ok(None)
}
fn socket_into_blocking_tcp_stream(&mut self, sock: Socket) -> Option<TcpStream> {
match sock.set_nonblocking(false) {
Ok(()) => Some(sock.into()),
Err(err) => {
self.set_error(err);
None
}
}
}
}
impl Drop for HappyEyeballs {
fn drop(&mut self) {
for (_, socket) in self.attempts.drain(..).flatten() {
_ = self.poller.delete(&socket);
}
}
}
enum AddOutcome {
Connected(TcpStream),
InProgress,
Error,
}