From e4662b77f1e042a55422164f5b008efc6a233347 Mon Sep 17 00:00:00 2001
From: Yousong Zhou <zhouyousong@yunionyun.com>
Date: Fri, 2 Apr 2021 15:33:19 +0800
Subject: [PATCH] util: ssh: add ConnectContext method for ClientConfig

---
 pkg/util/ssh/ssh.go | 33 ++++++++++++++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/pkg/util/ssh/ssh.go b/pkg/util/ssh/ssh.go
index 438478030d..a5ccd01848 100644
--- a/pkg/util/ssh/ssh.go
+++ b/pkg/util/ssh/ssh.go
@@ -16,8 +16,10 @@ package ssh
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
+	"net"
 	"os"
 	"strings"
 	"time"
@@ -29,6 +31,12 @@ import (
 	"yunion.io/x/pkg/errors"
 )
 
+const (
+	ErrBadConfig = errors.Error("bad config")
+	ErrNetwork   = errors.Error("network error")
+	ErrProtocol  = errors.Error("ssh protocol error")
+)
+
 type ClientConfig struct {
 	Username   string
 	Password   string
@@ -54,7 +62,7 @@ func (conf ClientConfig) ToSshConfig() (*ssh.ClientConfig, error) {
 	if conf.PrivateKey != "" {
 		signer, err := parsePrivateKey(conf.PrivateKey)
 		if err != nil {
-			return nil, err
+			return nil, errors.Wrapf(ErrBadConfig, "parse private key: %v", err)
 		}
 		auths = append(auths, ssh.PublicKeys(signer))
 	}
@@ -75,6 +83,29 @@ func (conf ClientConfig) Connect() (*ssh.Client, error) {
 	return client, nil
 }
 
+func (conf ClientConfig) ConnectContext(ctx context.Context) (*ssh.Client, error) {
+	cliConfig, err := conf.ToSshConfig()
+	if err != nil {
+		return nil, err
+	}
+
+	addr := fmt.Sprintf("%s:%d", conf.Host, conf.Port)
+	d := &net.Dialer{}
+	netconn, err := d.DialContext(ctx, "tcp", addr)
+	if err != nil {
+		return nil, errors.Wrapf(ErrNetwork, "tcp dial: %v", err)
+	}
+
+	sshconn, chans, reqs, err := ssh.NewClientConn(netconn, addr, cliConfig)
+	if err != nil {
+		netconn.Close()
+		return nil, errors.Wrap(ErrProtocol, err.Error())
+	}
+
+	sshc := ssh.NewClient(sshconn, chans, reqs)
+	return sshc, nil
+}
+
 type Client struct {
 	config ClientConfig
 	client *ssh.Client