From 9330025d07f1c595db8c3d8bae266b87b265e038 Mon Sep 17 00:00:00 2001
From: ChristianVisintin <christian.visintin1997@gmail.com>
Date: Sat, 16 Jan 2021 10:58:07 +0100
Subject: [PATCH] Connection timeout for SFTP/SCP clients

---
 CHANGELOG.md                      |  3 +++
 src/filetransfer/scp_transfer.rs  | 36 +++++++++++++++++++++++++------
 src/filetransfer/sftp_transfer.rs | 34 +++++++++++++++++++++++------
 3 files changed, 60 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index eb11ce5..2a20922 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,9 @@
 
 FIXME: Released on ???
 
+- Enhancements:
+  - Added connection timeout to 30 seconds to SFTP/SCP clients and improved name lookup system.
+
 ## 0.3.0
 
  Released on 10/01/2021
diff --git a/src/filetransfer/scp_transfer.rs b/src/filetransfer/scp_transfer.rs
index 5cd2351..5b7ffda 100644
--- a/src/filetransfer/scp_transfer.rs
+++ b/src/filetransfer/scp_transfer.rs
@@ -37,9 +37,9 @@ use crate::utils::parser::parse_lstime;
 use regex::Regex;
 use ssh2::{Channel, Session};
 use std::io::{BufReader, BufWriter, Read, Write};
-use std::net::TcpStream;
+use std::net::{SocketAddr, TcpStream, ToSocketAddrs};
 use std::path::{Path, PathBuf};
-use std::time::SystemTime;
+use std::time::{Duration, SystemTime};
 
 /// ## ScpFileTransfer
 ///
@@ -311,12 +311,34 @@ impl FileTransfer for ScpFileTransfer {
         password: Option<String>,
     ) -> Result<Option<String>, FileTransferError> {
         // Setup tcp stream
-        let tcp: TcpStream = match TcpStream::connect(format!("{}:{}", address, port)) {
-            Ok(stream) => stream,
-            Err(err) => {
+        let socket_addresses: Vec<SocketAddr> =
+            match format!("{}:{}", address, port).to_socket_addrs() {
+                Ok(s) => s.collect(),
+                Err(err) => {
+                    return Err(FileTransferError::new_ex(
+                        FileTransferErrorType::BadAddress,
+                        format!("{}", err),
+                    ))
+                }
+            };
+        let mut tcp: Option<TcpStream> = None;
+        // Try addresses
+        for socket_addr in socket_addresses.iter() {
+            match TcpStream::connect_timeout(&socket_addr, Duration::from_secs(30)) {
+                Ok(stream) => {
+                    tcp = Some(stream);
+                    break;
+                }
+                Err(_) => continue,
+            }
+        }
+        // If stream is None, return connection timeout
+        let tcp: TcpStream = match tcp {
+            Some(t) => t,
+            None => {
                 return Err(FileTransferError::new_ex(
-                    FileTransferErrorType::BadAddress,
-                    format!("{}", err),
+                    FileTransferErrorType::ConnectionError,
+                    String::from("Connection timeout"),
                 ))
             }
         };
diff --git a/src/filetransfer/sftp_transfer.rs b/src/filetransfer/sftp_transfer.rs
index 9e2622b..38e1edb 100644
--- a/src/filetransfer/sftp_transfer.rs
+++ b/src/filetransfer/sftp_transfer.rs
@@ -34,7 +34,7 @@ use crate::system::sshkey_storage::SshKeyStorage;
 // Includes
 use ssh2::{FileStat, OpenFlags, OpenType, Session, Sftp};
 use std::io::{BufReader, BufWriter, Read, Write};
-use std::net::TcpStream;
+use std::net::{SocketAddr, TcpStream, ToSocketAddrs};
 use std::path::{Path, PathBuf};
 use std::time::{Duration, SystemTime};
 
@@ -203,12 +203,34 @@ impl FileTransfer for SftpFileTransfer {
         password: Option<String>,
     ) -> Result<Option<String>, FileTransferError> {
         // Setup tcp stream
-        let tcp: TcpStream = match TcpStream::connect(format!("{}:{}", address, port)) {
-            Ok(stream) => stream,
-            Err(err) => {
+        let socket_addresses: Vec<SocketAddr> =
+            match format!("{}:{}", address, port).to_socket_addrs() {
+                Ok(s) => s.collect(),
+                Err(err) => {
+                    return Err(FileTransferError::new_ex(
+                        FileTransferErrorType::BadAddress,
+                        format!("{}", err),
+                    ))
+                }
+            };
+        let mut tcp: Option<TcpStream> = None;
+        // Try addresses
+        for socket_addr in socket_addresses.iter() {
+            match TcpStream::connect_timeout(&socket_addr, Duration::from_secs(30)) {
+                Ok(stream) => {
+                    tcp = Some(stream);
+                    break;
+                }
+                Err(_) => continue,
+            }
+        }
+        // If stream is None, return connection timeout
+        let tcp: TcpStream = match tcp {
+            Some(t) => t,
+            None => {
                 return Err(FileTransferError::new_ex(
-                    FileTransferErrorType::BadAddress,
-                    format!("{}", err),
+                    FileTransferErrorType::ConnectionError,
+                    String::from("Connection timeout"),
                 ))
             }
         };