internal/exec: Fix stdio pipe problems

exec today has two problems with how it handles stdio pipes:

- When Wait completes, closeStdio() is called.
  This closes the parent-side stdio pipes for receiving IO from the
  process. This is a problem because once the process has completed, we
  still need to be able to receive any final output. Today data from the
  process could be lost because of this.
- The parent's handles to the child-side stdio pipes are not closed
  after starting the process. Leaving duplicates of these handles in the
  parent process means that the other ends of the pipes are never closed
  when the process exits.

This commit makes the following changes:

- The parent's handles to the child-side stdio pipes are now closed
  after the child is started. This is necessary so that once the child
  exits, the parent-side pipes will return EOF once the remaining output
  drains.
- When Wait completes, the parent-side stdio pipes are not closed. The
  responsibility for this is now left to the client of the exec package.
  Currently the only user of exec is jobcontainers.JobProcess, which
  closes handles these when Close is called.

Additionally, the ProcThreadAttributeList is now allocated and used only
in Start. Previously it was saved on the Exec object, even though it was
not needed elsewhere. This makes the code cleaner, simplifies the Wait
logic, and eliminates the chance of leaking memory if an Exec object
is GC'd without being Wait'd.

Signed-off-by: Kevin Parsons <kevpar@microsoft.com>

internal/exec/exec.go

@@ -15,8 +15,7 @@ import (
 )
 
 var (
-	errProcNotStarted  = errors.New("process has not started yet")
-	errProcNotFinished = errors.New("process has not finished yet")
+	errProcNotStarted = errors.New("process has not started yet")
 )
 
 // Exec is an object that represents an external process. A user should NOT initialize one manually and instead should
@@ -44,7 +43,6 @@ type Exec struct {
 	// stdioPipesProcSide are the stdio pipes that will be passed into the process. These should not be interacted with at all
 	// and aren't exposed in any way to a user of Exec.
 	stdioPipesProcSide [3]*os.File
-	attrList           *windows.ProcThreadAttributeListContainer
 	*execConfig
 }
 
@@ -122,10 +120,11 @@ func (e *Exec) Start() error {
 	// 2. Pseudo console setup if one was requested.
 	// 3. Inherit only stdio handles if ones were requested.
 	// Therefore we need a list of size 3.
-	e.attrList, err = windows.NewProcThreadAttributeList(3)
+	attrList, err := windows.NewProcThreadAttributeList(3)
 	if err != nil {
 		return fmt.Errorf("failed to initialize process thread attribute list: %w", err)
 	}
+	defer attrList.Delete()
 
 	// Need to know whether the process needs to inherit stdio handles. The below setup is so that we only inherit the
 	// stdio pipes and nothing else into the new process.
@@ -139,7 +138,7 @@ func (e *Exec) Start() error {
 		}
 
 		// Set up the process to only inherit stdio handles and nothing else.
-		err := e.attrList.Update(
+		err := attrList.Update(
 			windows.PROC_THREAD_ATTRIBUTE_HANDLE_LIST,
 			unsafe.Pointer(&handles[0]),
 			uintptr(len(handles))*unsafe.Sizeof(handles[0]),
@@ -161,13 +160,13 @@ func (e *Exec) Start() error {
 	}
 
 	if e.job != nil {
-		if err := e.job.UpdateProcThreadAttribute(e.attrList); err != nil {
+		if err := e.job.UpdateProcThreadAttribute(attrList); err != nil {
 			return err
 		}
 	}
 
 	if e.cpty != nil {
-		if err := e.cpty.UpdateProcThreadAttribute(e.attrList); err != nil {
+		if err := e.cpty.UpdateProcThreadAttribute(attrList); err != nil {
 			return err
 		}
 	}
@@ -176,7 +175,7 @@ func (e *Exec) Start() error {
 	pSec := &windows.SecurityAttributes{Length: uint32(unsafe.Sizeof(zeroSec)), InheritHandle: 1}
 	tSec := &windows.SecurityAttributes{Length: uint32(unsafe.Sizeof(zeroSec)), InheritHandle: 1}
 
-	siEx.ProcThreadAttributeList = e.attrList.List() //nolint:govet // unusedwrite: ProcThreadAttributeList will be read in syscall
+	siEx.ProcThreadAttributeList = attrList.List() //nolint:govet // unusedwrite: ProcThreadAttributeList will be read in syscall
 	siEx.Cb = uint32(unsafe.Sizeof(*siEx))
 	if e.execConfig.token != 0 {
 		err = windows.CreateProcessAsUser(
@@ -214,6 +213,17 @@ func (e *Exec) Start() error {
 		_ = windows.CloseHandle(windows.Handle(pi.Thread))
 	}()
 
+	// Now that the process has started, we need to close our copies of its stdio handles.
+	// If we keep these around, then our ends of the pipes won't ever reach "EOF" after the process exits.
+	for i := range e.stdioPipesProcSide {
+		if e.stdioPipesProcSide[i] != nil {
+			if err := e.stdioPipesProcSide[i].Close(); err != nil {
+				return fmt.Errorf("close process's stdio handle %d: %w", i, err)
+			}
+			e.stdioPipesProcSide[i] = nil
+		}
+	}
+
 	// Grab an *os.Process to avoid reinventing the wheel here. The stdlib has great logic around waiting, exit code status/cleanup after a
 	// process has been launched.
 	e.process, err = os.FindProcess(int(pi.ProcessId))
@@ -236,16 +246,6 @@ func (e *Exec) Run() error {
 	return e.Wait()
 }
 
-// Close will release resources tied to the process (stdio etc.)
-func (e *Exec) close() error {
-	if e.procState == nil {
-		return errProcNotFinished
-	}
-	e.attrList.Delete()
-	e.closeStdio()
-	return nil
-}
-
 // Pid returns the pid of the running process. If the process isn't running, this will return -1.
 func (e *Exec) Pid() int {
 	if e.process == nil {
@@ -284,7 +284,7 @@ func (e *Exec) Wait() (err error) {
 	if err != nil {
 		return err
 	}
-	return e.close()
+	return nil
 }
 
 // Kill will forcefully kill the process.
@@ -386,21 +386,6 @@ func (e *Exec) setupStdio() error {
 	return nil
 }
 
-func (e *Exec) closeStdio() {
-	for i, file := range e.stdioPipesOurSide {
-		if file != nil {
-			file.Close()
-		}
-		e.stdioPipesOurSide[i] = nil
-	}
-	for i, file := range e.stdioPipesProcSide {
-		if file != nil {
-			file.Close()
-		}
-		e.stdioPipesProcSide[i] = nil
-	}
-}
-
 //
 // Below are a bunch of helpers for working with Windows' CreateProcess family of functions. These are mostly exact copies of the same utilities
 // found in the go stdlib.