feat(virtualization): console/vnc reconnect

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

fix: rename var

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

feat(virtualization): vnc reconnect

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

fix

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

fix

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

go.mod

@@ -4,7 +4,7 @@ go 1.23.1
 
 require (
 	github.com/Masterminds/semver/v3 v3.3.0
-	github.com/deckhouse/virtualization/api v0.0.0-20241101085803-1002322cdb92
+	github.com/deckhouse/virtualization/api v0.0.0-20241205091855-6f05a202ade8
 	github.com/google/go-containerregistry v0.20.0
 	github.com/gorilla/websocket v1.5.3
 	github.com/hashicorp/go-cleanhttp v0.5.2

go.sum

@@ -394,6 +394,8 @@ github.com/daviddengcn/go-colortext v1.0.0 h1:ANqDyC0ys6qCSvuEK7l3g5RaehL/Xck9EX
 github.com/daviddengcn/go-colortext v1.0.0/go.mod h1:zDqEI5NVUop5QPpVJUxE9UO10hRnmkD5G4Pmri9+m4c=
 github.com/deckhouse/virtualization/api v0.0.0-20241101085803-1002322cdb92 h1:EeO0Ly13a3DhoY9Q/Fz39MdhE6q2A4+cLi/d3nAGjpQ=
 github.com/deckhouse/virtualization/api v0.0.0-20241101085803-1002322cdb92/go.mod h1:t+6i4NC43RfNLqcZqkEc5vxY1ypKceqmOOKlVEq0cYA=
+github.com/deckhouse/virtualization/api v0.0.0-20241205091855-6f05a202ade8 h1:GAz+0wP9q4iHs9ogHfNUApNCGlYsXomkRk7iKpx/fAg=
+github.com/deckhouse/virtualization/api v0.0.0-20241205091855-6f05a202ade8/go.mod h1:t+6i4NC43RfNLqcZqkEc5vxY1ypKceqmOOKlVEq0cYA=
 github.com/denisenkom/go-mssqldb v0.0.0-20191128021309-1d7a30a10f73/go.mod h1:xbL0rPBG9cCiLr28tMa8zpbdarY27NDyej4t/EjAShU=
 github.com/denisenkom/go-mssqldb v0.12.2 h1:1OcPn5GBIobjWNd+8yjfHNIaFX14B1pWI3F9HZy5KXw=
 github.com/denisenkom/go-mssqldb v0.12.2/go.mod h1:lnIw1mZukFRZDJYQ0Pb833QS2IaC3l5HkEfra2LJ+sk=

internal/virtualization/cmd/console/console.go

@@ -20,10 +20,11 @@ Initially copied from https://github.com/kubevirt/kubevirt/blob/main/pkg/virtctl
 package console
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"os"
-	"os/signal"
+	"strings"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -86,6 +87,64 @@ func (c *Console) Run(args []string) error {
 		return err
 	}
 
+	stdinCh := make(chan []byte)
+	go func() {
+		in := os.Stdin
+		defer close(stdinCh)
+		buf := make([]byte, 1024)
+		for {
+			// reading from stdin
+			n, err := in.Read(buf)
+			if err != nil && err != io.EOF {
+				return
+			}
+			if n == 0 && err == io.EOF {
+				return
+			}
+
+			// the escape sequence
+			if buf[0] == 29 {
+				return
+			}
+
+			stdinCh <- buf[0:n]
+		}
+	}()
+
+	go func() {
+		if _, ok := <-stdinCh; !ok {
+			os.Exit(0)
+		}
+	}()
+
+	for {
+		err := connect(name, namespace, virtCli, stdinCh)
+		if err != nil {
+			if errors.Is(err, util.ErrorInterrupt) || strings.Contains(err.Error(), "not found") {
+				return nil
+			}
+
+			if e, ok := err.(*websocket.CloseError); ok && e.Code == websocket.CloseGoingAway {
+				fmt.Fprint(os.Stderr, "\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")
+
+				return nil
+			}
+
+			if e, ok := err.(*websocket.CloseError); ok && e.Code == websocket.CloseAbnormalClosure {
+				fmt.Fprint(os.Stderr, "\nYou were disconnected from the console. This has one of the following reasons:"+
+					"\n - network issues"+
+					"\n - machine restart\n")
+			} else {
+				fmt.Fprintf(os.Stderr, "%s\n", err)
+			}
+
+			time.Sleep(time.Second)
+		}
+	}
+}
+
+func connect(name string, namespace string, virtCli kubeclient.Client, stdinCh chan []byte) error {
 	stdinReader, stdinWriter := io.Pipe()
 	stdoutReader, stdoutWriter := io.Pipe()
 
@@ -94,8 +153,6 @@ func (c *Console) Run(args []string) error {
 	// Wait until the virtual machine is in running phase, user interrupt or timeout
 	resChan := make(chan error)
 	runningChan := make(chan error)
-	waitInterrupt := make(chan os.Signal, 1)
-	signal.Notify(waitInterrupt, os.Interrupt)
 
 	go func() {
 		con, err := virtCli.VirtualMachines(namespace).SerialConsole(name, &kubeclient.SerialConsoleOptions{ConnectionTimeout: time.Duration(timeout) * time.Minute})
@@ -111,27 +168,11 @@ func (c *Console) Run(args []string) error {
 		})
 	}()
 
-	select {
-	case <-waitInterrupt:
-		// Make a new line in the terminal
-		fmt.Println()
-		return nil
-	case err = <-runningChan:
-		if err != nil {
-			return err
-		}
-	}
-	err = util.AttachConsole(stdinReader, stdoutReader, stdinWriter, stdoutWriter,
-		fmt.Sprint("Successfully connected to ", name, " console. The escape sequence is ^]\n"),
-		resChan)
-
+	err := <-runningChan
 	if err != nil {
-		if e, ok := err.(*websocket.CloseError); ok && e.Code == websocket.CloseAbnormalClosure {
-			fmt.Fprint(os.Stderr, "\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 - network issues\n")
-		}
 		return err
 	}
-	return nil
+
+	err = util.AttachConsole(stdinCh, stdinReader, stdoutReader, stdinWriter, stdoutWriter, name, resChan)
+	return err
 }

