From 30301a2dd2539e1a1948eb60ba8a5a3ed78d6676 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BD=98=E6=9D=B0?= <panjie@huanle.com>
Date: Tue, 24 Jun 2025 14:52:35 +0800
Subject: [PATCH] Add ssh proxy feature

---
 assets/etcdkeeper/index.html |  2 +-
 src/etcdkeeper/go.mod        |  3 +-
 src/etcdkeeper/go.sum        |  3 ++
 src/etcdkeeper/main.go       | 59 ++++++++++++++++++++++++++++++++++--
 4 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/assets/etcdkeeper/index.html b/assets/etcdkeeper/index.html
index ed03bc5..2d4491b 100644
--- a/assets/etcdkeeper/index.html
+++ b/assets/etcdkeeper/index.html
@@ -146,7 +146,7 @@ <h2><img src="../etcdkeeper/framework/logo.png" width="20" style="position:relat
 			$('#elayout').height(($(window).height() - 128) + 'px')
 		}
 
-		var timeout = 5000 // milliseconds
+		var timeout = 10000 // milliseconds
 		var separator = '';
 		var serverBase = '../v3';
 		var etcdBase = Cookies.get('etcd-endpoint');
diff --git a/src/etcdkeeper/go.mod b/src/etcdkeeper/go.mod
index 1bf6ea4..7107043 100644
--- a/src/etcdkeeper/go.mod
+++ b/src/etcdkeeper/go.mod
@@ -6,6 +6,8 @@ require (
 	github.com/coreos/etcd v3.3.15+incompatible
 	go.etcd.io/etcd/client/v2 v2.305.15
 	go.etcd.io/etcd/client/v3 v3.5.8
+	golang.org/x/crypto v0.21.0
+	google.golang.org/grpc v1.59.0
 )
 
 require (
@@ -28,6 +30,5 @@ require (
 	google.golang.org/genproto v0.0.0-20230822172742-b8732ec3820d // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20230822172742-b8732ec3820d // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d // indirect
-	google.golang.org/grpc v1.59.0 // indirect
 	google.golang.org/protobuf v1.33.0 // indirect
 )
diff --git a/src/etcdkeeper/go.sum b/src/etcdkeeper/go.sum
index b8e93dd..4a1cec3 100644
--- a/src/etcdkeeper/go.sum
+++ b/src/etcdkeeper/go.sum
@@ -52,6 +52,8 @@ go.uber.org/zap v1.17.0/go.mod h1:MXVU+bhUf/A7Xi2HNOnopQOrmycQ5Ih87HtOu4q5SSo=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
+golang.org/x/crypto v0.21.0 h1:X31++rzVUdKhX5sWmSOFZxx8UW/ldWx55cbf08iNAMA=
+golang.org/x/crypto v0.21.0/go.mod h1:0BP7YvVV9gBbVKyeTG0Gyn+gZm94bibOW5BjDEYAOMs=
 golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
@@ -68,6 +70,7 @@ golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=
 golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/term v0.18.0 h1:FcHjZXDMxI8mM3nwhX9HlKop4C0YQvCVCdwYl2wOtE8=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
diff --git a/src/etcdkeeper/main.go b/src/etcdkeeper/main.go
index a64e534..e90b934 100644
--- a/src/etcdkeeper/main.go
+++ b/src/etcdkeeper/main.go
@@ -8,11 +8,9 @@ import (
 	_ "etcdkeeper/session/providers/memory"
 	"flag"
 	"fmt"
-	"github.com/coreos/etcd/pkg/transport"
-	"go.etcd.io/etcd/client/v2"
-	"go.etcd.io/etcd/client/v3"
 	"io"
 	"log"
+	"net"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -21,6 +19,13 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/coreos/etcd/pkg/transport"
+	"go.etcd.io/etcd/client/v2"
+	clientv3 "go.etcd.io/etcd/client/v3"
+	"google.golang.org/grpc"
+
+	"golang.org/x/crypto/ssh"
 )
 
 var (
@@ -34,6 +39,7 @@ var (
 	useAuth        = flag.Bool("auth", false, "use auth")
 	connectTimeout = flag.Int("timeout", 5, "ETCD client connect timeout")
 	sendMsgSize    = flag.Int("sendMsgSize", 2*1024*1024, "ETCD client max send msg size")
+	sshProxy       = flag.String("ssh-proxy", "", "SSH proxy in username:password@host:port format")
 	rootUsers      = make(map[string]*userInfo) // host:rootUser
 	rootUsersV2    = make(map[string]*userInfo) // host:rootUser
 
@@ -1006,6 +1012,22 @@ func newClient(uinfo *userInfo) (*clientv3.Client, error) {
 		MaxCallSendMsgSize: *sendMsgSize,
 		//DialOptions:        []grpc.DialOption{grpc.WithBlock()},
 	}
+
+	if *sshProxy != "" {
+		// log.Println("Using SSH proxy:", *sshProxy)
+
+		client, err := connectSSHProxy(*sshProxy)
+		if err != nil {
+			panic(err)
+		}
+		dialer := func(ctx context.Context, address string) (net.Conn, error) {
+			return client.Dial("tcp", address)
+		}
+		conf.DialOptions = []grpc.DialOption{
+			grpc.WithContextDialer(dialer),
+		}
+	}
+
 	if *useAuth {
 		conf.Username = uinfo.uname
 		conf.Password = uinfo.passwd
@@ -1119,3 +1141,34 @@ func getInfo(host string) map[string]string {
 func size(num int, unit int) (n, rem int) {
 	return num / unit, num - (num/unit)*unit
 }
+
+func connectSSHProxy(proxy string) (*ssh.Client, error) {
+	// parse format：username:password@host:port
+	atIdx := strings.LastIndex(proxy, "@")
+	if atIdx == -1 {
+		return nil, fmt.Errorf("invalid ssh-proxy format, missing '@'")
+	}
+	userPass := proxy[:atIdx]
+	hostPort := proxy[atIdx+1:]
+
+	colonIdx := strings.Index(userPass, ":")
+	if colonIdx == -1 {
+		return nil, fmt.Errorf("invalid ssh-proxy format, missing ':' in user:pass")
+	}
+	user := userPass[:colonIdx]
+	password := userPass[colonIdx+1:]
+
+	config := &ssh.ClientConfig{
+		User: user,
+		Auth: []ssh.AuthMethod{
+			ssh.Password(password),
+		},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+		Timeout:         time.Second * time.Duration(*connectTimeout),
+	}
+	client, err := ssh.Dial("tcp", hostPort, config)
+	if err != nil {
+		return nil, err
+	}
+	return client, nil
+}