Simplify the ssh.Conn parameter passing

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
  }