use std::io; use std::net::SocketAddr; use std::collections::VecDeque; use std::time::Duration; use actix::Arbiter; use trust_dns_resolver::ResolverFuture; use trust_dns_resolver::config::{ResolverConfig, ResolverOpts}; use trust_dns_resolver::lookup_ip::LookupIpFuture; use futures::{Async, Future, Poll}; use tokio_core::reactor::Timeout; use tokio_core::net::{TcpStream, TcpStreamNew}; #[derive(Fail, Debug)] pub enum TcpConnectorError { /// Failed to resolve the hostname #[fail(display = "Failed resolving hostname: {}", _0)] Dns(String), /// Address is invalid #[fail(display = "Invalid input: {}", _0)] InvalidInput(&'static str), /// Connecting took too long #[fail(display = "Timeout out while establishing connection")] Timeout, /// Connection io error #[fail(display = "{}", _0)] IoError(io::Error), } pub struct TcpConnector { lookup: Option, port: u16, ips: VecDeque, error: Option, timeout: Timeout, stream: Option, } impl TcpConnector { pub fn new>(addr: S, port: u16, timeout: Duration) -> TcpConnector { println!("TES: {:?} {:?}", addr.as_ref(), port); // try to parse as a regular SocketAddr first if let Ok(addr) = addr.as_ref().parse() { let mut ips = VecDeque::new(); ips.push_back(addr); TcpConnector { lookup: None, port: port, ips: ips, error: None, stream: None, timeout: Timeout::new(timeout, Arbiter::handle()).unwrap() } } else { // we need to do dns resolution let resolve = match ResolverFuture::from_system_conf(Arbiter::handle()) { Ok(resolve) => resolve, Err(err) => { warn!("Can not create system dns resolver: {}", err); ResolverFuture::new( ResolverConfig::default(), ResolverOpts::default(), Arbiter::handle()) } }; TcpConnector { lookup: Some(resolve.lookup_ip(addr.as_ref())), port: port, ips: VecDeque::new(), error: None, stream: None, timeout: Timeout::new(timeout, Arbiter::handle()).unwrap() } } } } impl Future for TcpConnector { type Item = TcpStream; type Error = TcpConnectorError; fn poll(&mut self) -> Poll { if let Some(err) = self.error.take() { Err(err) } else { // timeout if let Ok(Async::Ready(_)) = self.timeout.poll() { return Err(TcpConnectorError::Timeout) } // lookip ips if let Some(mut lookup) = self.lookup.take() { match lookup.poll() { Ok(Async::NotReady) => { self.lookup = Some(lookup); return Ok(Async::NotReady) }, Ok(Async::Ready(ips)) => { let port = self.port; let ips = ips.iter().map(|ip| SocketAddr::new(ip, port)); self.ips.extend(ips); if self.ips.is_empty() { return Err(TcpConnectorError::Dns( "Expect at least one A dns record".to_owned())) } }, Err(err) => return Err(TcpConnectorError::Dns(format!("{}", err))), } } // connect loop { if let Some(mut new) = self.stream.take() { match new.poll() { Ok(Async::Ready(sock)) => return Ok(Async::Ready(sock)), Ok(Async::NotReady) => { self.stream = Some(new); return Ok(Async::NotReady) }, Err(err) => { if self.ips.is_empty() { return Err(TcpConnectorError::IoError(err)) } } } } // try to connect let addr = self.ips.pop_front().unwrap(); self.stream = Some(TcpStream::connect(&addr, Arbiter::handle())); } } } }