diff options
| -rw-r--r-- | rsh/main.go | 19 |
1 files changed, 18 insertions, 1 deletions
diff --git a/rsh/main.go b/rsh/main.go index dfcea0b..68b6e49 100644 --- a/rsh/main.go +++ b/rsh/main.go @@ -203,7 +203,24 @@ func dial(user, host string, hostKey ssh.HostKeyCallback) (*ssh.Client, error) { Timeout: dialTimeout, } addr := net.JoinHostPort(host, strconv.Itoa(port)) - return ssh.Dial("tcp", addr, cfg) + conn, err := net.DialTimeout("tcp", addr, dialTimeout) + if err != nil { + return nil, err + } + if err := conn.SetDeadline(time.Now().Add(dialTimeout)); err != nil { + conn.Close() + return nil, err + } + clientConn, chans, reqs, err := ssh.NewClientConn(conn, addr, cfg) + if err != nil { + conn.Close() + return nil, err + } + if err := conn.SetDeadline(time.Time{}); err != nil { + clientConn.Close() + return nil, err + } + return ssh.NewClient(clientConn, chans, reqs), nil } // parseTransport extracts the user, host, and remote command from the argument |