From 85bcb15cf50a149c643709348dc70332be2e607f Mon Sep 17 00:00:00 2001
From: Sascha Grunert <sgrunert@redhat.com>
Date: Thu, 21 Mar 2024 19:11:48 +0100
Subject: [PATCH] Support websocket transport for exec and attach

This allows end users to choose the transport to be used as well as
runtime developers to experiment on feature development.

Supersedes: https://github.com/kubernetes-sigs/cri-tools/pull/1305

Signed-off-by: Sascha Grunert <sgrunert@redhat.com>
---
 cmd/crictl/attach.go |  8 +++++++-
 cmd/crictl/exec.go   | 47 ++++++++++++++++++++++++++++++++++----------
 cmd/crictl/util.go   |  4 ++++
 3 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/cmd/crictl/attach.go b/cmd/crictl/attach.go
index 093241300e..79c12fdc1a 100644
--- a/cmd/crictl/attach.go
+++ b/cmd/crictl/attach.go
@@ -43,6 +43,12 @@ var runtimeAttachCommand = &cli.Command{
 			Aliases: []string{"i"},
 			Usage:   "Keep STDIN open",
 		},
+		&cli.StringFlag{
+			Name:    transportFlag,
+			Aliases: []string{"r"},
+			Value:   transportSpdy,
+			Usage:   fmt.Sprintf("Transport protocol to be used, must be one of: %s, %s", transportSpdy, transportWebsocket),
+		},
 	},
 	Action: func(c *cli.Context) error {
 		id := c.Args().First()
@@ -109,5 +115,5 @@ func Attach(ctx context.Context, client internalapi.RuntimeService, opts attachO
 	}
 
 	logrus.Debugf("Attach URL: %v", URL)
-	return stream(ctx, opts.stdin, opts.tty, URL)
+	return stream(ctx, opts.stdin, opts.tty, opts.transport, URL)
 }
diff --git a/cmd/crictl/exec.go b/cmd/crictl/exec.go
index 58b3068b32..57775481a3 100644
--- a/cmd/crictl/exec.go
+++ b/cmd/crictl/exec.go
@@ -26,7 +26,7 @@ import (
 	mobyterm "github.com/moby/term"
 	"github.com/sirupsen/logrus"
 	"github.com/urfave/cli/v2"
-	restclient "k8s.io/client-go/rest"
+	"k8s.io/client-go/rest"
 	remoteclient "k8s.io/client-go/tools/remotecommand"
 	internalapi "k8s.io/cri-api/pkg/apis"
 	pb "k8s.io/cri-api/pkg/apis/runtime/v1"
@@ -37,6 +37,10 @@ const (
 	// TODO: make this configurable in kubelet.
 	kubeletURLSchema = "http"
 	kubeletURLHost   = "http://127.0.0.1:10250"
+
+	transportFlag      = "transport"
+	transportWebsocket = "websocket"
+	transportSpdy      = "spdy"
 )
 
 const detachSequence = "ctrl-p,ctrl-q"
@@ -67,6 +71,12 @@ var runtimeExecCommand = &cli.Command{
 			Aliases: []string{"i"},
 			Usage:   "Keep STDIN open",
 		},
+		&cli.StringFlag{
+			Name:    transportFlag,
+			Aliases: []string{"r"},
+			Value:   transportSpdy,
+			Usage:   fmt.Sprintf("Transport protocol to be used, must be one of: %s, %s", transportSpdy, transportWebsocket),
+		},
 	},
 	Action: func(c *cli.Context) error {
 		if c.NArg() < 2 {
@@ -79,11 +89,12 @@ var runtimeExecCommand = &cli.Command{
 		}
 
 		var opts = execOptions{
-			id:      c.Args().First(),
-			timeout: c.Int64("timeout"),
-			tty:     c.Bool("tty"),
-			stdin:   c.Bool("interactive"),
-			cmd:     c.Args().Slice()[1:],
+			id:        c.Args().First(),
+			timeout:   c.Int64("timeout"),
+			tty:       c.Bool("tty"),
+			stdin:     c.Bool("interactive"),
+			cmd:       c.Args().Slice()[1:],
+			transport: c.String(transportFlag),
 		}
 		if c.Bool("sync") {
 			exitCode, err := ExecSync(runtimeClient, opts)
@@ -160,13 +171,13 @@ func Exec(ctx context.Context, client internalapi.RuntimeService, opts execOptio
 	}
 
 	logrus.Debugf("Exec URL: %v", URL)
-	return stream(ctx, opts.stdin, opts.tty, URL)
+	return stream(ctx, opts.stdin, opts.tty, opts.transport, URL)
 }
 
-func stream(ctx context.Context, in, tty bool, url *url.URL) error {
-	executor, err := remoteclient.NewSPDYExecutor(&restclient.Config{TLSClientConfig: restclient.TLSClientConfig{Insecure: true}}, "POST", url)
+func stream(ctx context.Context, in, tty bool, transport string, url *url.URL) error {
+	executor, err := getExecutor(transport, url)
 	if err != nil {
-		return err
+		return fmt.Errorf("get executor: %w", err)
 	}
 
 	stdin, stdout, stderr := mobyterm.StdStreams()
@@ -204,3 +215,19 @@ func stream(ctx context.Context, in, tty bool, url *url.URL) error {
 	streamOptions.TerminalSizeQueue = t.MonitorSize(t.GetSize())
 	return t.Safe(func() error { return executor.StreamWithContext(ctx, streamOptions) })
 }
+
+func getExecutor(transport string, url *url.URL) (exec remoteclient.Executor, err error) {
+	config := &rest.Config{TLSClientConfig: rest.TLSClientConfig{Insecure: true}}
+
+	switch transport {
+	case transportSpdy:
+		return remoteclient.NewSPDYExecutor(config, "POST", url)
+
+	case transportWebsocket:
+		return remoteclient.NewWebSocketExecutor(config, "GET", url.String())
+
+	default:
+		return nil, fmt.Errorf("unknown transport: %s", transport)
+
+	}
+}
diff --git a/cmd/crictl/util.go b/cmd/crictl/util.go
index 221b76264a..5d7cdcf340 100644
--- a/cmd/crictl/util.go
+++ b/cmd/crictl/util.go
@@ -114,6 +114,8 @@ type execOptions struct {
 	stdin bool
 	// Command to exec
 	cmd []string
+	// transport to be used
+	transport string
 }
 type attachOptions struct {
 	// id of container
@@ -122,6 +124,8 @@ type attachOptions struct {
 	tty bool
 	// Whether pass Stdin to container
 	stdin bool
+	// transport to be used
+	transport string
 }
 
 type portforwardOptions struct {