fix: refactor console and vnc for better consistency

Signed-off-by: Daniil Antoshin <[email protected]>

+

Signed-off-by: Daniil Antoshin <[email protected]>

Author: danilrwx

src/cli/go.mod

@@ -19,7 +19,6 @@ require (
 	k8s.io/client-go v0.32.2
 	k8s.io/component-base v0.29.3
 	k8s.io/klog/v2 v2.130.1
-	k8s.io/utils v0.0.0-20241210054802-24370beab758
 )
 
 require github.com/golang/protobuf v1.5.4 // indirect
@@ -69,6 +68,7 @@ require (
 	k8s.io/api v0.32.2 // indirect
 	k8s.io/apiextensions-apiserver v0.29.3 // indirect
 	k8s.io/kube-openapi v0.0.0-20241105132330-32ad38e42d3f // indirect
+	k8s.io/utils v0.0.0-20241210054802-24370beab758 // indirect
 	kubevirt.io/api v1.2.0 // indirect
 	kubevirt.io/containerized-data-importer-api v1.57.0-alpha1 // indirect
 	kubevirt.io/controller-lifecycle-operator-sdk/api v0.0.0-20220329064328-f3cc58c6ed90 // indirect

src/cli/internal/cmd/console/console.go

@@ -20,13 +20,17 @@ Initially copied from https://github.com/kubevirt/kubevirt/blob/main/pkg/virtctl
 package console
 
 import (
+	"context"
 	"errors"
+	"fmt"
 	"io"
+	"os"
 	"strings"
 	"time"
 
 	"github.com/gorilla/websocket"
 	"github.com/spf13/cobra"
+	"golang.org/x/term"
 
 	"github.com/deckhouse/virtualization/api/client/kubeclient"
 
@@ -79,64 +83,138 @@ func (c *Console) Run(cmd *cobra.Command, args []string) error {
 	}
 
 	for {
-		err := connect(name, namespace, client, c.timeout)
-		if err == nil {
-			continue
-		}
-
-		if errors.Is(err, util.ErrorInterrupt) || strings.Contains(err.Error(), "not found") {
-			return err
-		}
-
-		var e *websocket.CloseError
-		if errors.As(err, &e) {
-			switch e.Code {
-			case websocket.CloseGoingAway:
-				cmd.Printf("\nYou were disconnected from the console. This has one of the following reasons:" +
-					"\n - another user connected to the console of the target vm\n")
+		select {
+		case <-cmd.Context().Done():
+			return nil
+		default:
+			cmd.Printf("Connecting to %s via console...\n", name)
+
+			err := connect(cmd.Context(), name, namespace, client, c.timeout)
+			if err != nil {
+				if strings.Contains(err.Error(), "not found") {
+					return err
+				}
+
+				var e *websocket.CloseError
+				if errors.As(err, &e) {
+					switch e.Code {
+					case websocket.CloseGoingAway:
+						cmd.Printf(util.CloseGoingAwayMessage)
+						return nil
+					case websocket.CloseAbnormalClosure:
+						cmd.Printf(util.CloseAbnormalClosureMessage)
+					}
+				} else {
+					cmd.Printf("%s\n", err)
+				}
+
+				time.Sleep(time.Second)
+			} else {
 				return nil
-			case websocket.CloseAbnormalClosure:
-				cmd.Printf("\nYou were disconnected from the console. This has one of the following reasons:" +
-					"\n - network issues" +
-					"\n - machine restart\n")
 			}
-		} else {
-			cmd.Printf("%s\n", err)
 		}
-
-		time.Sleep(time.Second)
 	}
 }
 
-func connect(name string, namespace string, virtCli kubeclient.Client, timeout int) error {
+func connect(ctx context.Context, name string, namespace string, virtCli kubeclient.Client, timeout int) error {
+	// in -> stdinWriter | stdinReader -> console
+	// out <- stdoutReader | stdoutWriter <- console
 	stdinReader, stdinWriter := io.Pipe()
 	stdoutReader, stdoutWriter := io.Pipe()
 
-	// in -> stdinWriter | stdinReader -> console
-	// out <- stdoutReader | stdoutWriter <- console
-	// Wait until the virtual machine is in running phase, user interrupt or timeout
-	resChan := make(chan error)
-	runningChan := make(chan error)
+	doneChan := make(chan struct{}, 1)
 
-	go func() {
-		con, err := virtCli.VirtualMachines(namespace).SerialConsole(name, &kubeclient.SerialConsoleOptions{ConnectionTimeout: time.Duration(timeout) * time.Minute})
-		runningChan <- err
+	k8sResErr := make(chan error)
+	writeStopErr := make(chan error)
+	readStopErr := make(chan error)
 
-		if err != nil {
-			return
-		}
+	console, err := virtCli.VirtualMachines(namespace).SerialConsole(name, &kubeclient.SerialConsoleOptions{ConnectionTimeout: time.Duration(timeout) * time.Minute})
+	if err != nil {
+		return fmt.Errorf("can't access VM %s: %s", name, err.Error())
+	}
 
-		resChan <- con.Stream(kubeclient.StreamOptions{
+	go func() {
+		err := console.Stream(kubeclient.StreamOptions{
 			In:  stdinReader,
 			Out: stdoutWriter,
 		})
+		if err != nil {
+			k8sResErr <- err
+		}
 	}()
 
-	err := <-runningChan
-	if err != nil {
-		return err
+	if term.IsTerminal(int(os.Stdin.Fd())) {
+		state, err := term.MakeRaw(int(os.Stdin.Fd()))
+		if err != nil {
+			return fmt.Errorf("make raw terminal failed: %w", err)
+		}
+		defer term.Restore(int(os.Stdin.Fd()), state)
+	}
+
+	fmt.Fprintf(os.Stderr, "Successfully connected to %s console. The escape sequence is ^]\n", name)
+
+	out := os.Stdout
+	go func() {
+		_, err := io.Copy(out, stdoutReader)
+		if err != nil {
+			readStopErr <- err
+		}
+	}()
+
+	stdinCh := make(chan []byte)
+	go func() {
+		in := os.Stdin
+		buf := make([]byte, 1024)
+		for {
+			// reading from stdin
+			n, err := in.Read(buf)
+			if err != nil {
+				if err != io.EOF || n == 0 {
+					return
+				}
+
+				readStopErr <- err
+			}
+
+			// the escape sequence
+			if buf[0] == 29 {
+				doneChan <- struct{}{}
+				return
+			}
+
+			stdinCh <- buf[0:n]
+		}
+	}()
+
+	go func() {
+		_, err := stdinWriter.Write([]byte("\r"))
+		if err != nil {
+			if err == io.EOF {
+				return
+			}
+
+			writeStopErr <- err
+		}
+
+		for b := range stdinCh {
+			_, err = stdinWriter.Write(b)
+			if err != nil {
+				if err == io.EOF {
+					return
+				}
+
+				writeStopErr <- err
+			}
+		}
+	}()
+
+	select {
+	case <-ctx.Done():
+	case <-doneChan:
+	case err = <-k8sResErr:
+	case err = <-writeStopErr:
+	case err = <-readStopErr:
 	}
 
-	err = util.AttachConsole(stdinReader, stdoutReader, stdinWriter, stdoutWriter, name, resChan)
 	return err
 }