From fd259fd40352f04fc5e26def91a94e3b2273a61b Mon Sep 17 00:00:00 2001
From: Thomas Boerger <thomas@webhippie.de>
Date: Tue, 12 Jul 2016 13:31:38 +0200
Subject: [PATCH] Added a flag for connection timeouts

---
 DOCS.md   |  2 ++
 main.go   |  6 ++++++
 plugin.go | 16 +++++++++-------
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/DOCS.md b/DOCS.md
index c139f2c..f7ac44d 100644
--- a/DOCS.md
+++ b/DOCS.md
@@ -10,6 +10,7 @@ The following parameters are used to configure the plugin:
 * **user** - user to log in as on the remote machine
 * **key** - private SSH key for the remote machine
 * **sleep** - sleep for seconds between host connections
+* **timeout** - timeout for the tcp connection attempt
 * **commands** - list of commands to execute
 
 The following secret values can be set to configure the plugin.
@@ -19,6 +20,7 @@ The following secret values can be set to configure the plugin.
 * **SSH_USER** - corresponds to **user**
 * **SSH_KEY** - corresponds to **key**
 * **SSH_SLEEP** - corresponds to **sleep**
+* **SSH_TIMEOUT** - corresponds to **timeout**
 
 It is highly recommended to put the **SSH_KEY** into a secret so it is not
 exposed to users. This can be done using the drone-cli.
diff --git a/main.go b/main.go
index 92e760e..b83e2e9 100644
--- a/main.go
+++ b/main.go
@@ -44,6 +44,11 @@ func main() {
 			Usage:  "sleep between hosts",
 			EnvVar: "PLUGIN_SLEEP,SSH_SLEEP",
 		},
+		cli.DurationFlag{
+			Name:   "timeout",
+			Usage:  "connection timeout",
+			EnvVar: "PLUGIN_TIMEOUT,SSH_TIMEOUT",
+		},
 		cli.StringSliceFlag{
 			Name:   "commands",
 			Usage:  "execute commands",
@@ -72,6 +77,7 @@ func run(c *cli.Context) error {
 			Host:     c.StringSlice("host"),
 			Port:     c.Int("port"),
 			Sleep:    c.Int("sleep"),
+			Timeout:  c.Duration("timeout"),
 			Commands: c.StringSlice("commands"),
 		},
 	}
diff --git a/plugin.go b/plugin.go
index 3b8f13d..135c22f 100644
--- a/plugin.go
+++ b/plugin.go
@@ -13,12 +13,13 @@ import (
 
 type (
 	Config struct {
-		Key      string   `json:"key"`
-		User     string   `json:"user"`
-		Host     []string `json:"host"`
-		Port     int      `json:"port"`
-		Sleep    int      `json:"sleep"`
-		Commands []string `json:"commands"`
+		Key      string        `json:"key"`
+		User     string        `json:"user"`
+		Host     []string      `json:"host"`
+		Port     int           `json:"port"`
+		Sleep    int           `json:"sleep"`
+		Timeout  time.Duration `json:"timeout"`
+		Commands []string      `json:"commands"`
 	}
 
 	Plugin struct {
@@ -44,7 +45,8 @@ func (p Plugin) Exec() error {
 		}
 
 		config := &ssh.ClientConfig{
-			User: p.Config.User,
+			Timeout: p.Config.Timeout,
+			User:    p.Config.User,
 			Auth: []ssh.AuthMethod{
 				ssh.PublicKeys(signer),
 			},