Merge pull request #26 from AvivNaaman/25-support-receive-timeout

#25: support timeout in worker recv, use in tests/cli.

Author: afiffon

smb-cli/src/path.rs

@@ -19,6 +19,8 @@ impl UncPath {
         cli: &Cli,
     ) -> Result<(Connection, Session, Tree, Option<Resource>), Box<dyn Error>> {
         let mut smb = Connection::new();
+        smb.set_timeout(Some(std::time::Duration::from_secs(10)))
+            .await?;
         smb.connect(format!("{}:{}", self.server, cli.port).as_str())
             .await?;
         let mut session = smb

smb/src/connection.rs

@@ -25,22 +25,35 @@ use netbios_client::NetBiosClient;
 use std::cmp::max;
 use std::sync::atomic::{AtomicU16, AtomicU64};
 use std::sync::Arc;
+use std::time::Duration;
 pub use transformer::TransformError;
 use worker::{Worker, WorkerImpl};
 
 pub struct Connection {
     handler: HandlerReference<ConnectionMessageHandler>,
+    timeout: Option<std::time::Duration>,
 }
 
 impl Connection {
     pub fn new() -> Connection {
         Connection {
             handler: HandlerReference::new(ConnectionMessageHandler::new()),
+            timeout: None,
         }
     }
+
+    #[maybe_async]
+    pub async fn set_timeout(&mut self, timeout: Option<Duration>) -> crate::Result<()> {
+        self.timeout = timeout;
+        if let Some(worker) = self.handler.worker.get() {
+            worker.set_timeout(timeout).await?;
+        }
+        Ok(())
+    }
+
     #[maybe_async]
     pub async fn connect(&mut self, address: &str) -> crate::Result<()> {
-        let mut netbios_client = NetBiosClient::new();
+        let mut netbios_client = NetBiosClient::new(self.timeout);
 
         log::debug!("Connecting to {}...", address);
         netbios_client.connect(address).await?;
@@ -105,7 +118,7 @@ impl Connection {
             }
         }
 
-        Ok(WorkerImpl::start(netbios_client).await?)
+        Ok(WorkerImpl::start(netbios_client, self.timeout).await?)
     }
 
     #[maybe_async]

smb/src/connection/netbios_client.rs

@@ -1,15 +1,16 @@
 use maybe_async::*;
-use std::io::Cursor;
+use std::{io::Cursor, time::Duration};
 
 #[cfg(feature = "sync")]
 use std::{
     io::{self, Read, Write},
-    net::TcpStream,
+    net::{TcpStream, ToSocketAddrs},
 };
 #[cfg(feature = "async")]
 use tokio::{
     io::{self, AsyncReadExt, AsyncWriteExt},
     net::{tcp, TcpStream},
+    select,
 };
 
 use binrw::prelude::*;
@@ -37,26 +38,57 @@ type TcpWrite = TcpStream;
 pub struct NetBiosClient {
     reader: Option<TcpRead>,
     writer: Option<TcpWrite>,
+    timeout: Option<Duration>,
 }
 
 impl NetBiosClient {
-    pub fn new() -> NetBiosClient {
+    pub fn new(timeout: Option<Duration>) -> NetBiosClient {
         NetBiosClient {
             reader: None,
             writer: None,
+            timeout,
         }
     }
 
     /// Connects to a NetBios server in the specified address.
     #[maybe_async]
     pub async fn connect(&mut self, address: &str) -> crate::Result<()> {
-        let socket = TcpStream::connect(address).await?;
+        let socket = self.connect_timeout(address).await?;
         let (r, w) = Self::split_socket(socket);
         self.reader = Some(r);
         self.writer = Some(w);
         Ok(())
     }
 
+    #[cfg(feature = "sync")]
+    fn connect_timeout(&mut self, address: &str) -> crate::Result<TcpStream> {
+        if let Some(t) = self.timeout {
+            log::debug!("Connecting to {} with timeout {:?}.", address, t);
+            // convert to SocketAddr:
+            let address = address
+                .to_socket_addrs()?
+                .next()
+                .ok_or(crate::Error::InvalidAddress(address.to_string()))?;
+            TcpStream::connect_timeout(&address, t).map_err(Into::into)
+        } else {
+            log::debug!("Connecting to {}.", address);
+            TcpStream::connect(&address).map_err(Into::into)
+        }
+    }
+
+    #[cfg(feature = "async")]
+    async fn connect_timeout(&mut self, address: &str) -> crate::Result<TcpStream> {
+        if let None = self.timeout {
+            log::debug!("Connecting to {}.", address);
+            return TcpStream::connect(&address).await.map_err(Into::into);
+        }
+
+        select! {
+            res = TcpStream::connect(&address) => res.map_err(Into::into),
+            _ = tokio::time::sleep(self.timeout.unwrap()) => Err(crate::Error::OperationTimeout("Tcp connect".to_string(), self.timeout.unwrap())),
+        }
+    }
+
     pub fn is_connected(&self) -> bool {
         self.reader.is_some()
     }
@@ -197,10 +229,12 @@ impl NetBiosClient {
             NetBiosClient {
                 reader: self.reader,
                 writer: None,
+                timeout: self.timeout,
             },
             NetBiosClient {
                 reader: None,
                 writer: self.writer,
+                timeout: self.timeout,
             },
         ))
     }

smb/src/connection/worker/multi_worker/async_backend.rs

@@ -1,5 +1,6 @@
 use crate::sync_helpers::*;
 use std::sync::Arc;
+use std::time::Duration;
 use tokio::{select, sync::oneshot};
 
 use crate::{msg_handler::IncomingMessage, packets::netbios::NetBiosTcpMessage, Error};
@@ -190,10 +191,25 @@ impl MultiWorkerBackend for AsyncBackend {
         oneshot::channel()
     }
 
-    async fn wait_on_waiter(waiter: Self::AwaitingWaiter) -> crate::Result<IncomingMessage> {
-        waiter
-            .await
-            .map_err(|_| Error::MessageProcessingError("Failed to receive message.".to_string()))?
+    async fn wait_on_waiter(
+        waiter: Self::AwaitingWaiter,
+        timeout: Option<Duration>,
+    ) -> crate::Result<IncomingMessage> {
+        match timeout {
+            Some(timeout) => {
+                tokio::select! {
+                    msg = waiter => {
+                        msg.map_err(|_| Error::MessageProcessingError("Failed to receive message.".to_string()))?
+                    },
+                    _ = tokio::time::sleep(timeout) => {
+                        Err(Error::OperationTimeout("Waiting for message receive.".to_string(), timeout))
+                    }
+                }
+            }
+            None => waiter.await.map_err(|_| {
+                Error::MessageProcessingError("Failed to receive message.".to_string())
+            })?,
+        }
     }
 
     fn send_notify(

smb/src/connection/worker/multi_worker/backend_trait.rs

@@ -1,7 +1,7 @@
 use crate::connection::netbios_client::NetBiosClient;
 use crate::sync_helpers::*;
 use maybe_async::*;
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
 use crate::{msg_handler::IncomingMessage, packets::netbios::NetBiosTcpMessage};
 
@@ -32,7 +32,10 @@ pub trait MultiWorkerBackend {
         mpsc::Receiver<Self::SendMessage>,
     );
 
-    async fn wait_on_waiter(waiter: Self::AwaitingWaiter) -> crate::Result<IncomingMessage>;
+    async fn wait_on_waiter(
+        waiter: Self::AwaitingWaiter,
+        timeout: Option<Duration>,
+    ) -> crate::Result<IncomingMessage>;
     fn send_notify(
         tx: Self::AwaitingNotifier,
         msg: crate::Result<IncomingMessage>,

smb/src/connection/worker/multi_worker/base.rs

@@ -4,6 +4,7 @@ use crate::connection::worker::Worker;
 use crate::sync_helpers::*;
 use maybe_async::*;
 use std::sync::atomic::AtomicBool;
+use std::time::Duration;
 use std::{collections::HashMap, sync::Arc};
 
 use crate::{
@@ -34,6 +35,9 @@ where
     /// A channel to send messages to the worker.
     pub(crate) sender: mpsc::Sender<T::SendMessage>,
     stopped: AtomicBool,
+
+    /// atomic duration:
+    timeout: RwLock<Option<Duration>>,
 }
 
 /// Holds state for the worker, regarding messages to be received.
@@ -140,7 +144,10 @@ where
     T::AwaitingNotifier: std::fmt::Debug,
 {
     #[maybe_async]
-    async fn start(netbios_client: NetBiosClient) -> crate::Result<Arc<Self>> {
+    async fn start(
+        netbios_client: NetBiosClient,
+        timeout: Option<Duration>,
+    ) -> crate::Result<Arc<Self>> {
         // Build the worker
         let (tx, rx) = T::make_send_channel_pair();
         let worker = Arc::new(MultiWorkerBase::<T> {
@@ -149,6 +156,7 @@ where
             transformer: Transformer::default(),
             sender: tx,
             stopped: AtomicBool::new(false),
+            timeout: RwLock::new(timeout),
         });
 
         worker
@@ -235,7 +243,8 @@ where
             rx
         };
 
-        let wait_result = T::wait_on_waiter(wait_for_receive).await;
+        let timeout = { *self.timeout.read().await? };
+        let wait_result = T::wait_on_waiter(wait_for_receive, timeout).await;
 
         // Wait for the message to be received.
         Ok(wait_result.map_err(|_| {
@@ -246,6 +255,12 @@ where
     fn transformer(&self) -> &Transformer {
         &self.transformer
     }
+
+    #[maybe_async]
+    async fn set_timeout(&self, timeout: Option<Duration>) -> crate::Result<()> {
+        *self.timeout.write().await? = timeout;
+        Ok(())
+    }
 }
 
 impl<T> std::fmt::Debug for MultiWorkerBase<T>

smb/src/connection/worker/multi_worker/threading_backend.rs

@@ -171,10 +171,22 @@ impl MultiWorkerBackend for ThreadingBackend {
         std::sync::mpsc::channel()
     }
 
-    fn wait_on_waiter(waiter: Self::AwaitingWaiter) -> crate::Result<IncomingMessage> {
-        waiter
-            .recv()
-            .map_err(|_| Error::MessageProcessingError("Failed to receive message.".to_string()))?
+    fn wait_on_waiter(
+        waiter: Self::AwaitingWaiter,
+        timeout: Option<Duration>,
+    ) -> crate::Result<IncomingMessage> {
+        if let None = timeout {
+            return waiter.recv().map_err(|_| {
+                Error::MessageProcessingError("Failed to receive message.".to_string())
+            })?;
+        }
+        waiter.recv_timeout(timeout.unwrap()).map_err(|e| match e {
+            std::sync::mpsc::RecvTimeoutError::Timeout => Error::OperationTimeout(
+                "Waiting for message receive.".to_string(),
+                timeout.unwrap(),
+            ),
+            _ => Error::MessageProcessingError("Failed to receive message.".to_string()),
+        })?
     }
 
     fn send_notify(

smb/src/connection/worker/single_worker.rs

@@ -1,4 +1,4 @@
-use std::{cell::RefCell, sync::Arc};
+use std::{cell::RefCell, sync::Arc, time::Duration};
 
 use crate::{
     connection::{netbios_client::NetBiosClient, transformer::Transformer},
@@ -18,10 +18,11 @@ pub struct SingleWorker {
 }
 
 impl Worker for SingleWorker {
-    fn start(netbios_client: NetBiosClient) -> crate::Result<Arc<Self>> {
+    fn start(netbios_client: NetBiosClient, timeout: Option<Duration>) -> crate::Result<Arc<Self>> {
         if !netbios_client.is_connected() {
             Err(crate::Error::NotConnected)
         } else {
+            netbios_client.set_read_timeout(timeout)?;
             Ok(Arc::new(Self {
                 netbios_client: RefCell::new(netbios_client),
                 transformer: Transformer::default(),
@@ -66,4 +67,8 @@ impl Worker for SingleWorker {
     fn transformer(&self) -> &Transformer {
         &self.transformer
     }
+
+    fn set_timeout(&self, timeout: Option<Duration>) -> crate::Result<()> {
+        self.netbios_client.borrow_mut().set_read_timeout(timeout)
+    }
 }

smb/src/connection/worker/worker_trait.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
 use crate::sync_helpers::*;
 
@@ -21,10 +21,16 @@ use crate::{
 #[allow(async_fn_in_trait)]
 pub trait Worker: Sized + std::fmt::Debug {
     /// Instantiates a new connection worker.
-    async fn start(netbios_client: NetBiosClient) -> crate::Result<Arc<Self>>;
+    async fn start(
+        netbios_client: NetBiosClient,
+        timeout: Option<Duration>,
+    ) -> crate::Result<Arc<Self>>;
     /// Stops the worker, shutting down the connection.
     async fn stop(&self) -> crate::Result<()>;
 
+    /// Sets the timeout for the worker.
+    async fn set_timeout(&self, timeout: Option<Duration>) -> crate::Result<()>;
+
     async fn send(self: &Self, msg: OutgoingMessage) -> crate::Result<SendMessageResult>;
     /// Receive a message from the server.
     /// This is a user function that will wait for the message to be received.

smb/src/error.rs

@@ -52,6 +52,8 @@ pub enum Error {
     UsernameError(String),
     #[error("Message processing failed. {0}")]
     MessageProcessingError(String),
+    #[error("Operation timed out: {0}, took >{1:?}")]
+    OperationTimeout(String, std::time::Duration),
     #[error("Lock error.")]
     LockError,
     #[cfg(feature = "async")]
@@ -69,6 +71,8 @@ pub enum Error {
     UnexpectedMessageId(u64, u64),
     #[error("Expected info of type {0} but got {1}")]
     UnexpectedInformationType(u8, u8),
+    #[error("Invalid address {0}")]
+    InvalidAddress(String),
 }
 
 impl<T> From<PoisonError<T>> for Error {