From 5f1a89a5a43c29f513c84e3854e7ab79420f550b Mon Sep 17 00:00:00 2001 From: Nikolay Edigaryev Date: Mon, 10 Aug 2026 19:34:09 +0100 Subject: [PATCH] Clean up exec processes on stream failures --- internal/rpc/exec.go | 60 +++++++++++++++++++--------- internal/rpc/exec_test.go | 83 +++++++++++++++++++++++++++++++++++++-- 2 files changed, 120 insertions(+), 23 deletions(-) diff --git a/internal/rpc/exec.go b/internal/rpc/exec.go index 78fcb62..31ea673 100644 --- a/internal/rpc/exec.go +++ b/internal/rpc/exec.go @@ -71,18 +71,17 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse]) return sendStartFailure(stream) } + // Release ownership before sending responses so failures do not leak the process handle + if err := cmd.Process.Release(); err != nil { + return err + } + // Explicitly notify the client that the process was started err = sendStartSuccess(stream) if err != nil { return err } - if cmd.Process != nil { - if err := cmd.Process.Release(); err != nil { - return err - } - } - if err := stream.Send(&ExecResponse{ Type: &ExecResponse_Exit_{ Exit: &ExecResponse_Exit{ @@ -147,25 +146,39 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse]) return sendStartFailure(stream) } - // Explicitly notify the client that the process was started - err = sendStartSuccess(stream) - if err != nil { - return err - } - + // Ensure the PTY is closed if sending the Started response fails if ptmx != nil { defer ptmx.Close() } + // Explicitly notify the client that the process was started + err = sendStartSuccess(stream) + if err != nil { + // Output readers have not started yet, so cancel and reap directly + _ = cmd.Cancel() + _ = cmd.Wait() + + return err + } + // Handle standard input and terminal resize events from the client fromClientErrCh := make(chan error, 1) + reportClientError := func(err error) { + fromClientErrCh <- err + _ = cmd.Cancel() + } go func() { for { request, err := stream.Recv() if err != nil { + // Allow the client to close its sending side while continuing to receive responses + if errors.Is(err, io.EOF) { + return + } + if !errors.Is(err, context.Canceled) { - fromClientErrCh <- err + reportClientError(err) } return @@ -192,7 +205,7 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse]) } else { // Close the standard input if err := stdin.Close(); err != nil { - fromClientErrCh <- err + reportClientError(err) return } @@ -202,7 +215,7 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse]) } if _, err := stdin.Write(dataToWrite); err != nil { - fromClientErrCh <- err + reportClientError(err) return } @@ -217,7 +230,7 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse]) Rows: uint16(typedAction.TerminalResize.GetRows()), Cols: uint16(typedAction.TerminalResize.GetCols()), }); err != nil { - fromClientErrCh <- err + reportClientError(err) return } @@ -230,13 +243,13 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse]) case ExecRequest_SIGNAL_SIGKILL: signal = syscall.SIGKILL default: - fromClientErrCh <- fmt.Errorf("unsupported exec signal %q", typedAction.Signal.String()) + reportClientError(fmt.Errorf("unsupported exec signal %q", typedAction.Signal.String())) return } if err := signalProcessGroup(cmd.Process, signal); err != nil && !errors.Is(err, os.ErrProcessDone) { - fromClientErrCh <- err + reportClientError(err) return } @@ -323,9 +336,18 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse]) } // Wait for the command to finish + err = cmd.Wait() + + // Prefer a client error over the command exit result + select { + case err := <-fromClientErrCh: + return err + default: + } + exitCode := 0 - if err := cmd.Wait(); err != nil { + if err != nil { var exitError *exec.ExitError if !errors.As(err, &exitError) { return err diff --git a/internal/rpc/exec_test.go b/internal/rpc/exec_test.go index c57fe3c..9f49e6a 100644 --- a/internal/rpc/exec_test.go +++ b/internal/rpc/exec_test.go @@ -5,7 +5,11 @@ package rpc import ( "context" + "errors" "io" + "os" + "path/filepath" + "strconv" "syscall" "testing" "time" @@ -14,7 +18,10 @@ import ( "google.golang.org/grpc" ) -const execTestTimeout = 5 * time.Second +const ( + execTestShell = "/bin/sh" + execTestTimeout = 5 * time.Second +) type execTestStream struct { grpc.ServerStream @@ -22,6 +29,7 @@ type execTestStream struct { ctx context.Context requests chan *ExecRequest responses chan *ExecResponse + sendHook func(*ExecResponse) error } var _ grpc.BidiStreamingServer[ExecRequest, ExecResponse] = (*execTestStream)(nil) @@ -35,6 +43,12 @@ func newExecTestStream(ctx context.Context) *execTestStream { } func (stream *execTestStream) Send(response *ExecResponse) error { + if stream.sendHook != nil { + if err := stream.sendHook(response); err != nil { + return err + } + } + select { case stream.responses <- response: return nil @@ -59,7 +73,7 @@ func (stream *execTestStream) Context() context.Context { return stream.ctx } func TestExecSendsStartedBeforeOutputAndExit(t *testing.T) { stream, result := startExecTest(t, &ExecRequest_Command{ - Name: "/bin/sh", + Name: execTestShell, Args: []string{"-c", "printf hello"}, }) @@ -97,7 +111,7 @@ func TestExecReportsStartFailureBeforeStarted(t *testing.T) { { name: "missing workdir", command: &ExecRequest_Command{ - Name: "/bin/sh", + Name: execTestShell, Workdir: "/definitely/missing/tart-guest-agent-test-workdir", }, }, @@ -119,6 +133,7 @@ func TestExecSignalsProcess(t *testing.T) { name string signal ExecRequest_Signal code int32 + err string }{ { name: "SIGTERM", @@ -130,6 +145,11 @@ func TestExecSignalsProcess(t *testing.T) { signal: ExecRequest_SIGNAL_SIGKILL, code: int32(signalExitCodeOffset + syscall.SIGKILL), }, + { + name: "unsupported", + signal: ExecRequest_SIGNAL_UNSPECIFIED, + err: `unsupported exec signal "SIGNAL_UNSPECIFIED"`, + }, } for _, test := range tests { @@ -144,6 +164,11 @@ func TestExecSignalsProcess(t *testing.T) { Type: &ExecRequest_Signal_{Signal: test.signal}, } + if test.err != "" { + require.EqualError(t, receiveExecResult(t, result), test.err) + return + } + response := receiveExecResponse(t, stream) require.NotNil(t, response.GetExit()) require.Equal(t, test.code, response.GetExit().GetCode()) @@ -154,7 +179,7 @@ func TestExecSignalsProcess(t *testing.T) { func TestExecSignalsProcessGroup(t *testing.T) { stream, result := startExecTest(t, &ExecRequest_Command{ - Name: "/bin/sh", + Name: execTestShell, Args: []string{"-c", "sleep 30 & printf ready; wait"}, }) require.NotNil(t, receiveExecResponse(t, stream).GetStarted()) @@ -169,15 +194,48 @@ func TestExecSignalsProcessGroup(t *testing.T) { require.NoError(t, receiveExecResult(t, result)) } +func TestExecReapsProcessWhenStartedCannotBeSent(t *testing.T) { + pidPath := filepath.Join(t.TempDir(), "pid") + sendErr := errors.New("failed to send Started") + var processPID int + + _, result := startExecTest(t, &ExecRequest_Command{ + Name: execTestShell, + Args: []string{"-c", `printf %d "$$" > "$PID_FILE"; exec sleep 30`}, + Env: map[string]string{"PID_FILE": pidPath}, + }, func(stream *execTestStream) { + stream.sendHook = func(response *ExecResponse) error { + if response.GetStarted() == nil { + return nil + } + + var err error + processPID, err = waitForExecTestPID(pidPath) + if err != nil { + return err + } + + return sendErr + } + }) + + require.ErrorIs(t, receiveExecResult(t, result), sendErr) + require.ErrorIs(t, syscall.Kill(processPID, 0), syscall.ESRCH) +} + func startExecTest( t *testing.T, command *ExecRequest_Command, + configure ...func(*execTestStream), ) (*execTestStream, <-chan error) { t.Helper() ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) stream := newExecTestStream(ctx) + for _, configureStream := range configure { + configureStream(stream) + } result := make(chan error, 1) go func() { result <- (&RPC{}).Exec(stream) @@ -211,3 +269,20 @@ func receiveExecResult(t *testing.T, result <-chan error) error { return nil } } + +func waitForExecTestPID(path string) (int, error) { + deadline := time.Now().Add(execTestTimeout) + for time.Now().Before(deadline) { + data, err := os.ReadFile(path) + if err == nil { + return strconv.Atoi(string(data)) + } + if !errors.Is(err, os.ErrNotExist) { + return 0, err + } + + time.Sleep(10 * time.Millisecond) + } + + return 0, context.DeadlineExceeded +}