update references to exec.command

sreya · sreya · commit b8fc6bb6a139 · 2024-11-21T21:21:12.000Z
diff --git a/agent/agentexec/exec.go b/agent/agentexec/exec.go
@@ -10,6 +10,8 @@ import (
 	"strconv"
 
 	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/pty"
 )
 
 const (
@@ -25,19 +27,35 @@ const (
 // is returned. All instances of exec.Cmd should flow through this function to ensure
 // proper resource constraints are applied to the child process.
 func CommandContext(ctx context.Context, cmd string, args ...string) (*exec.Cmd, error) {
+	cmd, args, err := agentExecCmd(cmd, args...)
+	if err != nil {
+		return nil, xerrors.Errorf("agent exec cmd: %w", err)
+	}
+	return exec.CommandContext(ctx, cmd, args...), nil
+}
+
+func PTYCommandContext(ctx context.Context, cmd string, args ...string) (*pty.Cmd, error) {
+	cmd, args, err := agentExecCmd(cmd, args...)
+	if err != nil {
+		return nil, xerrors.Errorf("agent exec cmd: %w", err)
+	}
+	return pty.CommandContext(ctx, cmd, args...), nil
+}
+
+func agentExecCmd(cmd string, args ...string) (string, []string, error) {
 	_, enabled := os.LookupEnv(EnvProcPrioMgmt)
 	if runtime.GOOS != "linux" || !enabled {
-		return exec.CommandContext(ctx, cmd, args...), nil
+		return cmd, args, nil
 	}
 
 	executable, err := os.Executable()
 	if err != nil {
-		return nil, xerrors.Errorf("get executable: %w", err)
+		return "", nil, xerrors.Errorf("get executable: %w", err)
 	}
 
 	bin, err := filepath.EvalSymlinks(executable)
 	if err != nil {
-		return nil, xerrors.Errorf("eval symlinks: %w", err)
+		return "", nil, xerrors.Errorf("eval symlinks: %w", err)
 	}
 
 	execArgs := []string{"agent-exec"}
@@ -51,7 +69,7 @@ func CommandContext(ctx context.Context, cmd string, args ...string) (*exec.Cmd,
 	execArgs = append(execArgs, "--", cmd)
 	execArgs = append(execArgs, args...)
 
-	return exec.CommandContext(ctx, bin, execArgs...), nil
+	return bin, execArgs, nil
 }
 
 // envValInt searches for a key in a list of environment variables and parses it to an int.
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -30,6 +30,7 @@ import (
 
 	"cdr.dev/slog"
 
+	"github.com/coder/coder/v2/agent/agentexec"
 	"github.com/coder/coder/v2/agent/usershell"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/pty"
@@ -725,7 +726,10 @@ func (s *Server) CreateCommand(ctx context.Context, script string, env []string)
 		}
 	}
 
-	cmd := pty.CommandContext(ctx, name, args...)
+	cmd, err := agentexec.PTYCommandContext(ctx, name, args...)
+	if err != nil {
+		return nil, xerrors.Errorf("pty command context: %w", err)
+	}
 	cmd.Dir = s.config.WorkingDirectory()
 
 	// If the metadata directory doesn't exist, we run the command
diff --git a/agent/reconnectingpty/buffered.go b/agent/reconnectingpty/buffered.go
@@ -14,6 +14,7 @@ import (
 
 	"cdr.dev/slog"
 
+	"github.com/coder/coder/v2/agent/agentexec"
 	"github.com/coder/coder/v2/pty"
 )
 
@@ -58,7 +59,11 @@ func newBuffered(ctx context.Context, cmd *pty.Cmd, options *Options, logger slo
 
 	// Add TERM then start the command with a pty.  pty.Cmd duplicates Path as the
 	// first argument so remove it.
-	cmdWithEnv := pty.CommandContext(ctx, cmd.Path, cmd.Args[1:]...)
+	cmdWithEnv, err := agentexec.PTYCommandContext(ctx, cmd.Path, cmd.Args[1:]...)
+	if err != nil {
+		rpty.state.setState(StateDone, xerrors.Errorf("pty command context: %w", err))
+		return rpty
+	}
 	cmdWithEnv.Env = append(rpty.command.Env, "TERM=xterm-256color")
 	cmdWithEnv.Dir = rpty.command.Dir
 	ptty, process, err := pty.Start(cmdWithEnv)
diff --git a/agent/reconnectingpty/screen.go b/agent/reconnectingpty/screen.go
@@ -9,7 +9,6 @@ import (
 	"io"
 	"net"
 	"os"
-	"os/exec"
 	"path/filepath"
 	"strings"
 	"sync"
@@ -20,6 +19,7 @@ import (
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
+	"github.com/coder/coder/v2/agent/agentexec"
 	"github.com/coder/coder/v2/pty"
 )
 
@@ -210,7 +210,7 @@ func (rpty *screenReconnectingPTY) doAttach(ctx context.Context, conn net.Conn,
 	logger.Debug(ctx, "spawning screen client", slog.F("screen_id", rpty.id))
 
 	// Wrap the command with screen and tie it to the connection's context.
-	cmd := pty.CommandContext(ctx, "screen", append([]string{
+	cmd, err := agentexec.PTYCommandContext(ctx, "screen", append([]string{
 		// -S is for setting the session's name.
 		"-S", rpty.id,
 		// -U tells screen to use UTF-8 encoding.
@@ -223,6 +223,9 @@ func (rpty *screenReconnectingPTY) doAttach(ctx context.Context, conn net.Conn,
 		rpty.command.Path,
 		// pty.Cmd duplicates Path as the first argument so remove it.
 	}, rpty.command.Args[1:]...)...)
+	if err != nil {
+		return nil, nil, xerrors.Errorf("pty command context: %w", err)
+	}
 	cmd.Env = append(rpty.command.Env, "TERM=xterm-256color")
 	cmd.Dir = rpty.command.Dir
 	ptty, process, err := pty.Start(cmd, pty.WithPTYOption(
@@ -327,29 +330,32 @@ func (rpty *screenReconnectingPTY) sendCommand(ctx context.Context, command stri
 	defer cancel()
 
 	var lastErr error
-	run := func() bool {
+	run := func() (bool, error) {
 		var stdout bytes.Buffer
 		//nolint:gosec
-		cmd := exec.CommandContext(ctx, "screen",
+		cmd, err := agentexec.CommandContext(ctx, "screen",
 			// -x targets an attached session.
 			"-x", rpty.id,
 			// -c is the flag for the config file.
 			"-c", rpty.configFile,
 			// -X runs a command in the matching session.
 			"-X", command,
 		)
+		if err != nil {
+			return false, xerrors.Errorf("command context: %w", err)
+		}
 		cmd.Env = append(rpty.command.Env, "TERM=xterm-256color")
 		cmd.Dir = rpty.command.Dir
 		cmd.Stdout = &stdout
-		err := cmd.Run()
+		err = cmd.Run()
 		if err == nil {
-			return true
+			return true, nil
 		}
 
 		stdoutStr := stdout.String()
 		for _, se := range successErrors {
 			if strings.Contains(stdoutStr, se) {
-				return true
+				return true, nil
 			}
 		}
 
@@ -359,11 +365,15 @@ func (rpty *screenReconnectingPTY) sendCommand(ctx context.Context, command stri
 			lastErr = xerrors.Errorf("`screen -x %s -X %s`: %w: %s", rpty.id, command, err, stdoutStr)
 		}
 
-		return false
+		return false, nil
 	}
 
 	// Run immediately.
-	if done := run(); done {
+	done, err := run()
+	if err != nil {
+		return err
+	}
+	if done {
 		return nil
 	}
 
@@ -379,7 +389,11 @@ func (rpty *screenReconnectingPTY) sendCommand(ctx context.Context, command stri
 			}
 			return errors.Join(ctx.Err(), lastErr)
 		case <-ticker.C:
-			if done := run(); done {
+			done, err := run()
+			if err != nil {
+				return err
+			}
+			if done {
 				return nil
 			}
 		}
diff --git a/scripts/rules.go b/scripts/rules.go
@@ -498,7 +498,11 @@ func noExecInAgent(m dsl.Matcher) {
 		`exec.Command($*_)`,
 		`exec.CommandContext($*_)`,
 	).
-		Where(m.File().PkgPath.Matches("/agent/")).
+		Where(
+			m.File().PkgPath.Matches("/agent/") &&
+				!m.File().PkgPath.Matches("/agentexec") &&
+				!m.File().Name.Matches(`_test\.go$`),
+		).
 		Report("The agent and its subpackages should not use exec.Command or exec.CommandContext directly. Consider using agentexec.CommandContext instead.")
 }
 
@@ -512,6 +516,10 @@ func noPTYInAgent(m dsl.Matcher) {
 		`pty.Command($*_)`,
 		`pty.CommandContext($*_)`,
 	).
-		Where(m.File().PkgPath.Matches(`/agent/`)).
+		Where(
+			m.File().PkgPath.Matches(`/agent/`) &&
+				!m.File().PkgPath.Matches(`/agentexec`) &&
+				!m.File().Name.Matches(`_test\.go$`),
+		).
 		Report("The agent and its subpackages should not use pty.Command or pty.CommandContext directly. Consider using agentexec.PTYCommandContext instead.")
 }
