Wrap timeouts around network operations

src/client.rs

@@ -5,6 +5,7 @@
 //! ```no_run
 //! use std::collections::HashMap;
 //! use std::net::ToSocketAddrs;
+//! use std::time::Duration;
 //!
 //! use epp_client::EppClient;
 //! use epp_client::domain::DomainCheck;
@@ -16,7 +17,8 @@
 //! // Create an instance of EppClient
 //! let host = "example.com";
 //! let addr = (host, 7000).to_socket_addrs().unwrap().next().unwrap();
-//! let mut client = match EppClient::connect("registry_name".to_string(), addr, host, None).await {
+//! let timeout = Duration::from_secs(5);
+//! let mut client = match EppClient::connect("registry_name".to_string(), addr, host, None, timeout).await {
 //!     Ok(client) => client,
 //!     Err(e) => panic!("Failed to create EppClient: {}",  e)
 //! };
@@ -37,6 +39,7 @@ use std::convert::TryInto;
 use std::io;
 use std::net::SocketAddr;
 use std::sync::Arc;
+use std::time::Duration;
 
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::TcpStream;
@@ -49,7 +52,7 @@ use tokio_rustls::TlsConnector;
 use tracing::info;
 
 use crate::common::{Certificate, NoExtension, PrivateKey};
-use crate::connection::EppConnection;
+use crate::connection::{self, EppConnection};
 use crate::error::Error;
 use crate::hello::{Greeting, GreetingDocument, HelloDocument};
 use crate::request::{Command, Extension, Transaction};
