✨ subprocess Add support for overriding the stdin/stdout/stderr of a subprocess (#705)

<!--
Copyright (C) 2020-2022 Arm Limited or its affiliates and Contributors.
All rights reserved.
SPDX-License-Identifier: Apache-2.0
-->
### Description

<!--
Please add any detail or context that would be useful to a reviewer.
-->

Add support for overriding the stdin/stdout/stderr of a subprocess.

Factories are used so that a context can be injected for context aware
readers and writers.

### Test Coverage

<!--
Please put an `x` in the correct box e.g. `[x]` to indicate the testing
coverage of this change.
-->

- [x]  This change is covered by existing or additional automated tests.
- [ ] Manual testing has been performed (and evidence provided) as
automated testing was not feasible.
- [ ] Additional tests are not required for this change (e.g.
documentation update).

Author: joshjennings98

changes/20250911113619.feature

@@ -0,0 +1 @@
+:sparkles: `subprocess` Add support for overriding the stdin/stdout/stderr of a subprocess

utils/mocks/mock_subprocess.go

@@ -127,6 +127,7 @@ type command struct {
 	as         *commandUtils.CommandAsDifferentUser
 	loggers    logs.Loggers
 	cmdWrapper cmdWrapper
+	io         ICommandIO
 }
 
 func (c *command) createCommand(cmdCtx context.Context) *exec.Cmd {
@@ -136,8 +137,7 @@ func (c *command) createCommand(cmdCtx context.Context) *exec.Cmd {
 	if err == nil {
 		cmd = cancellableCmd
 	}
-	cmd.Stdout = newOutStreamer(cmdCtx, c.loggers)
-	cmd.Stderr = newErrLogStreamer(cmdCtx, c.loggers)
+	cmd.Stdin, cmd.Stdout, cmd.Stderr = c.io.Register(cmdCtx)
 	cmd.Env = cmd.Environ()
 	cmd.Env = append(cmd.Env, c.env...)
 	setGroupAttrToCmd(cmd)
@@ -181,6 +181,20 @@ func newCommand(loggers logs.Loggers, as *commandUtils.CommandAsDifferentUser, e
 		as:         as,
 		loggers:    loggers,
 		cmdWrapper: cmdWrapper{},
+		io:         NewIOFromLoggers(loggers),
+	}
+	return
+}
+
+func newCommandWithCustomIO(loggers logs.Loggers, io ICommandIO, as *commandUtils.CommandAsDifferentUser, env []string, cmd string, args ...string) (osCmd *command) {
+	osCmd = &command{
+		cmd:        cmd,
+		args:       args,
+		env:        env,
+		as:         as,
+		loggers:    loggers,
+		cmdWrapper: cmdWrapper{},
+		io:         io,
 	}
 	return
 }

utils/subprocess/command_wrapper.go

@@ -47,7 +47,7 @@ func newSubProcess(ctx context.Context, loggers logs.Loggers, env []string, mess
 
 func newPlainSubProcess(ctx context.Context, loggers logs.Loggers, env []string, as *commandUtils.CommandAsDifferentUser, cmd string, args ...string) (p *Subprocess, err error) {
 	p = new(Subprocess)
-	err = p.setup(ctx, loggers, env, false, "", "", "", as, cmd, args...)
+	err = p.setup(ctx, loggers, nil, env, false, "", "", "", as, cmd, args...)
 	return
 }
 
@@ -130,7 +130,7 @@ func (s *Subprocess) Setup(ctx context.Context, loggers logs.Loggers, messageOnS
 
 // SetupWithEnvironment sets up a sub-process i.e. defines the command cmd and the messages on start, success and failure. Compared to Setup, it allows specifying additional environment variables to be used by the process.
 func (s *Subprocess) SetupWithEnvironment(ctx context.Context, loggers logs.Loggers, additionalEnv []string, messageOnStart string, messageOnSuccess, messageOnFailure string, cmd string, args ...string) (err error) {
-	return s.setup(ctx, loggers, additionalEnv, true, messageOnStart, messageOnSuccess, messageOnFailure, commandUtils.Me(), cmd, args...)
+	return s.setup(ctx, loggers, nil, additionalEnv, true, messageOnStart, messageOnSuccess, messageOnFailure, commandUtils.Me(), cmd, args...)
 }
 
 // SetupAs is similar to Setup but allows the command to be run as a different user.
@@ -140,11 +140,31 @@ func (s *Subprocess) SetupAs(ctx context.Context, loggers logs.Loggers, messageO
 
 // SetupAsWithEnvironment is similar to Setup but allows the command to be run as a different user. Compared to SetupAs, it allows specifying additional environment variables to be used by the process.
 func (s *Subprocess) SetupAsWithEnvironment(ctx context.Context, loggers logs.Loggers, additionalEnv []string, messageOnStart string, messageOnSuccess, messageOnFailure string, as *commandUtils.CommandAsDifferentUser, cmd string, args ...string) (err error) {
-	return s.setup(ctx, loggers, additionalEnv, true, messageOnStart, messageOnSuccess, messageOnFailure, as, cmd, args...)
+	return s.setup(ctx, loggers, nil, additionalEnv, true, messageOnStart, messageOnSuccess, messageOnFailure, as, cmd, args...)
 }
 
-// Setup sets up a sub-process i.e. defines the command cmd and the messages on start, success and failure.
-func (s *Subprocess) setup(ctx context.Context, loggers logs.Loggers, env []string, withAdditionalMessages bool, messageOnStart, messageOnSuccess, messageOnFailure string, as *commandUtils.CommandAsDifferentUser, cmd string, args ...string) (err error) {
+// SetupWithCustomIO sets up a sub-process i.e. defines the command cmd and the messages on start, success and failure. It allows the stdin, stdout, and stderr to be overridden.
+func (s *Subprocess) SetupWithCustomIO(ctx context.Context, loggers logs.Loggers, io ICommandIO, messageOnStart string, messageOnSuccess, messageOnFailure string, cmd string, args ...string) (err error) {
+	return s.SetupWithEnvironmentWithCustomIO(ctx, loggers, io, nil, messageOnStart, messageOnSuccess, messageOnFailure, cmd, args...)
+}
+
+// SetupWithEnvironmentWithCustomIO sets up a sub-process i.e. defines the command cmd and the messages on start, success and failure. Compared to SetupWithCustomIO, it allows specifying additional environment variables to be used by the process. It allows the stdin, stdout, and stderr to be overridden.
+func (s *Subprocess) SetupWithEnvironmentWithCustomIO(ctx context.Context, loggers logs.Loggers, io ICommandIO, additionalEnv []string, messageOnStart string, messageOnSuccess, messageOnFailure string, cmd string, args ...string) (err error) {
+	return s.setup(ctx, loggers, io, additionalEnv, true, messageOnStart, messageOnSuccess, messageOnFailure, commandUtils.Me(), cmd, args...)
+}
+
+// SetupAsWithCustomIO is similar to SetupWithCustomIO but allows the command to be run as a different user. It allows the stdin, stdout, and stderr to be overridden.
+func (s *Subprocess) SetupAsWithCustomIO(ctx context.Context, loggers logs.Loggers, io ICommandIO, messageOnStart string, messageOnSuccess, messageOnFailure string, as *commandUtils.CommandAsDifferentUser, cmd string, args ...string) (err error) {
+	return s.SetupAsWithEnvironmentWithCustomIO(ctx, loggers, io, nil, messageOnStart, messageOnSuccess, messageOnFailure, as, cmd, args...)
+}
+
+// SetupAsWithEnvironmentWithCustomIO is similar to SetupWithCustomIO but allows the command to be run as a different user. Compared to SetupAsWithCustomIO, it allows specifying additional environment variables to be used by the process. It allows the stdin, stdout, and stderr to be overridden.
+func (s *Subprocess) SetupAsWithEnvironmentWithCustomIO(ctx context.Context, loggers logs.Loggers, io ICommandIO, additionalEnv []string, messageOnStart string, messageOnSuccess, messageOnFailure string, as *commandUtils.CommandAsDifferentUser, cmd string, args ...string) (err error) {
+	return s.setup(ctx, loggers, io, additionalEnv, true, messageOnStart, messageOnSuccess, messageOnFailure, as, cmd, args...)
+}
+
+// Setup sets up a sub-process i.e. defines the command cmd and the messages on start, success and failure as well as the stdin, stdout, and stderr.
+func (s *Subprocess) setup(ctx context.Context, loggers logs.Loggers, io ICommandIO, env []string, withAdditionalMessages bool, messageOnStart, messageOnSuccess, messageOnFailure string, as *commandUtils.CommandAsDifferentUser, cmd string, args ...string) (err error) {
 	if s.IsOn() {
 		err = s.Stop()
 		if err != nil {
@@ -155,7 +175,11 @@ func (s *Subprocess) setup(ctx context.Context, loggers logs.Loggers, env []stri
 	defer s.mu.Unlock()
 	s.isRunning.Store(false)
 	s.processMonitoring = newSubprocessMonitoring(ctx)
-	s.command = newCommand(loggers, as, env, cmd, args...)
+	if io != nil {
+		s.command = newCommandWithCustomIO(loggers, io, as, env, cmd, args...)
+	} else {
+		s.command = newCommand(loggers, as, env, cmd, args...)
+	}
 	s.messaging = newSubprocessMessaging(loggers, withAdditionalMessages, messageOnSuccess, messageOnFailure, messageOnStart, s.command.GetPath())
 	s.reset()
 	return s.check()

utils/subprocess/executor.go

@@ -5,8 +5,10 @@
 package subprocess
 
 import (
+	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"os"
 	"os/exec"
 	"regexp"
@@ -27,6 +29,45 @@ import (
 	"github.com/ARM-software/golang-utils/utils/platform"
 )
 
+type testIO struct {
+	in  io.Reader
+	out *bytes.Buffer
+	err *bytes.Buffer
+}
+
+func newTestIO() *testIO {
+	return &testIO{
+		in:  strings.NewReader(""),
+		out: &bytes.Buffer{},
+		err: &bytes.Buffer{},
+	}
+}
+
+func (t *testIO) Register(context.Context) (io.Reader, io.Writer, io.Writer) {
+	return t.in, t.out, t.err
+}
+
+type execFunc func(ctx context.Context, l logs.Loggers, cmd string, args ...string) error
+
+func newDefaultExecutor(t *testing.T) execFunc {
+	t.Helper()
+	return func(ctx context.Context, l logs.Loggers, cmd string, args ...string) error {
+		return Execute(ctx, l, "", "", "", cmd, args...)
+	}
+}
+
+func newCustomIOExecutor(t *testing.T, customIO *testIO) execFunc {
+	t.Helper()
+	return func(ctx context.Context, l logs.Loggers, cmd string, args ...string) (err error) {
+		p := &Subprocess{}
+		err = p.SetupWithEnvironmentWithCustomIO(ctx, l, customIO, nil, "", "", "", cmd, args...)
+		if err != nil {
+			return
+		}
+		return p.Execute()
+	}
+}
+
 func TestExecuteEmptyLines(t *testing.T) {
 	t.Skip("would need to be reinstated when fixed")
 	defer goleak.VerifyNone(t)
@@ -265,52 +306,95 @@ func TestStartInterrupt(t *testing.T) {
 func TestExecute(t *testing.T) {
 	currentDir, err := os.Getwd()
 	require.NoError(t, err)
+
 	tests := []struct {
 		name       string
 		cmdWindows string
 		argWindows []string
 		cmdOther   string
 		argOther   []string
+		expectIO   bool
 	}{
 		{
 			name:       "ShortProcess",
 			cmdWindows: "cmd",
 			argWindows: []string{"/c", "dir", currentDir},
 			cmdOther:   "ls",
 			argOther:   []string{"-l", currentDir},
+			expectIO:   true,
 		},
 		{
 			name:       "LongProcess",
 			cmdWindows: "cmd",
 			argWindows: []string{"/c", fmt.Sprintf("ping -n 2 -w %v localhost > nul", time.Second.Milliseconds())}, // See https://stackoverflow.com/a/79268314/45375
 			cmdOther:   "sleep",
 			argOther:   []string{"1"},
+			expectIO:   false,
 		},
 	}
 
-	for i := range tests {
-		test := tests[i]
+	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			defer goleak.VerifyNone(t)
-			var loggers logs.Loggers = &logs.GenericLoggers{}
-			err := loggers.Check()
-			assert.Error(t, err)
 
-			err = Execute(context.Background(), loggers, "", "", "", "ls")
-			assert.Error(t, err)
+			customIO := newTestIO()
+			executors := []struct {
+				name string
+				run  execFunc
+				io   *testIO
+			}{
+				{"normal", newDefaultExecutor(t), nil},
+				{"with IO", newCustomIOExecutor(t, customIO), customIO},
+			}
 
-			loggers, err = logs.NewLogrLogger(logstest.NewTestLogger(t), "test")
-			require.NoError(t, err)
-			if platform.IsWindows() {
-				err = Execute(context.Background(), loggers, "", "", "", test.cmdWindows, test.argWindows...)
-			} else {
-				err = Execute(context.Background(), loggers, "", "", "", test.cmdOther, test.argOther...)
+			for _, executor := range executors {
+				t.Run(executor.name, func(t *testing.T) {
+					var loggers logs.Loggers = &logs.GenericLoggers{}
+					err := loggers.Check()
+					assert.Error(t, err)
+
+					err = executor.run(context.Background(), loggers, "ls")
+					assert.Error(t, err)
+
+					loggers, err = logs.NewLogrLogger(logstest.NewTestLogger(t), "test")
+					require.NoError(t, err)
+
+					if platform.IsWindows() {
+						err = executor.run(context.Background(), loggers, test.cmdWindows, test.argWindows...)
+					} else {
+						err = executor.run(context.Background(), loggers, test.cmdOther, test.argOther...)
+					}
+					require.NoError(t, err)
+
+					if executor.io != nil && test.expectIO {
+						assert.NotZero(t, executor.io.out.Len()+executor.io.err.Len()) // expect some output
+					}
+				})
 			}
-			require.NoError(t, err)
 		})
 	}
 }
 
+func TestExecuteWithCustomIO_Stderr(t *testing.T) {
+	if platform.IsWindows() {
+		t.Skip("Uses bash and redirection so can't run on Windows")
+	}
+	defer goleak.VerifyNone(t)
+
+	loggers, err := logs.NewLogrLogger(logstest.NewTestLogger(t), "test")
+	require.NoError(t, err)
+
+	customIO := newTestIO()
+	run := newCustomIOExecutor(t, customIO)
+
+	msg := "hello adrien"
+	err = run(context.Background(), loggers, "bash", "-c", fmt.Sprintf("echo %s 1>&2", msg))
+	require.NoError(t, err)
+
+	require.Empty(t, customIO.out.String()) // should be no stdout
+	require.Equal(t, fmt.Sprintln(msg), customIO.err.String())
+}
+
 func TestOutput(t *testing.T) {
 	loggers, err := logs.NewLogrLogger(logstest.NewTestLogger(t), "testOutput")
 	require.NoError(t, err)

utils/subprocess/executor_test.go

@@ -0,0 +1,66 @@
+package subprocess
+
+import (
+	"context"
+	"io"
+	"os"
+	"sync"
+
+	"github.com/ARM-software/golang-utils/utils/logs"
+)
+
+//go:generate go tool mockgen -destination=../mocks/mock_$GOPACKAGE.go -package=mocks github.com/ARM-software/golang-utils/utils/$GOPACKAGE ICommandIO
+
+// ICommandIO allows you to set the stdin, stdout, and stderr that will be used in a subprocess. A context can be injected for context aware readers and writers
+type ICommandIO interface {
+	// Register creates new readers and writers based on the constructor methods in the ICommandIO implementation. If the constructors are not specified then it will default to os.Stdin, os.Stdout, and os.Stderr
+	Register(context.Context) (in io.Reader, out, errs io.Writer)
+}
+
+type commandIO struct {
+	newInFunc    func(context.Context) io.Reader
+	newOutFunc   func(context.Context) io.Writer
+	newErrorFunc func(context.Context) io.Writer
+	mu           sync.Mutex
+}
+
+func NewIO(
+	newInFunc func(context.Context) io.Reader,
+	newOutFunc func(context.Context) io.Writer,
+	newErrorFunc func(context.Context) io.Writer,
+) ICommandIO {
+	return &commandIO{
+		mu:           sync.Mutex{},
+		newInFunc:    newInFunc,
+		newOutFunc:   newOutFunc,
+		newErrorFunc: newErrorFunc,
+	}
+}
+
+func NewIOFromLoggers(loggers logs.Loggers) ICommandIO {
+	return NewIO(
+		nil,
+		func(ctx context.Context) io.Writer { return newOutStreamer(ctx, loggers) },
+		func(ctx context.Context) io.Writer { return newErrLogStreamer(ctx, loggers) },
+	)
+}
+
+func NewDefaultIO() ICommandIO {
+	return NewIO(nil, nil, nil)
+}
+
+func (c *commandIO) Register(ctx context.Context) (in io.Reader, out, errs io.Writer) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	in, out, errs = os.Stdin, os.Stdout, os.Stderr
+	if c.newInFunc != nil {
+		in = c.newInFunc(ctx)
+	}
+	if c.newOutFunc != nil {
+		out = c.newOutFunc(ctx)
+	}
+	if c.newErrorFunc != nil {
+		errs = c.newErrorFunc(ctx)
+	}
+	return
+}