fix: Handle platform specific PTY polling for terminal copy (#279)

* fix: Handle platform specific PTY polling for terminal copy

* fix: Code review fixes

* fix: Code review fixes

internal/pty/output_copy.go

@@ -0,0 +1,40 @@
+package pty
+
+import (
+	"context"
+	"errors"
+	"io"
+)
+
+const outputCopyBufferSize = 32 * 1024
+
+// copyWithContext is a plain context-aware copy used as a fallback when the
+// source is not a *os.File (e.g. mocks in tests) and on platforms without a
+// poll-based reader. Cancellation is observed between reads.
+func copyWithContext(ctx context.Context, dst io.Writer, src io.Reader) error {
+	buf := make([]byte, outputCopyBufferSize)
+	for {
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
+		}
+
+		nr, err := src.Read(buf)
+		if nr > 0 {
+			nw, werr := dst.Write(buf[:nr])
+			if werr != nil {
+				return werr
+			}
+			if nw < nr {
+				return io.ErrShortWrite
+			}
+		}
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				return nil
+			}
+			return err
+		}
+	}
+}

internal/pty/output_copy_test.go

@@ -0,0 +1,79 @@
+//go:build !windows
+
+package pty
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/safedep/ptyx"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func spawnSh(t *testing.T, script string) ptyx.Session {
+	t.Helper()
+	sess, err := ptyx.Spawn(context.Background(), ptyx.SpawnOpts{
+		Prog: "/bin/sh",
+		Args: []string{"-c", script},
+		Cols: 80,
+		Rows: 24,
+		Env:  os.Environ(),
+	})
+	require.NoError(t, err)
+	return sess
+}
+
+func TestCopyPTYOutput_DataAndEOF(t *testing.T) {
+	sess := spawnSh(t, "printf 'hello-pty-output'")
+	defer func() { _ = sess.Close() }()
+
+	var buf bytes.Buffer
+	err := copyPTYOutput(context.Background(), &buf, sess.PtyReader())
+	require.NoError(t, err)
+	assert.Contains(t, buf.String(), "hello-pty-output")
+	_ = sess.Wait()
+}
+
+// TestCopyPTYOutput_CancelDoesNotHang guards the core invariant: a reader
+// blocked with no data available must return promptly on cancellation, never
+// relying on Close() to interrupt a blocked read (which does not work on macOS).
+func TestCopyPTYOutput_CancelDoesNotHang(t *testing.T) {
+	sess := spawnSh(t, "sleep 10")
+	defer func() { _ = sess.Close() }()
+	defer func() { _ = sess.Kill() }()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan error, 1)
+	go func() { done <- copyPTYOutput(ctx, io.Discard, sess.PtyReader()) }()
+
+	time.Sleep(200 * time.Millisecond)
+	cancel()
+
+	select {
+	case err := <-done:
+		assert.NoError(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("copyPTYOutput did not return after cancel: goroutine/thread leak")
+	}
+}
+
+// TestCopyPTYOutput_FallbackNonFile exercises the copyWithContext path used for
+// readers that are not *os.File (e.g. test doubles).
+func TestCopyPTYOutput_FallbackNonFile(t *testing.T) {
+	pr, pw := io.Pipe()
+	var buf bytes.Buffer
+	done := make(chan error, 1)
+	go func() { done <- copyPTYOutput(context.Background(), &buf, pr) }()
+
+	_, err := pw.Write([]byte("piped-data"))
+	require.NoError(t, err)
+	require.NoError(t, pw.Close())
+
+	require.NoError(t, <-done)
+	assert.Equal(t, "piped-data", buf.String())
+}

internal/pty/output_copy_unix.go

@@ -0,0 +1,98 @@
+//go:build !windows
+
+package pty
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"os"
+
+	"golang.org/x/sys/unix"
+)
+
+const ptyPollTimeoutMs = 100
+
+// copyPTYOutput copies child output from the PTY master to dst, driving the
+// read with a manual poll(2) loop instead of relying on Go's kqueue-based
+// netpoller.
+//
+// On some hardened / MDM-managed macOS hosts the runtime cannot register the
+// PTY master with the netpoller. os.OpenFile then leaves the master in
+// non-blocking mode without poller backing, so a plain io.Copy read returns
+// raw EAGAIN ("resource temporarily unavailable") on the very first read.
+// Forcing the master into blocking mode is not an option either: on macOS a
+// concurrent Close() does not interrupt a blocked PTY read(), leaking the
+// reader goroutine and its OS thread. poll(2) sidesteps both: it works without
+// the netpoller and the timeout lets us honor ctx cancellation without ever
+// depending on Close() to unblock a read.
+func copyPTYOutput(ctx context.Context, dst io.Writer, src io.Reader) error {
+	file, ok := src.(*os.File)
+	if !ok {
+		return copyWithContext(ctx, dst, src)
+	}
+
+	fd := int32(file.Fd())
+	pollFds := []unix.PollFd{{Fd: fd, Events: unix.POLLIN}}
+	buf := make([]byte, outputCopyBufferSize)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
+		}
+
+		n, err := unix.Poll(pollFds, ptyPollTimeoutMs)
+		if err != nil {
+			if errors.Is(err, unix.EINTR) {
+				continue
+			}
+			return fmt.Errorf("failed to poll pty master: %w", err)
+		}
+
+		if n == 0 {
+			continue // timeout: loop back to re-check ctx
+		}
+
+		revents := pollFds[0].Revents
+
+		// Drain readable data before acting on a hangup so that output written
+		// just before the child exited (POLLIN and POLLHUP can be reported
+		// together) is not lost.
+		if revents&unix.POLLIN != 0 {
+			nr, rerr := file.Read(buf)
+			if nr > 0 {
+				nw, werr := dst.Write(buf[:nr])
+				if werr != nil {
+					return werr
+				}
+				if nw < nr {
+					return io.ErrShortWrite
+				}
+			}
+			if rerr != nil {
+				if isReadEOF(rerr) {
+					return nil
+				}
+				if errors.Is(rerr, unix.EAGAIN) {
+					continue
+				}
+				return rerr
+			}
+			continue
+		}
+
+		if revents&(unix.POLLHUP|unix.POLLERR|unix.POLLNVAL) != 0 {
+			return nil
+		}
+	}
+}
+
+// isReadEOF reports whether a read error signals end of stream on the PTY
+// master. The master surfaces the child closing the slave as io.EOF on darwin
+// and as EIO on Linux.
+func isReadEOF(err error) bool {
+	return errors.Is(err, io.EOF) || errors.Is(err, unix.EIO)
+}

