From a769bbb7139d3f0bd8f04d636c135ee2f7657454 Mon Sep 17 00:00:00 2001
From: KS Chan <mrkschan@gmail.com>
Date: Wed, 1 Mar 2017 00:33:17 +0800
Subject: [PATCH] Add support for timeout

---
 main.go | 17 +++++++++++++----
 ssh.go  | 21 +++++++++++++++++----
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/main.go b/main.go
index 6772906..3a93d2a 100644
--- a/main.go
+++ b/main.go
@@ -140,17 +140,26 @@ func runSSH(c command, host string, section *SSHConfigFileSection, agt agent.Age
 
 	// Try using each available AuthMethod to establish SSH session
 	var (
+		config  *sshClientConfig
 		session *sshSession
 		err     error
 	)
 	for k, m := range methods {
-		config := newSSHClientConfig(user, host, section, agt, m)
+		config, err = newSSHClientConfig(user, host, section, agt, m, options)
+		if err != nil {
+			log.Debugf("Failed to parse SSH config - %v", err)
+			continue
+		}
+
 		session, err = config.NewSession(options)
-		if err == nil {
-			break // Session established, quit trying the next AuthMethod
+		if err != nil {
+			log.Debugf("Failed to establish session with %v - %v", k, err)
+			continue
 		}
 
-		log.Debugf("Failed to establish session with %v - %v", k, err)
+		if session != nil {
+			break // Session established, quit trying the next AuthMethod
+		}
 	}
 
 	if session == nil {
diff --git a/ssh.go b/ssh.go
index 50a347f..e807726 100644
--- a/ssh.go
+++ b/ssh.go
@@ -10,7 +10,9 @@ import (
 	"os"
 	"os/user"
 	"path/filepath"
+	"strconv"
 	"syscall"
+	"time"
 
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
@@ -64,16 +66,27 @@ func updateFromSSHConfigFile(section *SSHConfigFileSection, host, user *string,
 }
 
 // newSSHClientConfig initializes per-host SSH configuration.
-func newSSHClientConfig(user, host string, section *SSHConfigFileSection, agt agent.Agent, method ssh.AuthMethod) *sshClientConfig {
+func newSSHClientConfig(user, host string, section *SSHConfigFileSection, agt agent.Agent, method ssh.AuthMethod, options map[string]string) (*sshClientConfig, error) {
+	var timeout time.Duration
+
+	if ct, ok := options["ConnectTimeout"]; ok {
+		t, err := strconv.Atoi(ct)
+		if err != nil {
+			return nil, fmt.Errorf("ConnectTimeout is not an integer")
+		}
+		timeout = time.Second * time.Duration(t)
+	}
+
 	config := &ssh.ClientConfig{
-		User: user,
-		Auth: []ssh.AuthMethod{method},
+		User:    user,
+		Auth:    []ssh.AuthMethod{method},
+		Timeout: timeout,
 	}
 	return &sshClientConfig{
 		agent:        agt,
 		host:         host,
 		ClientConfig: config,
-	}
+	}, nil
 }
 
 // NewSession creates a new ssh session with the host.