fix(pty): close output writer before reader on Windows to unblock close (#8299)

mafredri · web-flow · commit 88c35d3f04db · 2023-07-05T15:25:07.000+03:00
diff --git a/pty/pty_windows.go b/pty/pty_windows.go
@@ -65,10 +65,13 @@ func newPty(opt ...Option) (*ptyWindows, error) {
 		0,
 		uintptr(unsafe.Pointer(&pty.console)),
 	)
-	if int32(ret) < 0 {
+	// CreatePseudoConsole returns S_OK on success, as per:
+	// https://learn.microsoft.com/en-us/windows/console/createpseudoconsole
+	if windows.Handle(ret) != windows.S_OK {
 		_ = pty.Close()
 		return nil, xerrors.Errorf("create pseudo console (%d): %w", int32(ret), err)
 	}
+
 	return pty, nil
 }
 
@@ -134,39 +137,65 @@ func (p *ptyWindows) Resize(height uint16, width uint16) error {
 		Y: int16(height),
 		X: int16(width),
 	})))))
-	if ret != 0 {
+	if windows.Handle(ret) != windows.S_OK {
 		return err
 	}
 	return nil
 }
 
-func (p *ptyWindows) Close() error {
-	p.closeMutex.Lock()
-	defer p.closeMutex.Unlock()
-	if p.closed {
-		return nil
-	}
-	p.closed = true
-
+// closeConsoleNoLock closes the console handle, and sets it to
+// windows.InvalidHandle. It must be called with p.closeMutex held.
+func (p *ptyWindows) closeConsoleNoLock() error {
 	// if we are running a command in the PTY, the corresponding *windowsProcess
 	// may have already closed the PseudoConsole when the command exited, so that
 	// output reads can get to EOF.  In that case, we don't need to close it
 	// again here.
 	if p.console != windows.InvalidHandle {
+		// ClosePseudoConsole has no return value and typically the syscall
+		// returns S_FALSE (a success value). We could ignore the return value
+		// and error here but we handle anyway, it just in case.
+		//
+		// Note that ClosePseudoConsole is a blocking system call and may write
+		// a final frame to the output buffer (p.outputWrite), so there must be
+		// a consumer (p.outputRead) to ensure we don't block here indefinitely.
+		//
+		// https://docs.microsoft.com/en-us/windows/console/closepseudoconsole
 		ret, _, err := procClosePseudoConsole.Call(uintptr(p.console))
-		if ret < 0 {
-			return xerrors.Errorf("close pseudo console: %w", err)
+		if winerrorFailed(ret) {
+			return xerrors.Errorf("close pseudo console (%d): %w", ret, err)
 		}
 		p.console = windows.InvalidHandle
 	}
 
-	// We always have these files
-	_ = p.outputRead.Close()
-	_ = p.inputWrite.Close()
-	// These get closed & unset if we Start() a new process.
+	return nil
+}
+
+func (p *ptyWindows) Close() error {
+	p.closeMutex.Lock()
+	defer p.closeMutex.Unlock()
+	if p.closed {
+		return nil
+	}
+
+	// Close the pseudo console, this will also terminate the process attached
+	// to this pty. If it was created via Start(), this also unblocks close of
+	// the readers below.
+	err := p.closeConsoleNoLock()
+	if err != nil {
+		return err
+	}
+
+	// Only set closed after the console has been successfully closed.
+	p.closed = true
+
+	// Close the pipes ensuring that the writer is closed before the respective
+	// reader, otherwise closing the reader may block indefinitely. Note that
+	// outputWrite and inputRead are unset when we Start() a new process.
 	if p.outputWrite != nil {
 		_ = p.outputWrite.Close()
 	}
+	_ = p.outputRead.Close()
+	_ = p.inputWrite.Close()
 	if p.inputRead != nil {
 		_ = p.inputRead.Close()
 	}
@@ -184,15 +213,13 @@ func (p *windowsProcess) waitInternal() {
 		// c.f. https://devblogs.microsoft.com/commandline/windows-command-line-introducing-the-windows-pseudo-console-conpty/
 		p.pw.closeMutex.Lock()
 		defer p.pw.closeMutex.Unlock()
-		if p.pw.console != windows.InvalidHandle {
-			ret, _, err := procClosePseudoConsole.Call(uintptr(p.pw.console))
-			if ret < 0 && p.cmdErr == nil {
-				// if we already have an error from the command, prefer that error
-				// but if the command succeeded and closing the PseudoConsole fails
-				// then record that error so that we have a chance to see it
-				p.cmdErr = err
-			}
-			p.pw.console = windows.InvalidHandle
+
+		err := p.pw.closeConsoleNoLock()
+		// if we already have an error from the command, prefer that error
+		// but if the command succeeded and closing the PseudoConsole fails
+		// then record that error so that we have a chance to see it
+		if err != nil && p.cmdErr == nil {
+			p.cmdErr = err
 		}
 	}()
 
@@ -225,3 +252,11 @@ func (p *windowsProcess) killOnContext(ctx context.Context) {
 		p.Kill()
 	}
 }
+
+// winerrorFailed returns true if the syscall failed, this function
+// assumes the return value is a 32-bit integer, like HRESULT.
+//
+// https://learn.microsoft.com/en-us/windows/win32/api/winerror/nf-winerror-failed
+func winerrorFailed(r1 uintptr) bool {
+	return int32(r1) < 0
+}
diff --git a/pty/ptytest/ptytest.go b/pty/ptytest/ptytest.go
@@ -354,7 +354,13 @@ type PTY struct {
 func (p *PTY) Close() error {
 	p.t.Helper()
 	pErr := p.PTY.Close()
-	eErr := p.outExpecter.close("close")
+	if pErr != nil {
+		p.logf("PTY: Close failed: %v", pErr)
+	}
+	eErr := p.outExpecter.close("PTY close")
+	if eErr != nil {
+		p.logf("PTY: close expecter failed: %v", eErr)
+	}
 	if pErr != nil {
 		return pErr
 	}
@@ -398,7 +404,13 @@ type PTYCmd struct {
 func (p *PTYCmd) Close() error {
 	p.t.Helper()
 	pErr := p.PTYCmd.Close()
-	eErr := p.outExpecter.close("close")
+	if pErr != nil {
+		p.logf("PTYCmd: Close failed: %v", pErr)
+	}
+	eErr := p.outExpecter.close("PTYCmd close")
+	if eErr != nil {
+		p.logf("PTYCmd: close expecter failed: %v", eErr)
+	}
 	if pErr != nil {
 		return pErr
 	}
