fix(platform): validate socks5 proxy (#2212)

tkestack · Dec 15, 2022 · 5289a85 · 5289a85
1 parent 966b72e
commit 5289a85
Showing 1 changed file with 18 additions and 7 deletions.
diff --git a/pkg/util/ssh/socks5_proxy.go b/pkg/util/ssh/socks5_proxy.go
@@ -27,12 +27,15 @@ import (
 	"tkestack.io/tke/pkg/util/log"
 )
 
+const defaultCheckTargetAddr = "ccr.ccs.tencentyun.com:443"
+
 type SOCKS5 struct {
-	Host        string
-	Port        int
-	User        string
-	Password    string
-	DialTimeOut time.Duration
+	Host            string
+	Port            int
+	User            string
+	Password        string
+	DialTimeOut     time.Duration
+	CheckTargetAddr string
 }
 
 func (sk SOCKS5) ProxyConn(targetAddr string) (net.Conn, func(), error) {
@@ -71,6 +74,14 @@ func (sk SOCKS5) ProxyConn(targetAddr string) (net.Conn, func(), error) {
 }
 
 func (sk SOCKS5) CheckTunnel() error {
-	_, err := proxy.SOCKS5("tcp", net.JoinHostPort(sk.Host, fmt.Sprintf("%d", sk.Port)), nil, proxy.Direct)
-	return err
+	targetAddr := sk.CheckTargetAddr
+	if len(targetAddr) == 0 {
+		targetAddr = defaultCheckTargetAddr
+	}
+	_, closer, err := sk.ProxyConn(targetAddr)
+	if err != nil {
+		return fmt.Errorf("tunnel is unavailable: %v", err)
+	}
+	defer closer()
+	return nil
 }