public_git
/
tart-guest-agent
mirror of https://github.com/cirruslabs/tart-guest-agent.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Clean up exec processes on stream failures

This commit is contained in:
Nikolay Edigaryev 2026-08-10 19:34:09 +01:00
parent 1eb9e26b2b
commit 5f1a89a5a4
2 changed files with 120 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

60
internal/rpc/exec.go
View File

@ -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

83
internal/rpc/exec_test.go
View File

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