From 8b058cae07e51089cf517d8265b7e833b97e84ec Mon Sep 17 00:00:00 2001
From: Amir Raminfar <findamir@gmail.com>
Date: Wed, 2 Apr 2025 09:27:57 -0700
Subject: [PATCH] feat: adds shell support for k8s mode (#3747)

---
 .../ContainerActionsToolbar.vue               |  4 +-
 docs/guide/shell.md                           |  2 +-
 go.mod                                        |  2 +
 go.sum                                        |  6 ++
 internal/k8s/client.go                        | 87 ++++++++++++++++++-
 internal/support/k8s/k8s_service.go           | 61 ++++++++++++-
 6 files changed, 154 insertions(+), 8 deletions(-)

diff --git a/assets/components/ContainerViewer/ContainerActionsToolbar.vue b/assets/components/ContainerViewer/ContainerActionsToolbar.vue
index c816af13..28dcfb0e 100644
--- a/assets/components/ContainerViewer/ContainerActionsToolbar.vue
+++ b/assets/components/ContainerViewer/ContainerActionsToolbar.vue
@@ -135,7 +135,7 @@
         </li>
       </template>
 
-      <template v-if="enableShell && host.type !== 'k8s'">
+      <template v-if="enableShell">
         <li class="line"></li>
         <li>
           <a @click.prevent="showDrawer(Terminal, { container, action: 'attach' }, 'lg')">
@@ -157,7 +157,6 @@
 <script lang="ts" setup>
 import { Container } from "@/models/Container";
 import { allLevels } from "@/composable/logContext";
-const { hosts } = useHosts();
 import LogAnalytics from "../LogViewer/LogAnalytics.vue";
 import Terminal from "@/components/Terminal.vue";
 
@@ -169,7 +168,6 @@ const showDrawer = useDrawer();
 const { container } = defineProps<{ container: Container }>();
 const clear = defineEmit();
 const { actionStates, start, stop, restart } = useContainerActions(toRef(() => container));
-const host = computed(() => hosts.value[container.host]);
 
 onKeyStroke("f", (e) => {
   if (hasComplexLogs.value) {
diff --git a/docs/guide/shell.md b/docs/guide/shell.md
index a30a3f01..7794bc47 100644
--- a/docs/guide/shell.md
+++ b/docs/guide/shell.md
@@ -27,4 +27,4 @@ services:
 :::
 
 > [!NOTE]
-> Shell access only works in Docker containers. This includes agents, remote, swarm and local mode. Support for Kubernetes is planned, but it is not yet available.
+> Shell access should work across all container types, including Docker, Kubernetes, and other orchestration platforms.
diff --git a/go.mod b/go.mod
index e413eb0f..eea9b5c4 100644
--- a/go.mod
+++ b/go.mod
@@ -79,9 +79,11 @@ require (
 	github.com/mattn/go-colorable v0.1.14 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/moby/docker-image-spec v1.3.1 // indirect
+	github.com/moby/spdystream v0.5.0 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
+	github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/segmentio/asm v1.2.0 // indirect
diff --git a/go.sum b/go.sum
index 3e0099a0..496e2392 100644
--- a/go.sum
+++ b/go.sum
@@ -10,6 +10,8 @@ github.com/alexflint/go-scalar v1.2.0 h1:WR7JPKkeNpnYIOfHRa7ivM21aWAdHD0gEWHCx+W
 github.com/alexflint/go-scalar v1.2.0/go.mod h1:LoFvNMqS1CPrMVltza4LvnGKhaSpc3oyLEBUZVhhS2o=
 github.com/andybalholm/cascadia v1.3.3 h1:AG2YHrzJIm4BZ19iwJ/DAua6Btl3IwJX+VI4kktS1LM=
 github.com/andybalholm/cascadia v1.3.3/go.mod h1:xNd9bqTn98Ln4DwST8/nG+H0yuB8Hmgu1YHNnWw0GeA=
+github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio=
+github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs=
 github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
 github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
 github.com/beme/abide v0.0.0-20190723115211-635a09831760 h1:FvTM5NSN5HYvfKpgL+8x73U5v063vHsd7AX05eV1DnM=
@@ -124,6 +126,8 @@ github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWE
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0=
 github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo=
+github.com/moby/spdystream v0.5.0 h1:7r0J1Si3QO/kjRitvSLVVFUjxMEb/YLj6S9FF62JBCU=
+github.com/moby/spdystream v0.5.0/go.mod h1:xBAYlnt/ay+11ShkdFKNAG7LsyK/tmNBVvVOwrfMgdI=
 github.com/moby/term v0.0.0-20201216013528-df9cb8a40635 h1:rzf0wL0CHVc8CEsgyygG0Mn9CNCCPZqOPaz8RiiHYQk=
 github.com/moby/term v0.0.0-20201216013528-df9cb8a40635/go.mod h1:FBS0z0QWA44HXygs7VXDUOGoN/1TV3RuWkLO04am3wc=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
@@ -135,6 +139,8 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
 github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
+github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f h1:y5//uYreIhSUg3J1GEMiLbxo1LJaP8RfCpH6pymGZus=
+github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f/go.mod h1:ZdcZmHo+o7JKHSa8/e818NopupXU1YMK5fe1lsApnBw=
 github.com/onsi/ginkgo/v2 v2.21.0 h1:7rg/4f3rB88pb5obDgNZrNHrQ4e6WpjonchcpuBRnZM=
 github.com/onsi/ginkgo/v2 v2.21.0/go.mod h1:7Du3c42kxCUegi0IImZ1wUQzMBVecgIHjR1C+NkhLQo=
 github.com/onsi/gomega v1.35.1 h1:Cwbd75ZBPxFSuZ6T+rN/WCb/gOc6YgFBXLlZLhC7Ds4=
diff --git a/internal/k8s/client.go b/internal/k8s/client.go
index ddbaa8b9..50cbee79 100644
--- a/internal/k8s/client.go
+++ b/internal/k8s/client.go
@@ -17,8 +17,10 @@ import (
 	"github.com/rs/zerolog/log"
 
 	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/kubernetes/scheme"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/clientcmd"
+	"k8s.io/client-go/tools/remotecommand"
 )
 
 type K8sClient struct {
@@ -245,11 +247,92 @@ func (k *K8sClient) ContainerActions(ctx context.Context, action container.Conta
 }
 
 func (k *K8sClient) ContainerAttach(ctx context.Context, id string) (io.WriteCloser, io.Reader, error) {
-	panic("not implemented")
+	namespace, podName, containerName := parsePodContainerID(id)
+	log.Debug().Str("container", containerName).Str("pod", podName).Msg("Executing command in pod")
+	req := k.Clientset.CoreV1().RESTClient().Post().
+		Resource("pods").
+		Name(podName).
+		Namespace(namespace).
+		SubResource("attach")
+
+	option := &corev1.PodAttachOptions{
+		Container: containerName,
+		Stdin:     true,
+		Stdout:    true,
+		Stderr:    true,
+		TTY:       true,
+	}
+
+	req.VersionedParams(
+		option,
+		scheme.ParameterCodec,
+	)
+
+	exec, err := remotecommand.NewSPDYExecutor(k.config, "POST", req.URL())
+	if err != nil {
+		return nil, nil, err
+	}
+
+	stdinReader, stdinWriter := io.Pipe()
+	stdoutReader, stdoutWriter := io.Pipe()
+
+	go func() {
+		err := exec.StreamWithContext(ctx, remotecommand.StreamOptions{
+			Stdin:  stdinReader,
+			Stdout: stdoutWriter,
+			Tty:    true,
+		})
+		if err != nil {
+			log.Error().Err(err).Msg("Error streaming command")
+		}
+	}()
+
+	return stdinWriter, stdoutReader, nil
 }
 
 func (k *K8sClient) ContainerExec(ctx context.Context, id string, cmd []string) (io.WriteCloser, io.Reader, error) {
-	panic("not implemented")
+	namespace, podName, containerName := parsePodContainerID(id)
+	log.Debug().Str("container", containerName).Str("pod", podName).Msg("Executing command in pod")
+	req := k.Clientset.CoreV1().RESTClient().Post().
+		Resource("pods").
+		Name(podName).
+		Namespace(namespace).
+		SubResource("exec")
+
+	option := &corev1.PodExecOptions{
+		Command:   cmd,
+		Container: containerName,
+		Stdin:     true,
+		Stdout:    true,
+		Stderr:    true,
+		TTY:       true,
+	}
+
+	req.VersionedParams(
+		option,
+		scheme.ParameterCodec,
+	)
+
+	exec, err := remotecommand.NewSPDYExecutor(k.config, "POST", req.URL())
+	if err != nil {
+		return nil, nil, err
+	}
+
+	stdinReader, stdinWriter := io.Pipe()
+	stdoutReader, stdoutWriter := io.Pipe()
+
+	go func() {
+		err := exec.StreamWithContext(ctx, remotecommand.StreamOptions{
+			Stdin:  stdinReader,
+			Stdout: stdoutWriter,
+			Tty:    true,
+		})
+		if err != nil {
+			log.Error().Err(err).Msg("Error streaming command")
+		}
+	}()
+
+	return stdinWriter, stdoutReader, nil
 }
 
 // Helper function to parse pod and container names from container ID
diff --git a/internal/support/k8s/k8s_service.go b/internal/support/k8s/k8s_service.go
index 3efcebb6..1b748db2 100644
--- a/internal/support/k8s/k8s_service.go
+++ b/internal/support/k8s/k8s_service.go
@@ -3,6 +3,7 @@ package k8s_support
 import (
 	"context"
 	"io"
+	"sync"
 
 	"github.com/rs/zerolog/log"
 
@@ -92,9 +93,65 @@ func (k *K8sClientService) SubscribeContainersStarted(ctx context.Context, conta
 }
 
 func (k *K8sClientService) Attach(ctx context.Context, container container.Container, stdin io.Reader, stdout io.Writer) error {
-	panic("not implemented")
+	cancelCtx, cancel := context.WithCancel(ctx)
+	writer, reader, err := k.client.ContainerAttach(cancelCtx, container.ID)
+	if err != nil {
+		cancel()
+		return err
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer writer.Close()
+		defer cancel()
+		defer wg.Done()
+		if _, err := io.Copy(writer, stdin); err != nil {
+			log.Error().Err(err).Msg("error copying stdin")
+		}
+	}()
+
+	go func() {
+		defer cancel()
+		defer wg.Done()
+		if _, err := io.Copy(stdout, reader); err != nil {
+			log.Error().Err(err).Msg("error copying stdout")
+		}
+	}()
+
+	wg.Wait()
+	return nil
 }
 
 func (k *K8sClientService) Exec(ctx context.Context, container container.Container, cmd []string, stdin io.Reader, stdout io.Writer) error {
-	panic("not implemented")
+	cancelCtx, cancel := context.WithCancel(ctx)
+	writer, reader, err := k.client.ContainerExec(cancelCtx, container.ID, cmd)
+	if err != nil {
+		cancel()
+		return err
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer writer.Close()
+		defer cancel()
+		defer wg.Done()
+		if _, err := io.Copy(writer, stdin); err != nil {
+			log.Error().Err(err).Msg("error copying stdin")
+		}
+	}()
+
+	go func() {
+		defer cancel()
+		defer wg.Done()
+		if _, err := io.Copy(stdout, reader); err != nil {
+			log.Error().Err(err).Msg("error copying stdout")
+		}
+	}()
+
+	wg.Wait()
+	return nil
 }