Fix: Runtime bugs

cli/packages/util/exec.go

@@ -1,24 +1,22 @@
 package util
 
 import (
+	"fmt"
 	"os"
 	"os/exec"
 	"os/signal"
 	"runtime"
-	"strings"
 	"syscall"
-
-	"github.com/mattn/go-isatty"
 )
 
-func RunCommand(singleCommand string, args []string, env []string) (*exec.Cmd, error) {
+func RunCommand(singleCommand string, args []string, env []string, waitForExit bool) (*exec.Cmd, error) {
 	var c *exec.Cmd
 	var err error
 
 	if singleCommand != "" {
-		c, err = RunCommandFromString(singleCommand, env)
+		c, err = RunCommandFromString(singleCommand, env, waitForExit)
 	} else {
-		c, err = RunCommandFromArgs(args, env)
+		c, err = RunCommandFromArgs(args, env, waitForExit)
 	}
 
 	return c, err
@@ -30,66 +28,65 @@ func IsProcessRunning(p *os.Process) bool {
 }
 
 // For "infisical run -- COMMAND"
-func RunCommandFromArgs(command []string, env []string) (*exec.Cmd, error) {
-	cmd := exec.Command(command[0], command[1:]...)
-	cmd.Env = env
+func RunCommandFromArgs(args []string, env []string, waitForExit bool) (*exec.Cmd, error) {
+	cmd := exec.Command(args[0], args[1:]...)
 	cmd.Stdin = os.Stdin
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
+	cmd.Env = env
 
-	err := execCommand(cmd)
+	err := execCommand(cmd, waitForExit)
 
 	return cmd, err
 }
 
-func execCommand(cmd *exec.Cmd) error {
-
-	shouldForward := !isatty.IsTerminal(os.Stdout.Fd())
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan)
+func execCommand(cmd *exec.Cmd, waitForExit bool) error {
+	sigChannel := make(chan os.Signal, 1)
+	signal.Notify(sigChannel)
 
 	if err := cmd.Start(); err != nil {
 		return err
 	}
 
-	// handle all signals
 	go func() {
 		for {
-			if shouldForward {
-				// forward to process
-				sig := <-sigChan
-				cmd.Process.Signal(sig)
-			} else {
-				<-sigChan
-			}
+			sig := <-sigChannel
+			_ = cmd.Process.Signal(sig) // process all sigs
 		}
 	}()
 
+	if !waitForExit {
+		return nil
+	}
+
+	if err := cmd.Wait(); err != nil {
+		_ = cmd.Process.Signal(os.Kill)
+		return fmt.Errorf("failed to wait for command termination: %v", err)
+	}
+
+	waitStatus := cmd.ProcessState.Sys().(syscall.WaitStatus)
+	os.Exit(waitStatus.ExitStatus())
 	return nil
 }
 
 // For "infisical run --command=COMMAND"
-func RunCommandFromString(command string, env []string) (*exec.Cmd, error) {
+func RunCommandFromString(command string, env []string, waitForExit bool) (*exec.Cmd, error) {
 	shell := [2]string{"sh", "-c"}
 	if runtime.GOOS == "windows" {
 		shell = [2]string{"cmd", "/C"}
 	} else {
-		// these shells all support the same options we use for sh
-		shells := []string{"/bash", "/dash", "/fish", "/zsh", "/ksh", "/csh", "/tcsh"}
-		envShell := os.Getenv("SHELL")
-		for _, s := range shells {
-			if strings.HasSuffix(envShell, s) {
-				shell[0] = envShell
-				break
-			}
+		currentShell := os.Getenv("SHELL")
+		if currentShell != "" {
+			shell[0] = currentShell
 		}
 	}
+
 	cmd := exec.Command(shell[0], shell[1], command) // #nosec G204 nosemgrep: semgrep_configs.prohibit-exec-command
 	cmd.Env = env
 	cmd.Stdin = os.Stdin
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 
-	err := execCommand(cmd)
+	err := execCommand(cmd, waitForExit)
 	return cmd, err
 }