internal/virtualization/cmd/lifecycle/vmop/vmop.go

@@ -223,7 +223,7 @@ func (v VirtualMachineOperation) isPhaseOrFailed(vmop *v1alpha2.VirtualMachineOp
 func (v VirtualMachineOperation) newVMOP(vmName, vmNamespace string, t v1alpha2.VMOPType, force bool) *v1alpha2.VirtualMachineOperation {
 	return &v1alpha2.VirtualMachineOperation{
 		TypeMeta: metav1.TypeMeta{
-			Kind:       v1alpha2.VMOPKind,
+			Kind:       v1alpha2.VirtualMachineOperationKind,
 			APIVersion: v1alpha2.Version,
 		},
 		ObjectMeta: metav1.ObjectMeta{

internal/virtualization/cmd/vnc/vnc.go

@@ -20,6 +20,7 @@ Initially copied from https://github.com/kubevirt/kubevirt/blob/main/pkg/virtctl
 package vnc
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -30,11 +31,15 @@ import (
 	"os/signal"
 	"path/filepath"
 	"runtime"
+	"strings"
 	"time"
 
 	"github.com/deckhouse/deckhouse-cli/internal/virtualization/templates"
 	"github.com/deckhouse/virtualization/api/client/kubeclient"
+	"github.com/deckhouse/virtualization/api/core/v1alpha2"
+	"github.com/gorilla/websocket"
 	"github.com/spf13/cobra"
+	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/tools/clientcmd"
 	"k8s.io/klog/v2"
 )
@@ -109,11 +114,6 @@ func (o *VNC) Run(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	// setup connection with VM
-	vnc, err := virtCli.VirtualMachines(namespace).VNC(vmName)
-	if err != nil {
-		return fmt.Errorf("can't access VM %s: %s", vmName, err.Error())
-	}
 	// Format the listening address to account for the port (ex: 127.0.0.0:5900)
 	// Set listenAddress to localhost if proxy-only flag is not set
 	if !proxyOnly {
@@ -134,6 +134,52 @@ func (o *VNC) Run(cmd *cobra.Command, args []string) error {
 	// End of pre-flight checks. Everything looks good, we can start
 	// the goroutines and let the data flow
 
+	for {
+		err := connect(ln, virtCli, cmd, namespace, vmName)
+		if err != nil {
+			if strings.Contains(err.Error(), "not found") {
+				return err
+			}
+
+			if e, ok := err.(*websocket.CloseError); ok && e.Code == websocket.CloseGoingAway {
+				fmt.Fprint(os.Stderr, "\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")
+
+				return nil
+			}
+
+			if e, ok := err.(*websocket.CloseError); ok && e.Code == websocket.CloseAbnormalClosure {
+				fmt.Fprint(os.Stderr, "\nYou were disconnected from the console. This has one of the following reasons:"+
+					"\n - network issues"+
+					"\n - machine restart\n")
+			} else {
+				fmt.Fprintf(os.Stderr, "%s\n", err)
+			}
+
+			time.Sleep(time.Second)
+			continue
+		}
+
+		return nil
+	}
+}
+
+func connect(ln *net.TCPListener, virtCli kubeclient.Client, cmd *cobra.Command, namespace, vmName string) (err error) {
+	vm, err := virtCli.VirtualMachines(namespace).Get(context.Background(), vmName, v1.GetOptions{})
+	if err != nil {
+		return err
+	}
+
+	if vm.Status.Phase != v1alpha2.MachineRunning {
+		return errors.New("VM is not running")
+	}
+
+	// setup connection with VM
+	vnc, err := virtCli.VirtualMachines(namespace).VNC(vmName)
+	if err != nil {
+		return fmt.Errorf("can't access VM %s: %s", vmName, err.Error())
+	}
+
 	//                                       -> pipeInWriter  -> pipeInReader
 	// remote-viewer -> unix sock connection
 	//                                       <- pipeOutReader <- pipeOutWriter
@@ -197,6 +243,8 @@ func (o *VNC) Run(cmd *cobra.Command, args []string) error {
 
 	port := ln.Addr().(*net.TCPAddr).Port
 
+	ctx, cancelCtx := context.WithCancel(context.Background())
+
 	if proxyOnly {
 		defer close(doneChan)
 		optionString, err := json.Marshal(struct {
@@ -208,7 +256,7 @@ func (o *VNC) Run(cmd *cobra.Command, args []string) error {
 		fmt.Fprintln(cmd.OutOrStdout(), string(optionString))
 	} else {
 		// execute VNC Viewer
-		go checkAndRunVNCViewer(doneChan, viewResChan, port)
+		go checkAndRunVNCViewer(ctx, doneChan, viewResChan, port)
 	}
 
 	go func() {
@@ -227,13 +275,12 @@ func (o *VNC) Run(cmd *cobra.Command, args []string) error {
 	case err = <-listenResChan:
 	}
 
-	if err != nil {
-		return fmt.Errorf("error encountered: %s", err.Error())
-	}
-	return nil
+	cancelCtx()
+
+	return err
 }
 
-func checkAndRunVNCViewer(doneChan chan struct{}, viewResChan chan error, port int) {
+func checkAndRunVNCViewer(ctx context.Context, doneChan chan struct{}, viewResChan chan error, port int) {
 	defer close(doneChan)
 	var err error
 	args := []string{}
@@ -293,7 +340,7 @@ func checkAndRunVNCViewer(doneChan chan struct{}, viewResChan chan error, port i
 	} else {
 		klog.V(4).Infof("Executing commandline: '%s %v'", vncBin, args)
 		// #nosec No risk for attacket injection. vncBin and args include predefined strings
-		cmd := exec.Command(vncBin, args...)
+		cmd := exec.CommandContext(ctx, vncBin, args...)
 		output, err := cmd.CombinedOutput()
 		if err != nil {
 			klog.Errorf("%s execution failed: %v, output: %v", vncBin, err, string(output))

internal/virtualization/util/console.go

@@ -20,16 +20,17 @@ Initially copied from https://github.com/kubevirt/kubevirt/blob/main/pkg/virtctl
 package util
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"os"
-	"os/signal"
 
 	"golang.org/x/term"
 )
 
-func AttachConsole(stdinReader, stdoutReader *io.PipeReader, stdinWriter, stdoutWriter *io.PipeWriter, message string, resChan <-chan error) (err error) {
-	stopChan := make(chan struct{}, 1)
+var ErrorInterrupt = errors.New("interrupt")
+
+func AttachConsole(stdinCh chan []byte, stdinReader, stdoutReader *io.PipeReader, stdinWriter, stdoutWriter *io.PipeWriter, name string, resChan <-chan error) (err error) {
 	writeStop := make(chan error)
 	readStop := make(chan error)
 	if term.IsTerminal(int(os.Stdin.Fd())) {
@@ -39,55 +40,38 @@ func AttachConsole(stdinReader, stdoutReader *io.PipeReader, stdinWriter, stdout
 		}
 		defer term.Restore(int(os.Stdin.Fd()), state)
 	}
-	fmt.Fprint(os.Stderr, message)
-
-	in := os.Stdin
-	out := os.Stdout
 
-	go func() {
-		interrupt := make(chan os.Signal, 1)
-		signal.Notify(interrupt, os.Interrupt)
-		<-interrupt
-		close(stopChan)
-	}()
+	fmt.Fprintf(os.Stderr, "Successfully connected to %s console. The escape sequence is ^]\n", name)
 
+	out := os.Stdout
 	go func() {
+		defer close(readStop)
 		_, err := io.Copy(out, stdoutReader)
 		readStop <- err
 	}()
 
 	go func() {
 		defer close(writeStop)
-		buf := make([]byte, 1024)
-		for {
-			// reading from stdin
-			n, err := in.Read(buf)
-			if err != nil && err != io.EOF {
-				writeStop <- err
-				return
-			}
-			if n == 0 && err == io.EOF {
-				return
-			}
 
-			// the escape sequence
-			if buf[0] == 29 {
-				return
-			}
-			// Writing out to the console connection
-			_, err = stdinWriter.Write(buf[0:n])
+		stdinWriter.Write([]byte("\r"))
+		if err == io.EOF {
+			return
+		}
+
+		for b := range stdinCh {
+			_, err = stdinWriter.Write(b)
 			if err == io.EOF {
 				return
 			}
 		}
 	}()
 
 	select {
-	case <-stopChan:
-	case err = <-readStop:
 	case err = <-writeStop:
+		return ErrorInterrupt
+	case err = <-readStop:
+		return ErrorInterrupt
 	case err = <-resChan:
+		return err
 	}
-
-	return err
 }