From 603fde3cfffc4894933251ff35eacbe0d851df26 Mon Sep 17 00:00:00 2001
From: haozi007 <duguhaotian@gmail.com>
Date: Tue, 5 Dec 2023 19:31:02 +0800
Subject: [PATCH] support websocket for exec and attach

Signed-off-by: haozi007 <duguhaotian@gmail.com>
---
 cmd/crictl/exec.go | 22 +++++++++++++++++++++-
 cmd/crictl/util.go |  4 ++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/cmd/crictl/exec.go b/cmd/crictl/exec.go
index 58b3068b32..59a0d59677 100644
--- a/cmd/crictl/exec.go
+++ b/cmd/crictl/exec.go
@@ -26,6 +26,7 @@ import (
 	mobyterm "github.com/moby/term"
 	"github.com/sirupsen/logrus"
 	"github.com/urfave/cli/v2"
+	"k8s.io/apimachinery/pkg/util/httpstream"
 	restclient "k8s.io/client-go/rest"
 	remoteclient "k8s.io/client-go/tools/remotecommand"
 	internalapi "k8s.io/cri-api/pkg/apis"
@@ -163,8 +164,27 @@ func Exec(ctx context.Context, client internalapi.RuntimeService, opts execOptio
 	return stream(ctx, opts.stdin, opts.tty, URL)
 }
 
+func getExecutor(url *url.URL) (exec remoteclient.Executor, err error) {
+	config := &restclient.Config{TLSClientConfig: restclient.TLSClientConfig{Insecure: true}}
+
+	exec, err = remoteclient.NewSPDYExecutor(config, "POST", url)
+	if err != nil {
+		return
+	}
+	if !IsEnabledWebsockets() {
+		return
+	}
+	wexec, err := remoteclient.NewWebSocketExecutor(config, "GET", url.String())
+	if err != nil {
+		return
+	}
+
+	exec, err = remoteclient.NewFallbackExecutor(wexec, exec, httpstream.IsUpgradeFailure)
+	return
+}
+
 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)
+	executor, err := getExecutor(url)
 	if err != nil {
 		return err
 	}
diff --git a/cmd/crictl/util.go b/cmd/crictl/util.go
index d9be5444cb..4bb1cf182c 100644
--- a/cmd/crictl/util.go
+++ b/cmd/crictl/util.go
@@ -70,6 +70,10 @@ func SetupInterruptSignalHandler() <-chan struct{} {
 	return signalIntStopCh
 }
 
+func IsEnabledWebsockets() bool {
+	return strings.ToLower(os.Getenv("CRICTL_REMOTE_COMMAND_WEBSOCKETS")) == "true"
+}
+
 type listOptions struct {
 	// id of container or sandbox
 	id string