use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; use actix_rt::net::TcpStream; use actix_service::{Service, ServiceFactory}; use either::Either; use futures::future::{ok, Ready}; use trust_dns_resolver::AsyncResolver; use crate::connect::{Address, Connect, Connection}; use crate::connector::{TcpConnector, TcpConnectorFactory}; use crate::error::ConnectError; use crate::resolver::{Resolver, ResolverFactory}; pub struct ConnectServiceFactory { tcp: TcpConnectorFactory, resolver: ResolverFactory, } impl ConnectServiceFactory { /// Construct new ConnectService factory pub fn new() -> Self { ConnectServiceFactory { tcp: TcpConnectorFactory::default(), resolver: ResolverFactory::default(), } } /// Construct new connect service with custom dns resolver pub fn with_resolver(resolver: AsyncResolver) -> Self { ConnectServiceFactory { tcp: TcpConnectorFactory::default(), resolver: ResolverFactory::new(resolver), } } /// Construct new service pub fn service(&self) -> ConnectService { ConnectService { tcp: self.tcp.service(), resolver: self.resolver.service(), } } /// Construct new tcp stream service pub fn tcp_service(&self) -> TcpConnectService { TcpConnectService { tcp: self.tcp.service(), resolver: self.resolver.service(), } } } impl Default for ConnectServiceFactory { fn default() -> Self { ConnectServiceFactory { tcp: TcpConnectorFactory::default(), resolver: ResolverFactory::default(), } } } impl Clone for ConnectServiceFactory { fn clone(&self) -> Self { ConnectServiceFactory { tcp: self.tcp.clone(), resolver: self.resolver.clone(), } } } impl ServiceFactory for ConnectServiceFactory { type Request = Connect; type Response = Connection; type Error = ConnectError; type Config = (); type Service = ConnectService; type InitError = (); type Future = Ready>; fn new_service(&self, _: &()) -> Self::Future { ok(self.service()) } } #[derive(Clone)] pub struct ConnectService { tcp: TcpConnector, resolver: Resolver, } impl Service for ConnectService { type Request = Connect; type Response = Connection; type Error = ConnectError; type Future = ConnectServiceResponse; fn poll_ready(&mut self, _: &mut Context) -> Poll> { Poll::Ready(Ok(())) } fn call(&mut self, req: Connect) -> Self::Future { ConnectServiceResponse { state: ConnectState::Resolve(self.resolver.call(req)), tcp: self.tcp.clone(), } } } enum ConnectState { Resolve( as Service>::Future), Connect( as Service>::Future), } impl ConnectState { fn poll( &mut self, cx: &mut Context, ) -> Either, ConnectError>>, Connect> { match self { ConnectState::Resolve(ref mut fut) => match Pin::new(fut).poll(cx) { Poll::Pending => Either::Left(Poll::Pending), Poll::Ready(Ok(res)) => Either::Right(res), Poll::Ready(Err(err)) => Either::Left(Poll::Ready(Err(err))), }, ConnectState::Connect(ref mut fut) => Either::Left(Pin::new(fut).poll(cx)), } } } pub struct ConnectServiceResponse { state: ConnectState, tcp: TcpConnector, } impl Future for ConnectServiceResponse { type Output = Result, ConnectError>; fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll { let res = match self.state.poll(cx) { Either::Right(res) => { self.state = ConnectState::Connect(self.tcp.call(res)); self.state.poll(cx) } Either::Left(res) => return res, }; match res { Either::Left(res) => res, Either::Right(_) => panic!(), } } } #[derive(Clone)] pub struct TcpConnectService { tcp: TcpConnector, resolver: Resolver, } impl Service for TcpConnectService { type Request = Connect; type Response = TcpStream; type Error = ConnectError; type Future = TcpConnectServiceResponse; fn poll_ready(&mut self, _: &mut Context) -> Poll> { Poll::Ready(Ok(())) } fn call(&mut self, req: Connect) -> Self::Future { TcpConnectServiceResponse { state: TcpConnectState::Resolve(self.resolver.call(req)), tcp: self.tcp.clone(), } } } enum TcpConnectState { Resolve( as Service>::Future), Connect( as Service>::Future), } impl TcpConnectState { fn poll( &mut self, cx: &mut Context, ) -> Either>, Connect> { match self { TcpConnectState::Resolve(ref mut fut) => match Pin::new(fut).poll(cx) { Poll::Pending => (), Poll::Ready(Ok(res)) => return Either::Right(res), Poll::Ready(Err(err)) => return Either::Left(Poll::Ready(Err(err))), }, TcpConnectState::Connect(ref mut fut) => { if let Poll::Ready(res) = Pin::new(fut).poll(cx) { return match res { Ok(conn) => Either::Left(Poll::Ready(Ok(conn.into_parts().0))), Err(err) => Either::Left(Poll::Ready(Err(err))), }; } } } Either::Left(Poll::Pending) } } pub struct TcpConnectServiceResponse { state: TcpConnectState, tcp: TcpConnector, } impl Future for TcpConnectServiceResponse { type Output = Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll { let res = match self.state.poll(cx) { Either::Right(res) => { self.state = TcpConnectState::Connect(self.tcp.call(res)); self.state.poll(cx) } Either::Left(res) => return res, }; match res { Either::Left(res) => res, Either::Right(_) => panic!(), } } }