fix: ssh portmap redo ssh dial

Author: wencaiwulue

pkg/daemon/handler/ssh.go

@@ -82,7 +82,7 @@ func (w *wsHandler) handle(ctx context.Context) {
 	if err != nil {
 		return
 	}
-	err = util.PortMapUntil(ctx, cli, remote, local)
+	err = util.PortMapUntil(ctx, w.sshConfig, remote, local)
 	if err != nil {
 		w.Log("Port map error: %v", err)
 		return

pkg/handler/connect.go

@@ -893,7 +893,7 @@ func SshJump(ctx context.Context, conf *util.SshConfig, flags *pflag.FlagSet, pr
 	if print {
 		log.Infof("wait jump to bastion host...")
 	}
-	err = util.PortMapUntil(ctx, cli, remote, local)
+	err = util.PortMapUntil(ctx, conf, remote, local)
 	if err != nil {
 		log.Errorf("ssh proxy err: %v", err)
 		return

pkg/util/ssh.go

@@ -67,37 +67,44 @@ func (s *SshConfig) ToRPC() *rpc.SshJump {
 	}
 }
 
-func PortMap(ctx context.Context, sshClient *ssh.Client, remoteEndpoint, localEndpoint netip.AddrPort, done chan struct{}) error {
-	// Listen on remote server port
-	var lc net.ListenConfig
-	listen, err := lc.Listen(ctx, "tcp", localEndpoint.String())
-	if err != nil {
-		return err
-	}
-	defer listen.Close()
+func PortMap(ctx context.Context, sshClient *ssh.Client, remoteEndpoint netip.AddrPort, local net.Listener) error {
+	defer sshClient.Close()
 
-	select {
-	case done <- struct{}{}:
-	default:
-	}
+	errChan := make(chan error, 1)
 	// handle incoming connections on reverse forwarded tunnel
-	for ctx.Err() == nil {
-		localConn, err := listen.Accept()
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case err := <-errChan:
+			return err
+		default:
+		}
+
+		localConn, err := local.Accept()
 		if err != nil {
 			return err
 		}
-		go func(localConn net.Conn) {
-			defer localConn.Close()
-			remoteConn, err := sshClient.Dial("tcp", remoteEndpoint.String())
+		go func() {
+			err := func() error {
+				defer localConn.Close()
+				remoteConn, err := sshClient.Dial("tcp", remoteEndpoint.String())
+				if err != nil {
+					log.Errorf("Failed to dial %s: %s", remoteEndpoint.String(), err)
+					return err
+				}
+				defer remoteConn.Close()
+				copyStream(localConn, remoteConn)
+				return nil
+			}()
 			if err != nil {
-				log.Errorf("Failed to dial %s: %s", remoteEndpoint.String(), err)
-				return
+				select {
+				case errChan <- err:
+				default:
+				}
 			}
-			defer remoteConn.Close()
-			copyStream(localConn, remoteConn)
-		}(localConn)
+		}()
 	}
-	return ctx.Err()
 }
 
 // DialSshRemote https://github.com/golang/go/issues/21478
@@ -225,12 +232,12 @@ func publicKeyFile(file string) (ssh.AuthMethod, error) {
 	return ssh.PublicKeys(key), nil
 }
 
-func copyStream(client net.Conn, remote net.Conn) {
+func copyStream(local net.Conn, remote net.Conn) {
 	chDone := make(chan bool, 2)
 
 	// start remote -> local data transfer
 	go func() {
-		_, err := io.Copy(client, remote)
+		_, err := io.Copy(local, remote)
 		if err != nil && !errors.Is(err, net.ErrClosed) {
 			log.Debugf("error while copy remote->local: %s", err)
 		}
@@ -242,7 +249,7 @@ func copyStream(client net.Conn, remote net.Conn) {
 
 	// start local -> remote data transfer
 	go func() {
-		_, err := io.Copy(remote, client)
+		_, err := io.Copy(remote, local)
 		if err != nil && !errors.Is(err, net.ErrClosed) {
 			log.Debugf("error while copy local->remote: %s", err)
 		}
@@ -388,22 +395,44 @@ func init() {
 	})
 }
 
-func PortMapUntil(ctx context.Context, cli *ssh.Client, remote, local netip.AddrPort) error {
+func PortMapUntil(ctx context.Context, conf *SshConfig, remote, local netip.AddrPort) error {
+	// Listen on remote server port
+	var lc net.ListenConfig
+	localListen, err := lc.Listen(ctx, "tcp", local.String())
+	if err != nil {
+		return err
+	}
+	defer localListen.Close()
+
 	errChan := make(chan error, 1)
 	readyChan := make(chan struct{}, 1)
 	go func() {
 		for ctx.Err() == nil {
-			err := PortMap(ctx, cli, remote, local, readyChan)
-			if err != nil {
-				if !errors.Is(err, context.Canceled) {
-					log.Errorf("Ssh forward failed err: %v", err)
+			func() {
+				ctx2, cancelFunc := context.WithCancel(ctx)
+				defer cancelFunc()
+				cli, err := DialSshRemote(ctx2, conf)
+				if err != nil {
+					time.Sleep(time.Second * 2)
+					log.Errorf("failed to dial remote ssh server: %v", err)
+					return
 				}
 				select {
-				case errChan <- err:
+				case readyChan <- struct{}{}:
 				default:
 				}
-			}
-			time.Sleep(time.Second * 2)
+				err = PortMap(ctx, cli, remote, localListen)
+				if err != nil {
+					if !errors.Is(err, context.Canceled) {
+						log.Errorf("Ssh forward failed err: %v", err)
+					}
+					select {
+					case errChan <- err:
+					default:
+					}
+				}
+				time.Sleep(time.Second * 2)
+			}()
 		}
 	}()
 	select {