From 638994e52636849919341c84b541522d653c4db6 Mon Sep 17 00:00:00 2001
From: Artur Troian <troian.ap@gmail.com>
Date: Thu, 21 Mar 2024 23:37:47 +0100
Subject: [PATCH] fix(lease/shell): command with multiple arguments/flag to be
 properly processed by cobra (#219)

if lease shell command contains flags, cobra will treat them as own
flags and mostly likely will fail. enclosing lease command into double
brackets prevents it. in command itself it will be properly tokenized
and honor inner double quotes if any

Signed-off-by: Artur Troian <troian.ap@gmail.com>
---
 cmd/provider-services/cmd/shell.go | 23 +++++++++++++++++------
 gateway/rest/client.go             |  2 +-
 gateway/rest/client_shell.go       |  9 +--------
 3 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/cmd/provider-services/cmd/shell.go b/cmd/provider-services/cmd/shell.go
index a1bf1436..36d5327b 100644
--- a/cmd/provider-services/cmd/shell.go
+++ b/cmd/provider-services/cmd/shell.go
@@ -11,6 +11,7 @@ import (
 	"sync"
 	"syscall"
 
+	"github.com/go-andiamo/splitter"
 	dockerterm "github.com/moby/term"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -66,16 +67,14 @@ func LeaseShellCmd() *cobra.Command {
 }
 
 func doLeaseShell(cmd *cobra.Command, args []string) error {
-	var stdin io.ReadCloser
-	var stdout io.Writer
-	var stderr io.Writer
-	stdout = os.Stdout
-	stderr = os.Stderr
+	var stdin io.Reader
+	stdout := cmd.OutOrStdout()
+	stderr := cmd.ErrOrStderr()
 	connectStdin := viper.GetBool(FlagStdin)
 	setupTty := viper.GetBool(FlagTty)
 	podIndex := viper.GetUint(FlagReplicaIndex)
 	if connectStdin || setupTty {
-		stdin = os.Stdin
+		stdin = cmd.InOrStdin()
 	}
 
 	var tty term.TTY
@@ -138,6 +137,18 @@ func doLeaseShell(cmd *cobra.Command, args []string) error {
 	service := args[0]
 	remoteCmd := args[1:]
 
+	if len(remoteCmd) == 1 {
+		spaceSplitter, err := splitter.NewSplitter(' ', splitter.DoubleQuotes)
+		if err != nil {
+			return err
+		}
+
+		remoteCmd, err = spaceSplitter.Split(remoteCmd[0])
+		if err != nil {
+			return err
+		}
+	}
+
 	var terminalResizes chan remotecommand.TerminalSize
 	wg := &sync.WaitGroup{}
 	ctx, cancel := context.WithCancel(ctx)
diff --git a/gateway/rest/client.go b/gateway/rest/client.go
index 871aca2c..44b9f80d 100644
--- a/gateway/rest/client.go
+++ b/gateway/rest/client.go
@@ -52,7 +52,7 @@ type Client interface {
 	LeaseLogs(ctx context.Context, id mtypes.LeaseID, services string, follow bool, tailLines int64) (*ServiceLogs, error)
 	ServiceStatus(ctx context.Context, id mtypes.LeaseID, service string) (*cltypes.ServiceStatus, error)
 	LeaseShell(ctx context.Context, id mtypes.LeaseID, service string, podIndex uint, cmd []string,
-		stdin io.ReadCloser,
+		stdin io.Reader,
 		stdout io.Writer,
 		stderr io.Writer,
 		tty bool,
diff --git a/gateway/rest/client_shell.go b/gateway/rest/client_shell.go
index c1232d7f..f8ef6a3b 100644
--- a/gateway/rest/client_shell.go
+++ b/gateway/rest/client_shell.go
@@ -24,7 +24,7 @@ var (
 )
 
 func (c *client) LeaseShell(ctx context.Context, lID mtypes.LeaseID, service string, podIndex uint, cmd []string,
-	stdin io.ReadCloser,
+	stdin io.Reader,
 	stdout io.Writer,
 	stderr io.Writer,
 	tty bool,
@@ -155,13 +155,6 @@ loop:
 
 	subcancel()
 
-	if stdin != nil {
-		err := stdin.Close()
-		if err != nil {
-			saveError("closing stdin", err)
-		}
-	}
-
 	// Check to see if the remote end returned an error
 	if remoteErrorData != nil {
 		if err = processRemoteError(remoteErrorData); err != nil {