Merge pull request #91 from XuGuangnian/main

add connect_with_timeout method
1148118271 · Dec 19, 2023 · fbb651f · fbb651f
2 parents be05260 + 21ff11b
commit fbb651f
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/examples/timeout/src/main.rs b/examples/timeout/src/main.rs
@@ -21,7 +21,7 @@ fn main() {
         .password("password")
         .private_key_path("./id_rsa")
         .timeout(Some(Duration::from_secs(5)))
-        .connect("127.0.0.1:7777")
+        .connect_with_timeout("127.0.0.1:7777", Some(Duration::from_secs(2)))
     {
         Err(e) => println!("Got error {}", e),
         _ => unreachable!(),

diff --git a/src/session/mod.rs b/src/session/mod.rs
@@ -257,6 +257,26 @@ impl SessionBuilder {
         self.connect_bio(tcp)
     }
 
+    pub fn connect_with_timeout<A>(
+        self,
+        addr: A,
+        timeout: Option<Duration>,
+    ) -> SshResult<SessionConnector<TcpStream>>
+    where
+        A: ToSocketAddrs,
+    {
+        // connect tcp with custom connection timeout
+        let tcp = if let Some(ref to) = timeout {
+            TcpStream::connect_timeout(&addr.to_socket_addrs()?.next().unwrap(), *to)?
+        } else {
+            TcpStream::connect(addr)?
+        };
+
+        // default nonblocking
+        tcp.set_nonblocking(true).unwrap();
+        self.connect_bio(tcp)
+    }
+
     /// connect to target server w/ a bio object
     ///
     /// which requires to implement `std::io::{Read, Write}`