Clean up exec processes on stream failures

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
+	// Ensure the PTY is closed if sending the Started response fails
-	err = sendStartSuccess(stream)
-	if err != nil {
-		return err
-	}
-
 	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

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
+}