Implement async TLS via futures-rustls (BES-27)

rttp_client/Cargo.toml

@@ -40,14 +40,14 @@ async-native-tls = { optional = true, version = "0.5" }
 
 rustls           = { optional = true, version = "0.23" }
 webpki-roots     = { optional = true, version = "0.26" }
-async-rustls     = { optional = true, version = "0.4" }
+futures-rustls   = { optional = true, version = "0.26" }
 
 
 [features]
 default = []
 async = []
 tls-native  = ["native-tls", "async-native-tls"]
-tls-rustls  = ["rustls", "webpki-roots", "async-rustls"]
+tls-rustls  = ["rustls", "webpki-roots", "futures-rustls"]
 
 ##
 # warning: /data/rttp/rttp_client/Cargo.toml: Found `feature = ...` in `target.'cfg(...)'.dependencies`.

rttp_client/src/connection/async_connection.rs

@@ -1,15 +1,23 @@
-use std::net::TcpStream;
+use std::net::{TcpStream, ToSocketAddrs};
 
 use futures::io::{AllowStdIo, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 use socks::{Socks4Stream, Socks5Stream};
-use std::io::{Read, Write};
+use socket2::{Domain, Protocol, Socket, Type};
+use std::io::{self, Read, Write};
+use std::time;
 use url::Url;
 
+#[cfg(feature = "tls-rustls")]
+use std::sync::Arc;
+
 use crate::connection::connection::Connection;
 use crate::error;
 use crate::request::RawRequest;
 use crate::response::Response;
-use crate::types::{Proxy, ProxyType, ToUrl};
+use crate::types::{Proxy, ProxyType};
+
+const HEADER_END: &[u8] = b"\r\n\r\n";
+const CRLF: &[u8] = b"\r\n";
 
 pub struct AsyncConnection<'a> {
   conn: Connection<'a>,
@@ -38,8 +46,43 @@ impl<'a> AsyncConnection<'a> {
 }
 
 impl<'a> AsyncConnection<'a> {
-  async fn async_tcp_stream(&self, addr: &String) -> error::Result<TcpStream> {
-    self.conn.block_tcp_stream(addr)
+  async fn async_tcp_stream(&self, addr: &str) -> error::Result<TcpStream> {
+    let config = self.conn.config();
+    let timeout_read = time::Duration::from_millis(config.read_timeout());
+    let timeout_write = time::Duration::from_millis(config.write_timeout());
+    let mut last_err = None;
+
+    let addrs = addr.to_socket_addrs().map_err(error::request)?;
+    for addr in addrs {
+      let domain = Domain::for_address(addr);
+      let socket = match Socket::new(domain, Type::STREAM, Some(Protocol::TCP)) {
+        Ok(s) => s,
+        Err(e) => {
+          last_err = Some(e);
+          continue;
+        }
+      };
+
+      if let Err(e) = socket.set_read_timeout(Some(timeout_read)) {
+        last_err = Some(e);
+        continue;
+      }
+      if let Err(e) = socket.set_write_timeout(Some(timeout_write)) {
+        last_err = Some(e);
+        continue;
+      }
+
+      if let Err(e) = socket.connect(&addr.into()) {
+        last_err = Some(e);
+        continue;
+      }
+
+      return Ok(TcpStream::from(socket));
+    }
+
+    Err(error::request(
+      last_err.unwrap_or_else(|| io::Error::other("failed to connect")),
+    ))
   }
 
   async fn async_write_stream<S>(&self, stream: &mut S) -> error::Result<()>
@@ -68,12 +111,138 @@ impl<'a> AsyncConnection<'a> {
   where
     S: AsyncRead + Unpin,
   {
-    let mut buffer = Vec::new();
+    let mut binary = async_read_response_header(stream).await?;
+    if is_chunked_encoded(&binary) {
+      binary.extend_from_slice(&async_read_chunked_body(stream).await?);
+    } else {
+      stream
+        .read_to_end(&mut binary)
+        .await
+        .map_err(error::request)?;
+    }
+    Ok(binary)
+  }
+}
+
+async fn async_read_response_header<S>(stream: &mut S) -> error::Result<Vec<u8>>
+where
+  S: AsyncRead + Unpin,
+{
+  let mut header = Vec::new();
+  let mut byte = [0u8; 1];
+
+  loop {
+    let read = stream.read(&mut byte).await.map_err(error::request)?;
+    if read == 0 {
+      if header.is_empty() {
+        return Ok(header);
+      }
+      return Err(error::bad_response("Incomplete http response headers"));
+    }
+
+    header.push(byte[0]);
+    if header.ends_with(HEADER_END) {
+      return Ok(header);
+    }
+  }
+}
+
+fn is_chunked_encoded(header: &[u8]) -> bool {
+  String::from_utf8_lossy(header).lines().any(|line| {
+    let Some((name, value)) = line.split_once(':') else {
+      return false;
+    };
+
+    name.eq_ignore_ascii_case("Transfer-Encoding")
+      && value
+        .split(',')
+        .any(|token| token.trim().eq_ignore_ascii_case("chunked"))
+  })
+}
+
+async fn async_read_chunked_body<S>(stream: &mut S) -> error::Result<Vec<u8>>
+where
+  S: AsyncRead + Unpin,
+{
+  let mut body = Vec::new();
+
+  loop {
+    let line = async_read_crlf_line(stream).await?;
+    let chunk_size = parse_chunk_size(&line)?;
+
+    if chunk_size == 0 {
+      async_consume_trailers(stream).await?;
+      return Ok(body);
+    }
+
+    let current_len = body.len();
+    body.resize(current_len + chunk_size, 0);
     stream
-      .read_to_end(&mut buffer)
+      .read_exact(&mut body[current_len..])
       .await
       .map_err(error::request)?;
-    Ok(buffer)
+    async_consume_crlf(stream).await?;
+  }
+}
+
+async fn async_read_crlf_line<S>(stream: &mut S) -> error::Result<Vec<u8>>
+where
+  S: AsyncRead + Unpin,
+{
+  let mut line = Vec::new();
+  let mut byte = [0u8; 1];
+
+  loop {
+    let read = stream.read(&mut byte).await.map_err(error::request)?;
+    if read == 0 {
+      return Err(error::bad_response("Unexpected end of chunked body"));
+    }
+
+    line.push(byte[0]);
+    if line.ends_with(CRLF) {
+      return Ok(line);
+    }
+  }
+}
+
+fn parse_chunk_size(line: &[u8]) -> error::Result<usize> {
+  let line = std::str::from_utf8(line).map_err(error::response)?;
+  let size = line
+    .trim_end_matches("\r\n")
+    .split(';')
+    .next()
+    .map(str::trim)
+    .filter(|value| !value.is_empty())
+    .ok_or_else(|| error::bad_response("Chunk size line is empty"))?;
+
+  usize::from_str_radix(size, 16).map_err(|_| error::bad_response("Invalid chunk size"))
+}
+
+async fn async_consume_crlf<S>(stream: &mut S) -> error::Result<()>
+where
+  S: AsyncRead + Unpin,
+{
+  let mut suffix = [0u8; 2];
+  stream
+    .read_exact(&mut suffix)
+    .await
+    .map_err(error::request)?;
+  if suffix == *CRLF {
+    Ok(())
+  } else {
+    Err(error::bad_response("Invalid chunk terminator"))
+  }
+}
+
+async fn async_consume_trailers<S>(stream: &mut S) -> error::Result<()>
+where
+  S: AsyncRead + Unpin,
+{
+  loop {
+    let line = async_read_crlf_line(stream).await?;
+    if line == CRLF {
+      return Ok(());
+    }
   }
 }
 
@@ -105,9 +274,15 @@ impl<'a> AsyncConnection<'a> {
     self.async_read_stream(url, stream).await
   }
 
-  async fn async_send_https(&self, url: &Url, mut stream: TcpStream) -> error::Result<Vec<u8>> {
-    #[cfg(any(feature = "tls-native", feature = "tls-rustls"))]
+  async fn async_send_https(&self, url: &Url, stream: TcpStream) -> error::Result<Vec<u8>> {
+    #[cfg(feature = "tls-rustls")]
     {
+      return self.async_send_https_rustls(url, stream).await;
+    }
+
+    #[cfg(all(feature = "tls-native", not(feature = "tls-rustls")))]
+    {
+      let mut stream = stream;
       return self.conn.block_send_https(url, &mut stream);
     }
 
@@ -120,6 +295,55 @@ impl<'a> AsyncConnection<'a> {
       ));
     }
   }
+
+  #[cfg(feature = "tls-rustls")]
+  async fn async_send_https_rustls(
+    &self,
+    url: &Url,
+    stream: TcpStream,
+  ) -> error::Result<Vec<u8>> {
+    use futures_rustls::TlsConnector;
+    use rustls::pki_types::ServerName;
+    use rustls::{ClientConfig, RootCertStore};
+
+    use crate::connection::connection::NoCertificateVerification;
+
+    let config = self.conn.config();
+    let mut root_store = RootCertStore::empty();
+    if config.verify_ssl_cert() {
+      root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
+    }
+
+    let builder = ClientConfig::builder();
+    let rustls_config = if config.verify_ssl_cert() {
+      builder
+        .with_root_certificates(root_store)
+        .with_no_client_auth()
+    } else {
+      builder
+        .dangerous()
+        .with_custom_certificate_verifier(Arc::new(NoCertificateVerification))
+        .with_no_client_auth()
+    };
+
+    let host = self.conn.host(url)?;
+    let server_name: ServerName<'static> = match host.parse::<std::net::IpAddr>() {
+      Ok(ip) => ServerName::IpAddress(ip.into()),
+      Err(_) => ServerName::try_from(host.as_str())
+        .map_err(|_| error::bad_ssl(format!("Invalid server name: {}", host)))?
+        .to_owned(),
+    };
+
+    let connector = TlsConnector::from(Arc::new(rustls_config));
+    let async_tcp = AllowStdIo::new(stream);
+    let mut tls_stream = connector
+      .connect(server_name, async_tcp)
+      .await
+      .map_err(|e| error::bad_ssl(e.to_string()))?;
+
+    self.async_write_stream(&mut tls_stream).await?;
+    self.async_read_stream(url, &mut tls_stream).await
+  }
 }
 
 // proxy connection

rttp_client/src/connection/connection.rs

@@ -26,7 +26,7 @@ use rustls::{
 
 #[cfg(feature = "tls-rustls")]
 #[derive(Debug)]
-struct NoCertificateVerification;
+pub(crate) struct NoCertificateVerification;
 
 #[cfg(feature = "tls-rustls")]
 impl ServerCertVerifier for NoCertificateVerification {