internal/pty/output_copy_windows.go

@@ -0,0 +1,15 @@
+//go:build windows
+
+package pty
+
+import (
+	"context"
+	"io"
+)
+
+// copyPTYOutput on Windows uses a plain context-aware copy. The conpty backend
+// is not affected by the kqueue netpoller issue that requires a poll(2) loop on
+// unix hosts.
+func copyPTYOutput(ctx context.Context, dst io.Writer, src io.Reader) error {
+	return copyWithContext(ctx, dst, src)
+}

internal/pty/session.go

@@ -21,6 +21,12 @@ type InteractiveSession interface {
 	// PtyReader returns the reader to receive output from the child process
 	PtyReader() io.Reader
 
+	// CopyOutputContext copies child output to dst until the child exits (EOF)
+	// or ctx is cancelled. On unix it drives the read with a poll(2) loop so it
+	// works even when the Go netpoller cannot manage the PTY master, and so it
+	// can be cancelled without leaking a goroutine blocked in read().
+	CopyOutputContext(ctx context.Context, dst io.Writer) error
+
 	// SetRawMode puts terminal in raw mode (for PTY passthrough)
 	SetRawMode() error
 
@@ -129,6 +135,10 @@ func NewSession(ctx context.Context, cfg SessionConfig) (InteractiveSession, err
 func (s *session) PtyWriter() io.Writer { return s.spawn.PtyWriter() }
 func (s *session) PtyReader() io.Reader { return s.spawn.PtyReader() }
 
+func (s *session) CopyOutputContext(ctx context.Context, dst io.Writer) error {
+	return copyPTYOutput(ctx, dst, s.spawn.PtyReader())
+}
+
 func (s *session) SetRawMode() error {
 	_, err := s.console.MakeRaw()
 	return err

internal/runner/execute.go

@@ -7,7 +7,7 @@ import (
 	"os"
 	"os/exec"
 	"strings"
-	"sync"
+	"time"
 
 	"github.com/safedep/dry/log"
 	"github.com/safedep/dry/usefulerror"
@@ -27,6 +27,10 @@ const (
 	ExecutionModeAuto
 )
 
+// outputDrainGrace bounds how long we wait for the PTY output reader to finish
+// after the child exits before forcing it to stop.
+const outputDrainGrace = 2 * time.Second
+
 type ExecuteOptions struct {
 	PackageManagerName string
 	DryRun             bool
@@ -167,12 +171,20 @@ func runPTY(
 		return fmt.Errorf("failed to create output router: %w", err)
 	}
 
-	var wg sync.WaitGroup
-	wg.Go(func() {
-		if _, err := io.Copy(outputRouter, sess.PtyReader()); err != nil {
+	// The output reader normally ends on its own when the PTY master reports
+	// EOF after the child exits. copyCtx lets us stop it otherwise: on parent
+	// cancellation (Ctrl+C) and on the drain-grace path below, which guards
+	// against a lingering descendant keeping the slave open (no EOF).
+	copyCtx, stopCopy := context.WithCancel(ctx)
+	defer stopCopy()
+
+	copyDone := make(chan struct{})
+	go func() {
+		defer close(copyDone)
+		if err := sess.CopyOutputContext(copyCtx, outputRouter); err != nil {
 			log.Errorf("failed to copy output: %v", err)
 		}
-	})
+	}()
 
 	inputRouter, err := pty.NewInputRouter(sess.PtyWriter())
 	if err != nil {
@@ -217,7 +229,16 @@ func runPTY(
 	}
 
 	sessionError := sess.Wait()
-	wg.Wait()
+
+	// Child has exited. Let the reader drain to EOF, but bound the wait so a
+	// lingering descendant holding the slave open cannot block teardown.
+	select {
+	case <-copyDone:
+	case <-time.After(outputDrainGrace):
+		log.Debugf("output drain grace exceeded, stopping pty reader")
+		stopCopy()
+		<-copyDone
+	}
 
 	if sessionError != nil {
 		return wrapCommandExecutionError(sessionError, result)