Simplify the ssh.Conn parameter passing

Author: iawia002

pkg/proxy/http_proxy.go

@@ -5,9 +5,10 @@ import (
 	"fmt"
 	"net/http"
 
+	"golang.org/x/crypto/ssh"
 	"k8s.io/klog"
+
 	"kubesphere.io/tower/pkg/agent"
-	"kubesphere.io/tower/pkg/utils"
 )
 
 type HTTPProxy struct {
@@ -22,8 +23,8 @@ type HTTPProxy struct {
 	kubesphereAPIServerProxy *Server
 }
 
-func NewHTTPProxy(ssh utils.GetSSHConn, kubernetesPort uint16, kubespherePort uint16, config *agent.Config, ca, serverCert, serverKey []byte) (*HTTPProxy, *http.Transport, *http.Transport, error) {
-	k8stransPort, useBearerToken, servertlsConfig, err := buildServerData(ssh, config.KubernetesSvcHost, config.CAData, config.CertData, config.KeyData, ca, serverCert, serverKey)
+func NewHTTPProxy(sshConn ssh.Conn, kubernetesPort uint16, kubespherePort uint16, config *agent.Config, ca, serverCert, serverKey []byte) (*HTTPProxy, *http.Transport, *http.Transport, error) {
+	k8stransPort, useBearerToken, servertlsConfig, err := buildServerData(sshConn, config.KubernetesSvcHost, config.CAData, config.CertData, config.KeyData, ca, serverCert, serverKey)
 	if err != nil {
 		return nil, nil, nil, err
 	}
@@ -33,7 +34,7 @@ func NewHTTPProxy(ssh utils.GetSSHConn, kubernetesPort uint16, kubespherePort ui
 		return nil, nil, nil, err
 	}
 
-	kstransPort, useBearerToken, _, err := buildServerData(ssh, config.KubeSphereSvcHost, nil, nil, nil, nil, nil, nil)
+	kstransPort, useBearerToken, _, err := buildServerData(sshConn, config.KubeSphereSvcHost, nil, nil, nil, nil, nil, nil)
 	if err != nil {
 		return nil, nil, nil, err
 	}

pkg/proxy/proxy.go

@@ -216,7 +216,7 @@ func (s *Proxy) handleWebsocket(w http.ResponseWriter, req *http.Request) {
 	// if the agent has connected the server with the same cluster name, we don't need to create HttpProxy anymore
 	// we only create two new httpTransport objects, then put them into the server's httpClient set.
 	if proxy, ok = s.sessions[c.Name]; !ok {
-		proxy, k8sTransport, ksTransport, err = NewHTTPProxy(func() ssh.Conn { return sshConn }, client.Spec.Connection.KubernetesAPIServerPort, client.Spec.Connection.KubeSphereAPIServerPort, c, s.caCert, cert, key)
+		proxy, k8sTransport, ksTransport, err = NewHTTPProxy(sshConn, client.Spec.Connection.KubernetesAPIServerPort, client.Spec.Connection.KubeSphereAPIServerPort, c, s.caCert, cert, key)
 		if err != nil {
 			failed(err)
 			return
@@ -229,13 +229,13 @@ func (s *Proxy) handleWebsocket(w http.ResponseWriter, req *http.Request) {
 
 		s.sessions[c.Name] = proxy
 	} else {
-		k8sTransport, _, _, err = buildServerData(func() ssh.Conn { return sshConn }, c.KubernetesSvcHost, c.CAData, c.CertData, c.KeyData, s.caCert, cert, key)
+		k8sTransport, _, _, err = buildServerData(sshConn, c.KubernetesSvcHost, c.CAData, c.CertData, c.KeyData, s.caCert, cert, key)
 		if err != nil {
 			failed(err)
 			return
 		}
 
-		ksTransport, _, _, err = buildServerData(func() ssh.Conn { return sshConn }, c.KubeSphereSvcHost, c.CAData, c.CertData, c.KeyData, s.caCert, cert, key)
+		ksTransport, _, _, err = buildServerData(sshConn, c.KubeSphereSvcHost, c.CAData, c.CertData, c.KeyData, s.caCert, cert, key)
 		if err != nil {
 			failed(err)
 			return

pkg/proxy/proxy_server.go

@@ -11,6 +11,7 @@ import (
 	"sync"
 	"time"
 
+	"golang.org/x/crypto/ssh"
 	utilnet "k8s.io/apimachinery/pkg/util/net"
 	k8sproxy "k8s.io/apimachinery/pkg/util/proxy"
 	"k8s.io/klog"
@@ -68,15 +69,11 @@ func newProxyServer(name, host, scheme string, port uint16, useBearerToken bool,
 }
 
 // buildServerData returns http.Transport and tlsConfig, which are necessary for creating proxy server.
-func buildServerData(sshConn utils.GetSSHConn, host string, ca, cert, key, serverCa, serverCert, serverKey []byte) (*http.Transport, bool, *tls.Config, error) {
+func buildServerData(sshConn ssh.Conn, host string, ca, cert, key, serverCa, serverCert, serverKey []byte) (*http.Transport, bool, *tls.Config, error) {
 	useBearerToken := true
 
 	transport := &http.Transport{
 		DialContext: func(ctx context.Context, network, addr string) (conn net.Conn, err error) {
-			c := sshConn()
-			if c == nil {
-				return nil, fmt.Errorf("no remote connetion available")
-			}
 			return utils.NewSshConn(sshConn, host)
 		},
 	}

pkg/utils/ssh_conn.go

@@ -2,30 +2,25 @@ package utils
 
 import (
 	"errors"
-	"golang.org/x/crypto/ssh"
 	"io"
 	"net"
 	"time"
+
+	"golang.org/x/crypto/ssh"
 )
 
 var ErrorInvalidConnection = errors.New("invalid connection")
 
-//ErrorNoAvailableConn means there haven't available shh connection.
-var ErrorNoAvailableConn = errors.New("no available ssh connection")
-
-type GetSSHConn func() ssh.Conn
-
 type SshConn struct {
 	dst io.ReadWriteCloser
 }
 
-func NewSshConn(conn GetSSHConn, remote string) (net.Conn, error) {
-	c := conn()
-	if c == nil {
-		return nil, ErrorNoAvailableConn
+func NewSshConn(conn ssh.Conn, remote string) (net.Conn, error) {
+	if conn == nil {
+		return nil, errors.New("the ssh connection is nil")
 	}
 
-	dst, reqs, err := c.OpenChannel("kubesphere", []byte(remote))
+	dst, reqs, err := conn.OpenChannel("kubesphere", []byte(remote))
 	if err != nil {
 		return nil, err
 	}