From 999ada2e8444c7890d8e44211eec13fb5fd6e3fd Mon Sep 17 00:00:00 2001
From: txthinking <cloud@txthinking.com>
Date: Tue, 22 Nov 2022 21:16:58 +0800
Subject: [PATCH] 60s wait for command

---
 cli/hancock/main.go |  2 +-
 conn.go             | 39 +++++++++++++++++++++++++++++++++++++++
 instance.go         | 10 ++++++++--
 3 files changed, 48 insertions(+), 3 deletions(-)
 create mode 100644 conn.go

diff --git a/cli/hancock/main.go b/cli/hancock/main.go
index 122c0d8..c64744c 100644
--- a/cli/hancock/main.go
+++ b/cli/hancock/main.go
@@ -78,7 +78,7 @@ func main() {
 	}
 	app := cli.NewApp()
 	app.Name = "hancock"
-	app.Version = "20221106"
+	app.Version = "20221122"
 	app.Usage = "Deploy and run command on remote instances with built-in nami, joker"
 	app.Authors = []*cli.Author{
 		{
diff --git a/conn.go b/conn.go
new file mode 100644
index 0000000..01c021d
--- /dev/null
+++ b/conn.go
@@ -0,0 +1,39 @@
+// Copyright (c) 2020-present Cloud <cloud@txthinking.com>
+//
+// This program is free software; you can redistribute it and/or
+// modify it under the terms of version 3 of the GNU General Public
+// License as published by the Free Software Foundation.
+//
+// This program is distributed in the hope that it will be useful, but
+// WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+// General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+package hancock
+
+import (
+	"net"
+	"time"
+)
+
+type Conn struct {
+	net.Conn
+	Timeout int
+}
+
+func (c *Conn) Read(b []byte) (int, error) {
+	if err := c.Conn.SetDeadline(time.Now().Add(time.Duration(c.Timeout) * time.Second)); err != nil {
+		return 0, err
+	}
+	return c.Conn.Read(b)
+}
+
+func (c *Conn) Write(b []byte) (int, error) {
+	if err := c.Conn.SetDeadline(time.Now().Add(time.Duration(c.Timeout) * time.Second)); err != nil {
+		return 0, err
+	}
+	return c.Conn.Write(b)
+}
diff --git a/instance.go b/instance.go
index 36c357e..1820c96 100644
--- a/instance.go
+++ b/instance.go
@@ -16,9 +16,11 @@ package hancock
 
 import (
 	"io"
+	"net"
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
 
 	"github.com/pkg/sftp"
 	"golang.org/x/crypto/ssh"
@@ -47,16 +49,20 @@ func NewInstance(server, user, password string, privateKey []byte) (*Instance, e
 		l = append(l, ssh.PublicKeys(signer))
 	}
 
+	conn, err := net.DialTimeout("tcp", server, 60*time.Second)
+	if err != nil {
+		return nil, err
+	}
 	config := &ssh.ClientConfig{
 		User:            user,
 		Auth:            l,
 		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
-		// Timeout:         10 * time.Second,
 	}
-	client, err := ssh.Dial("tcp", server, config)
+	c, chans, reqs, err := ssh.NewClientConn(&Conn{Conn: conn, Timeout: 60}, server, config)
 	if err != nil {
 		return nil, err
 	}
+	client := ssh.NewClient(c, chans, reqs)
 	return &Instance{
 		Client: client,
 	}, nil