use std::collections::VecDeque; use std::io; use std::net::SocketAddr; use futures::{ future::{ok, FutureResult}, Async, Future, Poll, }; use tokio_tcp::{ConnectFuture, TcpStream}; use trust_dns_resolver::config::{ResolverConfig, ResolverOpts}; use trust_dns_resolver::system_conf::read_system_conf; use super::resolver::{HostAware, Resolver, ResolverError, ResolverFuture}; use super::service::{NewService, Service}; // #[derive(Fail, Debug)] #[derive(Debug)] pub enum ConnectorError { /// Failed to resolve the hostname // #[fail(display = "Failed resolving hostname: {}", _0)] Resolver(ResolverError), /// Not dns records // #[fail(display = "Invalid input: {}", _0)] NoRecords, /// Connection io error // #[fail(display = "{}", _0)] IoError(io::Error), } impl From for ConnectorError { fn from(err: ResolverError) -> Self { ConnectorError::Resolver(err) } } pub struct ConnectionInfo { pub host: String, pub addr: SocketAddr, } pub struct Connector { resolver: Resolver, } impl Default for Connector { fn default() -> Self { let (cfg, opts) = if let Ok((cfg, opts)) = read_system_conf() { (cfg, opts) } else { (ResolverConfig::default(), ResolverOpts::default()) }; Connector::new(cfg, opts) } } impl Connector { pub fn new(cfg: ResolverConfig, opts: ResolverOpts) -> Self { Connector { resolver: Resolver::new(cfg, opts), } } pub fn with_resolver( resolver: Resolver, ) -> impl Service< Request = T, Response = (T, ConnectionInfo, TcpStream), Error = ConnectorError, > + Clone { Connector { resolver } } pub fn new_service() -> impl NewService< Request = T, Response = (T, ConnectionInfo, TcpStream), Error = ConnectorError, InitError = E, > + Clone { || -> FutureResult, E> { ok(Connector::default()) } } pub fn new_service_with_config( cfg: ResolverConfig, opts: ResolverOpts, ) -> impl NewService< Request = T, Response = (T, ConnectionInfo, TcpStream), Error = ConnectorError, InitError = E, > + Clone { move || -> FutureResult, E> { ok(Connector::new(cfg.clone(), opts)) } } pub fn change_request(&self) -> Connector { Connector { resolver: self.resolver.change_request(), } } } impl Clone for Connector { fn clone(&self) -> Self { Connector { resolver: self.resolver.clone(), } } } impl Service for Connector { type Request = T; type Response = (T, ConnectionInfo, TcpStream); type Error = ConnectorError; type Future = ConnectorFuture; fn poll_ready(&mut self) -> Poll<(), Self::Error> { Ok(Async::Ready(())) } fn call(&mut self, req: Self::Request) -> Self::Future { ConnectorFuture { fut: self.resolver.call(req), fut2: None, } } } pub struct ConnectorFuture { fut: ResolverFuture, fut2: Option>, } impl Future for ConnectorFuture { type Item = (T, ConnectionInfo, TcpStream); type Error = ConnectorError; fn poll(&mut self) -> Poll { if let Some(ref mut fut) = self.fut2 { return fut.poll(); } match self.fut.poll().map_err(ConnectorError::from)? { Async::Ready((req, host, addrs)) => { if addrs.is_empty() { Err(ConnectorError::NoRecords) } else { self.fut2 = Some(TcpConnector::new(req, host, addrs)); self.poll() } } Async::NotReady => Ok(Async::NotReady), } } } #[derive(Clone)] pub struct DefaultConnector(Connector); impl Default for DefaultConnector { fn default() -> Self { DefaultConnector(Connector::default()) } } impl DefaultConnector { pub fn new(cfg: ResolverConfig, opts: ResolverOpts) -> Self { DefaultConnector(Connector::new(cfg, opts)) } } impl Service for DefaultConnector { type Request = T; type Response = TcpStream; type Error = ConnectorError; type Future = DefaultConnectorFuture; fn poll_ready(&mut self) -> Poll<(), Self::Error> { self.0.poll_ready() } fn call(&mut self, req: Self::Request) -> Self::Future { DefaultConnectorFuture { fut: self.0.call(req), } } } #[doc(hidden)] pub struct DefaultConnectorFuture { fut: ConnectorFuture, } impl Future for DefaultConnectorFuture { type Item = TcpStream; type Error = ConnectorError; fn poll(&mut self) -> Poll { Ok(Async::Ready(try_ready!(self.fut.poll()).2)) } } /// Tcp stream connector pub struct TcpConnector { req: Option, host: Option, addr: Option, addrs: VecDeque, stream: Option, } impl TcpConnector { pub fn new(req: T, host: String, addrs: VecDeque) -> TcpConnector { TcpConnector { addrs, req: Some(req), host: Some(host), addr: None, stream: None, } } } impl Future for TcpConnector { type Item = (T, ConnectionInfo, TcpStream); type Error = ConnectorError; fn poll(&mut self) -> Poll { // connect loop { if let Some(new) = self.stream.as_mut() { match new.poll() { Ok(Async::Ready(sock)) => { return Ok(Async::Ready(( self.req.take().unwrap(), ConnectionInfo { host: self.host.take().unwrap(), addr: self.addr.take().unwrap(), }, sock, ))) } Ok(Async::NotReady) => return Ok(Async::NotReady), Err(err) => { if self.addrs.is_empty() { return Err(ConnectorError::IoError(err)); } } } } // try to connect let addr = self.addrs.pop_front().unwrap(); self.stream = Some(TcpStream::connect(&addr)); self.addr = Some(addr) } } }