actix-tls release candidate prep (#422)

actix-tls/src/connect/tcp.rs

@@ -1,43 +1,202 @@
-use actix_rt::net::TcpStream;
+//! TCP connector service.
+//!
+//! See [`TcpConnector`] for main connector service factory docs.
+
+use std::{
+    collections::VecDeque,
+    future::Future,
+    io,
+    net::{IpAddr, SocketAddr, SocketAddrV4, SocketAddrV6},
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use actix_rt::net::{TcpSocket, TcpStream};
 use actix_service::{Service, ServiceFactory};
+use actix_utils::future::{ok, Ready};
+use futures_core::ready;
+use log::{error, trace};
+use tokio_util::sync::ReusableBoxFuture;
 
-use super::{Address, Connect, ConnectError, ConnectServiceFactory, Connection, Resolver};
+use super::{connect_addrs::ConnectAddrs, error::ConnectError, ConnectInfo, Connection, Host};
 
-/// Create TCP connector service.
-pub fn new_connector<T: Address + 'static>(
-    resolver: Resolver,
-) -> impl Service<Connect<T>, Response = Connection<T, TcpStream>, Error = ConnectError> + Clone
-{
-    ConnectServiceFactory::new(resolver).service()
+/// TCP connector service factory.
+#[derive(Debug, Copy, Clone)]
+pub struct TcpConnector;
+
+impl TcpConnector {
+    /// Returns a new TCP connector service.
+    pub fn service(&self) -> TcpConnectorService {
+        TcpConnectorService
+    }
 }
 
-/// Create TCP connector service factory.
-pub fn new_connector_factory<T: Address + 'static>(
-    resolver: Resolver,
-) -> impl ServiceFactory<
-    Connect<T>,
-    Config = (),
-    Response = Connection<T, TcpStream>,
-    Error = ConnectError,
-    InitError = (),
-> + Clone {
-    ConnectServiceFactory::new(resolver)
+impl<R: Host> ServiceFactory<ConnectInfo<R>> for TcpConnector {
+    type Response = Connection<R, TcpStream>;
+    type Error = ConnectError;
+    type Config = ();
+    type Service = TcpConnectorService;
+    type InitError = ();
+    type Future = Ready<Result<Self::Service, Self::InitError>>;
+
+    fn new_service(&self, _: ()) -> Self::Future {
+        ok(self.service())
+    }
 }
 
-/// Create TCP connector service with default parameters.
-pub fn default_connector<T: Address + 'static>(
-) -> impl Service<Connect<T>, Response = Connection<T, TcpStream>, Error = ConnectError> + Clone
-{
-    new_connector(Resolver::Default)
+/// TCP connector service.
+#[derive(Debug, Copy, Clone)]
+pub struct TcpConnectorService;
+
+impl<R: Host> Service<ConnectInfo<R>> for TcpConnectorService {
+    type Response = Connection<R, TcpStream>;
+    type Error = ConnectError;
+    type Future = TcpConnectorFut<R>;
+
+    actix_service::always_ready!();
+
+    fn call(&self, req: ConnectInfo<R>) -> Self::Future {
+        let port = req.port();
+
+        let ConnectInfo {
+            request: req,
+            addr,
+            local_addr,
+            ..
+        } = req;
+
+        TcpConnectorFut::new(req, port, local_addr, addr)
+    }
 }
 
-/// Create TCP connector service factory with default parameters.
-pub fn default_connector_factory<T: Address + 'static>() -> impl ServiceFactory<
-    Connect<T>,
-    Config = (),
-    Response = Connection<T, TcpStream>,
-    Error = ConnectError,
-    InitError = (),
-> + Clone {
-    new_connector_factory(Resolver::Default)
+/// Connect future for TCP service.
+#[doc(hidden)]
+pub enum TcpConnectorFut<R> {
+    Response {
+        req: Option<R>,
+        port: u16,
+        local_addr: Option<IpAddr>,
+        addrs: Option<VecDeque<SocketAddr>>,
+        stream: ReusableBoxFuture<Result<TcpStream, io::Error>>,
+    },
+
+    Error(Option<ConnectError>),
+}
+
+impl<R: Host> TcpConnectorFut<R> {
+    pub(crate) fn new(
+        req: R,
+        port: u16,
+        local_addr: Option<IpAddr>,
+        addr: ConnectAddrs,
+    ) -> TcpConnectorFut<R> {
+        if addr.is_unresolved() {
+            error!("TCP connector: unresolved connection address");
+            return TcpConnectorFut::Error(Some(ConnectError::Unresolved));
+        }
+
+        trace!(
+            "TCP connector: connecting to {} on port {}",
+            req.hostname(),
+            port
+        );
+
+        match addr {
+            ConnectAddrs::None => unreachable!("none variant already checked"),
+
+            ConnectAddrs::One(addr) => TcpConnectorFut::Response {
+                req: Some(req),
+                port,
+                local_addr,
+                addrs: None,
+                stream: ReusableBoxFuture::new(connect(addr, local_addr)),
+            },
+
+            // when resolver returns multiple socket addr for request they would be popped from
+            // front end of queue and returns with the first successful tcp connection.
+            ConnectAddrs::Multi(mut addrs) => {
+                let addr = addrs.pop_front().unwrap();
+
+                TcpConnectorFut::Response {
+                    req: Some(req),
+                    port,
+                    local_addr,
+                    addrs: Some(addrs),
+                    stream: ReusableBoxFuture::new(connect(addr, local_addr)),
+                }
+            }
+        }
+    }
+}
+
+impl<R: Host> Future for TcpConnectorFut<R> {
+    type Output = Result<Connection<R, TcpStream>, ConnectError>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self.get_mut() {
+            TcpConnectorFut::Error(err) => Poll::Ready(Err(err.take().unwrap())),
+
+            TcpConnectorFut::Response {
+                req,
+                port,
+                local_addr,
+                addrs,
+                stream,
+            } => loop {
+                match ready!(stream.poll(cx)) {
+                    Ok(sock) => {
+                        let req = req.take().unwrap();
+
+                        trace!(
+                            "TCP connector: successfully connected to {:?} - {:?}",
+                            req.hostname(),
+                            sock.peer_addr()
+                        );
+
+                        return Poll::Ready(Ok(Connection::new(req, sock)));
+                    }
+
+                    Err(err) => {
+                        trace!(
+                            "TCP connector: failed to connect to {:?} port: {}",
+                            req.as_ref().unwrap().hostname(),
+                            port,
+                        );
+
+                        if let Some(addr) = addrs.as_mut().and_then(|addrs| addrs.pop_front()) {
+                            stream.set(connect(addr, *local_addr));
+                        } else {
+                            return Poll::Ready(Err(ConnectError::Io(err)));
+                        }
+                    }
+                }
+            },
+        }
+    }
+}
+
+async fn connect(addr: SocketAddr, local_addr: Option<IpAddr>) -> io::Result<TcpStream> {
+    // use local addr if connect asks for it
+    match local_addr {
+        Some(ip_addr) => {
+            let socket = match ip_addr {
+                IpAddr::V4(ip_addr) => {
+                    let socket = TcpSocket::new_v4()?;
+                    let addr = SocketAddr::V4(SocketAddrV4::new(ip_addr, 0));
+                    socket.bind(addr)?;
+                    socket
+                }
+                IpAddr::V6(ip_addr) => {
+                    let socket = TcpSocket::new_v6()?;
+                    let addr = SocketAddr::V6(SocketAddrV6::new(ip_addr, 0, 0, 0));
+                    socket.bind(addr)?;
+                    socket
+                }
+            };
+
+            socket.connect(addr).await
+        }
+
+        None => TcpStream::connect(addr).await,
+    }
 }