@@ -69,7 +72,8 @@ impl EppClient<TlsStream<TcpStream>> {
     ///
     /// The `registry` is used as a name in internal logging; `addr` provides the address to
     /// connect to, `hostname` is sent as the TLS server name indication and `identity` provides
-    /// optional TLS client authentication. Uses rustls as the TLS implementation.
+    /// optional TLS client authentication (using) rustls as the TLS implementation.
+    /// The `timeout` limits the time spent on any underlying network operations.
     ///
     /// Alternatively, use `EppClient::new()` with any established `AsyncRead + AsyncWrite + Unpin`
     /// implementation.
@@ -78,6 +82,7 @@ impl EppClient<TlsStream<TcpStream>> {
         addr: SocketAddr,
         hostname: &str,
         identity: Option<(Vec<Certificate>, PrivateKey)>,
+        timeout: Duration,
     ) -> Result<Self, Error> {
         info!("Connecting to server: {:?}", addr);
 
@@ -115,17 +120,17 @@ impl EppClient<TlsStream<TcpStream>> {
         })?;
 
         let connector = TlsConnector::from(Arc::new(config));
-        let tcp = TcpStream::connect(&addr).await?;
+        let future = connector.connect(domain, TcpStream::connect(&addr).await?);
-        let stream = connector.connect(domain, tcp).await?;
+        let stream = connection::timeout(timeout, future).await?;
-        Self::new(registry, stream).await
+        Self::new(registry, stream, timeout).await
     }
 }
 
 impl<IO: AsyncRead + AsyncWrite + Unpin> EppClient<IO> {
     /// Create an `EppClient` from an already established connection
-    pub async fn new(registry: String, stream: IO) -> Result<Self, Error> {
+    pub async fn new(registry: String, stream: IO, timeout: Duration) -> Result<Self, Error> {
         Ok(Self {
-            connection: EppConnection::new(registry, stream).await?,
+            connection: EppConnection::new(registry, stream, timeout).await?,
         })
     }
 

src/connection.rs

@@ -1,6 +1,8 @@
 //! Manages registry connections and reading/writing to them
 
 use std::convert::TryInto;
+use std::future::Future;
+use std::time::Duration;
 use std::{io, str, u32};
 
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
@@ -13,14 +15,20 @@ pub(crate) struct EppConnection<IO> {
     registry: String,
     stream: IO,
     pub greeting: String,
+    timeout: Duration,
 }
 
 impl<IO: AsyncRead + AsyncWrite + Unpin> EppConnection<IO> {
-    pub(crate) async fn new(registry: String, stream: IO) -> Result<Self, Error> {
+    pub(crate) async fn new(
+        registry: String,
+        stream: IO,
+        timeout: Duration,
+    ) -> Result<Self, Error> {
         let mut this = Self {
             registry,
             stream,
             greeting: String::new(),
+            timeout,
         };
 
         this.greeting = this.get_epp_response().await?;
@@ -40,7 +48,7 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> EppConnection<IO> {
         buf[..4].clone_from_slice(&len_u32);
         buf[4..].clone_from_slice(content.as_bytes());
 
-        let wrote = self.stream.write(&buf).await?;
+        let wrote = timeout(self.timeout, self.stream.write(&buf)).await?;
         debug!("{}: Wrote {} bytes", self.registry, wrote);
         Ok(())
     }
@@ -48,7 +56,7 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> EppConnection<IO> {
     /// Receives response from the socket and converts it into an EPP XML string
     async fn get_epp_response(&mut self) -> Result<String, Error> {
         let mut buf = [0u8; 4];
-        self.stream.read_exact(&mut buf).await?;
+        timeout(self.timeout, self.stream.read_exact(&mut buf)).await?;
 
         let buf_size: usize = u32::from_be_bytes(buf).try_into()?;
 
@@ -59,7 +67,7 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> EppConnection<IO> {
         let mut read_size: usize = 0;
 
         loop {
-            let read = self.stream.read(&mut buf[read_size..]).await?;
+            let read = timeout(self.timeout, self.stream.read(&mut buf[read_size..])).await?;
             debug!("{}: Read: {} bytes", self.registry, read);
 
             read_size += read;
@@ -95,7 +103,18 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> EppConnection<IO> {
     pub(crate) async fn shutdown(&mut self) -> Result<(), Error> {
         info!("{}: Closing connection", self.registry);
 
-        self.stream.shutdown().await?;
+        timeout(self.timeout, self.stream.shutdown()).await?;
         Ok(())
     }
 }
+
+pub(crate) async fn timeout<T, E: Into<Error>>(
+    timeout: Duration,
+    fut: impl Future<Output = Result<T, E>>,
+) -> Result<T, Error> {
+    match tokio::time::timeout(timeout, fut).await {
+        Ok(Ok(t)) => Ok(t),
+        Ok(Err(e)) => Err(e.into()),
+        Err(_) => Err(Error::Timeout),
+    }
+}

src/error.rs

@@ -12,8 +12,9 @@ use crate::response::ResponseStatus;
 /// Error enum holding the possible error types
 #[derive(Debug)]
 pub enum Error {
-    Io(std::io::Error),
     Command(ResponseStatus),
+    Io(std::io::Error),
+    Timeout,
     Xml(Box<dyn StdError + Send + Sync>),
     Other(Box<dyn StdError + Send + Sync>),
 }
@@ -23,10 +24,11 @@ impl StdError for Error {}
 impl Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
-            Error::Io(e) => write!(f, "I/O error: {}", e),
             Error::Command(e) => {
                 write!(f, "command error: {}", e.result.message)
             }
+            Error::Io(e) => write!(f, "I/O error: {}", e),
+            Error::Timeout => write!(f, "timeout"),
             Error::Xml(e) => write!(f, "(de)serialization error: {}", e),
             Error::Other(e) => write!(f, "error: {}", e),
         }

src/lib.rs

@@ -41,6 +41,7 @@
 //! ```no_run
 //! use std::collections::HashMap;
 //! use std::net::ToSocketAddrs;
+//! use std::time::Duration;
 //!
 //! use epp_client::EppClient;
 //! use epp_client::domain::DomainCheck;
@@ -54,7 +55,8 @@
 //! // Create an instance of EppClient
 //! let host = "example.com";
 //! let addr = (host, 7000).to_socket_addrs().unwrap().next().unwrap();
-//! let mut client = match EppClient::connect("registry_name".to_string(), addr, host, None).await {
+//! let timeout = Duration::from_secs(5);
+//! let mut client = match EppClient::connect("registry_name".to_string(), addr, host, None, timeout).await {
 //!     Ok(client) => client,
 //!     Err(e) => panic!("Failed to create EppClient: {}",  e)
 //! };

tests/basic.rs

@@ -1,6 +1,7 @@
 use std::fs::File;
 use std::io::{self, Read, Write};
 use std::str;
+use std::time::Duration;
 
 use regex::Regex;
 use tokio_test::io::Builder;
@@ -86,7 +87,10 @@ async fn client() {
     ])
     .build();
 
-    let mut client = EppClient::new("test".into(), stream).await.unwrap();
+    let mut client = EppClient::new("test".into(), stream, Duration::from_secs(5))
+        .await
+        .unwrap();
+
     assert_eq!(client.xml_greeting(), xml("response/greeting.xml"));
     client
         .transact(