guest-agent: stream non-TTY exec output as it is produced

lib/system/guest_agent/exec.go

@@ -38,8 +38,8 @@ func (s *guestServer) Exec(stream pb.GuestService_ExecServer) error {
 	log.Printf("[guest-agent] exec: command=%v tty=%v cwd=%s timeout=%d",
 		command, start.Tty, start.Cwd, start.TimeoutSeconds)
 
-	// Create context with timeout if specified
-	ctx := context.Background()
+	// Tie the command to the stream so a client disconnect kills it.
+	ctx := stream.Context()
 	if start.TimeoutSeconds > 0 {
 		var cancel context.CancelFunc
 		ctx, cancel = context.WithTimeout(ctx, time.Duration(start.TimeoutSeconds)*time.Second)
@@ -52,13 +52,19 @@ func (s *guestServer) Exec(stream pb.GuestService_ExecServer) error {
 	return s.executeNoTTY(ctx, stream, start)
 }
 
+// execDrainTimeout bounds how long output may keep flowing after the command exits.
+var execDrainTimeout = 5 * time.Second
+
 // executeNoTTY executes command without TTY
 func (s *guestServer) executeNoTTY(ctx context.Context, stream pb.GuestService_ExecServer, start *pb.ExecStart) error {
 	// Run command directly - guest-agent is already running in container namespace
 	if len(start.Command) == 0 {
 		return fmt.Errorf("empty command")
 	}
 
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	cmd := exec.CommandContext(ctx, start.Command[0], start.Command[1:]...)
 
 	// Set up environment (no TTY defaults for non-TTY mode)
@@ -70,19 +76,56 @@ func (s *guestServer) executeNoTTY(ctx context.Context, stream pb.GuestService_E
 	}
 
 	stdin, _ := cmd.StdinPipe()
-	stdout, _ := cmd.StdoutPipe()
-	stderr, _ := cmd.StderrPipe()
 
-	if err := cmd.Start(); err != nil {
+	// Own the output pipes so cmd.Wait can reap the process without closing them mid-read.
+	stdoutR, stdoutW, err := os.Pipe()
+	if err != nil {
+		return fmt.Errorf("stdout pipe: %w", err)
+	}
+	defer stdoutR.Close()
+	stderrR, stderrW, err := os.Pipe()
+	if err != nil {
+		stdoutW.Close()
+		return fmt.Errorf("stderr pipe: %w", err)
+	}
+	defer stderrR.Close()
+	cmd.Stdout = stdoutW
+	cmd.Stderr = stderrW
+
+	err = cmd.Start()
+	stdoutW.Close()
+	stderrW.Close()
+	if err != nil {
 		return fmt.Errorf("start command: %w", err)
 	}
 
 	// Mutex to protect concurrent stream.Send calls (gRPC streams are not thread-safe)
 	var sendMu sync.Mutex
+	var sendErr error
 
-	// Use WaitGroup to ensure all output is read before sending
+	// Stream output as it is produced so long-lived commands (relays, tails)
+	// deliver bytes before exit.
 	var wg sync.WaitGroup
-	var stdoutData, stderrData []byte
+	pump := func(r io.Reader, wrap func([]byte) *pb.ExecResponse) {
+		defer wg.Done()
+		buf := make([]byte, 32*1024)
+		for {
+			n, err := r.Read(buf)
+			if n > 0 {
+				sendMu.Lock()
+				if sendErr == nil {
+					if sendErr = stream.Send(wrap(append([]byte(nil), buf[:n]...))); sendErr != nil {
+						// Client is gone: kill the command, keep draining so it can't block on a full pipe.
+						cancel()
+					}
+				}
+				sendMu.Unlock()
+			}
+			if err != nil {
+				return
+			}
+		}
+	}
 
 	// Handle stdin in background
 	go func() {
@@ -98,50 +141,31 @@ func (s *guestServer) executeNoTTY(ctx context.Context, stream pb.GuestService_E
 		}
 	}()
 
-	// Read all stdout/stderr BEFORE calling Wait() - Wait() closes the pipes!
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		data, _ := io.ReadAll(stdout)
-		stdoutData = data
-	}()
-
-	wg.Add(1)
+	wg.Add(2)
+	go pump(stdoutR, func(b []byte) *pb.ExecResponse {
+		return &pb.ExecResponse{Response: &pb.ExecResponse_Stdout{Stdout: b}}
+	})
+	go pump(stderrR, func(b []byte) *pb.ExecResponse {
+		return &pb.ExecResponse{Response: &pb.ExecResponse_Stderr{Stderr: b}}
+	})
+	pumpsDone := make(chan struct{})
 	go func() {
-		defer wg.Done()
-		data, _ := io.ReadAll(stderr)
-		stderrData = data
+		wg.Wait()
+		close(pumpsDone)
 	}()
 
-	// Wait for all reads to complete FIRST (before Wait closes pipes)
-	wg.Wait()
-
-	// Now safe to call Wait - pipes are fully drained
+	// Reap first: a stalled client backpressures the command but never leaks it.
 	waitErr := cmd.Wait()
 
-	// Now stream output in chunks (streaming compatible)
-	const chunkSize = 32 * 1024
-	for i := 0; i < len(stdoutData); i += chunkSize {
-		end := i + chunkSize
-		if end > len(stdoutData) {
-			end = len(stdoutData)
-		}
-		sendMu.Lock()
-		stream.Send(&pb.ExecResponse{
-			Response: &pb.ExecResponse_Stdout{Stdout: stdoutData[i:end]},
-		})
-		sendMu.Unlock()
+	select {
+	case <-pumpsDone:
+	case <-time.After(execDrainTimeout):
+		// Returning cancels the stream, which unblocks a Send stuck on flow control.
+		log.Printf("[guest-agent] exec output not drained %s after exit; abandoning stream", execDrainTimeout)
+		return fmt.Errorf("exec output not drained within %s of command exit", execDrainTimeout)
 	}
-	for i := 0; i < len(stderrData); i += chunkSize {
-		end := i + chunkSize
-		if end > len(stderrData) {
-			end = len(stderrData)
-		}
-		sendMu.Lock()
-		stream.Send(&pb.ExecResponse{
-			Response: &pb.ExecResponse_Stderr{Stderr: stderrData[i:end]},
-		})
-		sendMu.Unlock()
+	if sendErr != nil {
+		return fmt.Errorf("send output: %w", sendErr)
 	}
 
 	exitCode := int32(0)

lib/system/guest_agent/exec_stream_test.go

@@ -0,0 +1,176 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"io"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"sync"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+
+	pb "github.com/kernel/hypeman/lib/guest"
+)
+
+// recordingExecStream is a fake Exec server stream that timestamps stdout chunks.
+type recordingExecStream struct {
+	grpc.ServerStream
+	ctx    context.Context
+	mu     sync.Mutex
+	stdout []timedChunk
+}
+
+type timedChunk struct {
+	at   time.Time
+	data string
+}
+
+func (s *recordingExecStream) Context() context.Context { return s.ctx }
+func (s *recordingExecStream) Recv() (*pb.ExecRequest, error) {
+	return nil, io.EOF
+}
+
+func (s *recordingExecStream) Send(resp *pb.ExecResponse) error {
+	if out := resp.GetStdout(); out != nil {
+		s.mu.Lock()
+		s.stdout = append(s.stdout, timedChunk{at: time.Now(), data: string(out)})
+		s.mu.Unlock()
+	}
+	return nil
+}
+
+func TestExecuteNoTTYStreamsOutputBeforeExit(t *testing.T) {
+	stream := &recordingExecStream{ctx: t.Context()}
+	start := time.Now()
+	err := (&guestServer{}).executeNoTTY(t.Context(), stream, &pb.ExecStart{
+		Command: []string{"sh", "-c", "echo first; sleep 2; echo second"},
+	})
+	require.NoError(t, err)
+	exitedAt := time.Since(start)
+
+	require.NotEmpty(t, stream.stdout)
+	require.True(t, strings.HasPrefix(stream.stdout[0].data, "first"))
+	firstAt := stream.stdout[0].at.Sub(start)
+	require.Less(t, firstAt, exitedAt-time.Second,
+		"first chunk arrived at %s but command exited at %s; output is being buffered", firstAt, exitedAt)
+
+	var all strings.Builder
+	for _, c := range stream.stdout {
+		all.WriteString(c.data)
+	}
+	require.Equal(t, "first\nsecond\n", all.String())
+}
+
+// funcExecStream is a fake Exec server stream with pluggable Recv and Send.
+type funcExecStream struct {
+	grpc.ServerStream
+	ctx  context.Context
+	recv func() (*pb.ExecRequest, error)
+	send func(*pb.ExecResponse) error
+}
+
+func (s *funcExecStream) Context() context.Context       { return s.ctx }
+func (s *funcExecStream) Recv() (*pb.ExecRequest, error) { return s.recv() }
+func (s *funcExecStream) Send(r *pb.ExecResponse) error  { return s.send(r) }
+
+// runWithDeadline fails the test if fn has not returned within d.
+func runWithDeadline(t *testing.T, d time.Duration, fn func() error) error {
+	t.Helper()
+	done := make(chan error, 1)
+	go func() { done <- fn() }()
+	select {
+	case err := <-done:
+		return err
+	case <-time.After(d):
+		t.Fatalf("exec handler did not return within %s", d)
+		return nil
+	}
+}
+
+// requireReaped asserts the pid written to pidFile no longer exists (not even as a zombie).
+func requireReaped(t *testing.T, pidFile string) {
+	t.Helper()
+	raw, err := os.ReadFile(pidFile)
+	require.NoError(t, err)
+	pid, err := strconv.Atoi(strings.TrimSpace(string(raw)))
+	require.NoError(t, err)
+	require.ErrorIs(t, syscall.Kill(pid, 0), syscall.ESRCH, "pid %d was not reaped", pid)
+}
+
+func TestExecuteNoTTYReapsWhenClientStalls(t *testing.T) {
+	prev := execDrainTimeout
+	execDrainTimeout = 500 * time.Millisecond
+	t.Cleanup(func() { execDrainTimeout = prev })
+
+	// Send blocks like a gRPC stream with an exhausted flow-control window.
+	streamCtx, streamCancel := context.WithCancel(t.Context())
+	t.Cleanup(streamCancel)
+	stream := &funcExecStream{
+		ctx:  streamCtx,
+		recv: func() (*pb.ExecRequest, error) { <-streamCtx.Done(); return nil, io.EOF },
+		send: func(*pb.ExecResponse) error { <-streamCtx.Done(); return streamCtx.Err() },
+	}
+
+	pidFile := filepath.Join(t.TempDir(), "pid")
+	ctx, cancel := context.WithTimeout(t.Context(), time.Second)
+	defer cancel()
+	err := runWithDeadline(t, 10*time.Second, func() error {
+		return (&guestServer{}).executeNoTTY(ctx, stream, &pb.ExecStart{
+			Command: []string{"sh", "-c", `echo $$ > "$1"; exec yes`, "sh", pidFile},
+		})
+	})
+	require.ErrorContains(t, err, "not drained")
+	requireReaped(t, pidFile)
+}
+
+func TestExecuteNoTTYKillsCommandWhenSendFails(t *testing.T) {
+	sendErr := errors.New("client gone")
+	stream := &funcExecStream{
+		ctx:  t.Context(),
+		recv: func() (*pb.ExecRequest, error) { return nil, io.EOF },
+		send: func(*pb.ExecResponse) error { return sendErr },
+	}
+
+	pidFile := filepath.Join(t.TempDir(), "pid")
+	err := runWithDeadline(t, 10*time.Second, func() error {
+		return (&guestServer{}).executeNoTTY(t.Context(), stream, &pb.ExecStart{
+			Command: []string{"sh", "-c", `echo $$ > "$1"; echo started; exec sleep 30`, "sh", pidFile},
+		})
+	})
+	require.ErrorIs(t, err, sendErr)
+	requireReaped(t, pidFile)
+}
+
+func TestExecKillsCommandWhenClientDisconnects(t *testing.T) {
+	streamCtx, disconnect := context.WithCancel(t.Context())
+	defer disconnect()
+	pidFile := filepath.Join(t.TempDir(), "pid")
+	started := false
+	stream := &funcExecStream{
+		ctx: streamCtx,
+		recv: func() (*pb.ExecRequest, error) {
+			if !started {
+				started = true
+				return &pb.ExecRequest{Request: &pb.ExecRequest_Start{Start: &pb.ExecStart{
+					Command: []string{"sh", "-c", `echo $$ > "$1"; exec sleep 30`, "sh", pidFile},
+				}}}, nil
+			}
+			<-streamCtx.Done()
+			return nil, streamCtx.Err()
+		},
+		send: func(*pb.ExecResponse) error { return streamCtx.Err() },
+	}
+
+	time.AfterFunc(500*time.Millisecond, disconnect)
+	_ = runWithDeadline(t, 10*time.Second, func() error {
+		return (&guestServer{}).Exec(stream)
+	})
+	requireReaped(t, pidFile)
+}