diff --git a/changes/20250911113619.feature b/changes/20250911113619.feature new file mode 100644 index 0000000000..bd3da6b89d --- /dev/null +++ b/changes/20250911113619.feature @@ -0,0 +1 @@ +:sparkles: `subprocess` Add support for overriding the stdin/stdout/stderr of a subprocess diff --git a/utils/mocks/mock_subprocess.go b/utils/mocks/mock_subprocess.go new file mode 100644 index 0000000000..e03e5d7ac3 --- /dev/null +++ b/utils/mocks/mock_subprocess.go @@ -0,0 +1,58 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: github.com/ARM-software/golang-utils/utils/subprocess (interfaces: ICommandIO) +// +// Generated by this command: +// +// mockgen -destination=../mocks/mock_subprocess.go -package=mocks github.com/ARM-software/golang-utils/utils/subprocess ICommandIO +// + +// Package mocks is a generated GoMock package. +package mocks + +import ( + context "context" + io "io" + reflect "reflect" + + gomock "go.uber.org/mock/gomock" +) + +// MockICommandIO is a mock of ICommandIO interface. +type MockICommandIO struct { + ctrl *gomock.Controller + recorder *MockICommandIOMockRecorder + isgomock struct{} +} + +// MockICommandIOMockRecorder is the mock recorder for MockICommandIO. +type MockICommandIOMockRecorder struct { + mock *MockICommandIO +} + +// NewMockICommandIO creates a new mock instance. +func NewMockICommandIO(ctrl *gomock.Controller) *MockICommandIO { + mock := &MockICommandIO{ctrl: ctrl} + mock.recorder = &MockICommandIOMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockICommandIO) EXPECT() *MockICommandIOMockRecorder { + return m.recorder +} + +// Register mocks base method. +func (m *MockICommandIO) Register(arg0 context.Context) (io.Reader, io.Writer, io.Writer) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Register", arg0) + ret0, _ := ret[0].(io.Reader) + ret1, _ := ret[1].(io.Writer) + ret2, _ := ret[2].(io.Writer) + return ret0, ret1, ret2 +} + +// Register indicates an expected call of Register. +func (mr *MockICommandIOMockRecorder) Register(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Register", reflect.TypeOf((*MockICommandIO)(nil).Register), arg0) +} diff --git a/utils/subprocess/command_wrapper.go b/utils/subprocess/command_wrapper.go index c935bdb139..5d16ad347e 100644 --- a/utils/subprocess/command_wrapper.go +++ b/utils/subprocess/command_wrapper.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 } diff --git a/utils/subprocess/executor.go b/utils/subprocess/executor.go index 55b549d0a2..b32387ad14 100644 --- a/utils/subprocess/executor.go +++ b/utils/subprocess/executor.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() diff --git a/utils/subprocess/executor_test.go b/utils/subprocess/executor_test.go index fec2e45837..cee016c706 100644 --- a/utils/subprocess/executor_test.go +++ b/utils/subprocess/executor_test.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,12 +306,14 @@ 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", @@ -278,6 +321,7 @@ func TestExecute(t *testing.T) { argWindows: []string{"/c", "dir", currentDir}, cmdOther: "ls", argOther: []string{"-l", currentDir}, + expectIO: true, }, { name: "LongProcess", @@ -285,32 +329,72 @@ func TestExecute(t *testing.T) { 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) diff --git a/utils/subprocess/io.go b/utils/subprocess/io.go new file mode 100644 index 0000000000..d07070abf9 --- /dev/null +++ b/utils/subprocess/io.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 +} diff --git a/utils/subprocess/io_test.go b/utils/subprocess/io_test.go new file mode 100644 index 0000000000..ece9d7cde3 --- /dev/null +++ b/utils/subprocess/io_test.go @@ -0,0 +1,17 @@ +package subprocess + +import ( + "context" + "os" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestDefaultIO(t *testing.T) { + io := NewDefaultIO() + in, out, errs := io.Register(context.Background()) + assert.Equal(t, os.Stdin, in) + assert.Equal(t, os.Stdout, out) + assert.Equal(t, os.Stderr, errs) +}