gh-96522: Fix deadlock in pty.spawn (#96639)

4 files changed, 56 insertions, 22 deletions

diff --git a/Lib/pty.py b/Lib/pty.py
index 6571050..1d97994 100644
--- a/Lib/pty.py
+++ b/Lib/pty.py
@@ -115,12 +115,6 @@ def fork():
     # Parent and child process.
     return pid, master_fd
 
-def _writen(fd, data):
-    """Write all the data to a descriptor."""
-    while data:
-        n = os.write(fd, data)
-        data = data[n:]
-
 def _read(fd):
     """Default read function."""
     return os.read(fd, 1024)
@@ -130,9 +124,42 @@ def _copy(master_fd, master_read=_read, stdin_read=_read):
     Copies
             pty master -> standard output   (master_read)
             standard input -> pty master    (stdin_read)"""
-    fds = [master_fd, STDIN_FILENO]
-    while fds:
-        rfds, _wfds, _xfds = select(fds, [], [])
+    if os.get_blocking(master_fd):
+        # If we write more than tty/ndisc is willing to buffer, we may block
+        # indefinitely. So we set master_fd to non-blocking temporarily during
+        # the copy operation.
+        os.set_blocking(master_fd, False)
+        try:
+            _copy(master_fd, master_read=master_read, stdin_read=stdin_read)
+        finally:
+            # restore blocking mode for backwards compatibility
+            os.set_blocking(master_fd, True)
+        return
+    high_waterlevel = 4096
+    stdin_avail = master_fd != STDIN_FILENO
+    stdout_avail = master_fd != STDOUT_FILENO
+    i_buf = b''
+    o_buf = b''
+    while 1:
+        rfds = []
+        wfds = []
+        if stdin_avail and len(i_buf) < high_waterlevel:
+            rfds.append(STDIN_FILENO)
+        if stdout_avail and len(o_buf) < high_waterlevel:
+            rfds.append(master_fd)
+        if stdout_avail and len(o_buf) > 0:
+            wfds.append(STDOUT_FILENO)
+        if len(i_buf) > 0:
+            wfds.append(master_fd)
+
+        rfds, wfds, _xfds = select(rfds, wfds, [])
+
+        if STDOUT_FILENO in wfds:
+            try:
+                n = os.write(STDOUT_FILENO, o_buf)
+                o_buf = o_buf[n:]
+            except OSError:
+                stdout_avail = False
 
         if master_fd in rfds:
             # Some OSes signal EOF by returning an empty byte string,
@@ -144,15 +171,18 @@ def _copy(master_fd, master_read=_read, stdin_read=_read):
             if not data:  # Reached EOF.
                 return    # Assume the child process has exited and is
                           # unreachable, so we clean up.
-            else:
-                os.write(STDOUT_FILENO, data)
+            o_buf += data
+
+        if master_fd in wfds:
+            n = os.write(master_fd, i_buf)
+            i_buf = i_buf[n:]
 
-        if STDIN_FILENO in rfds:
+        if stdin_avail and STDIN_FILENO in rfds:
             data = stdin_read(STDIN_FILENO)
             if not data:
-                fds.remove(STDIN_FILENO)
+                stdin_avail = False
             else:
-                _writen(master_fd, data)
+                i_buf += data
 
 def spawn(argv, master_read=_read, stdin_read=_read):
     """Create a spawned process."""
diff --git a/Lib/test/test_pty.py b/Lib/test/test_pty.py
index c723bb3..c9c2b42 100644
--- a/Lib/test/test_pty.py
+++ b/Lib/test/test_pty.py
@@ -312,8 +312,8 @@ class SmallPtyTests(unittest.TestCase):
         self.orig_pty_waitpid = pty.waitpid
         self.fds = []  # A list of file descriptors to close.
         self.files = []
-        self.select_rfds_lengths = []
-        self.select_rfds_results = []
+        self.select_input = []
+        self.select_output = []
         self.tcsetattr_mode_setting = None
 
     def tearDown(self):
@@ -350,8 +350,8 @@ class SmallPtyTests(unittest.TestCase):
 
     def _mock_select(self, rfds, wfds, xfds):
         # This will raise IndexError when no more expected calls exist.
-        self.assertEqual(self.select_rfds_lengths.pop(0), len(rfds))
-        return self.select_rfds_results.pop(0), [], []
+        self.assertEqual((rfds, wfds, xfds), self.select_input.pop(0))
+        return self.select_output.pop(0)
 
     def _make_mock_fork(self, pid):
         def mock_fork():
@@ -374,11 +374,13 @@ class SmallPtyTests(unittest.TestCase):
         os.write(masters[1], b'from master')
         os.write(write_to_stdin_fd, b'from stdin')
 
-        # Expect two select calls, the last one will cause IndexError
+        # Expect three select calls, the last one will cause IndexError
         pty.select = self._mock_select
-        self.select_rfds_lengths.append(2)
-        self.select_rfds_results.append([mock_stdin_fd, masters[0]])
-        self.select_rfds_lengths.append(2)
+        self.select_input.append(([mock_stdin_fd, masters[0]], [], []))
+        self.select_output.append(([mock_stdin_fd, masters[0]], [], []))
+        self.select_input.append(([mock_stdin_fd, masters[0]], [mock_stdout_fd, masters[0]], []))
+        self.select_output.append(([], [mock_stdout_fd, masters[0]], []))
+        self.select_input.append(([mock_stdin_fd, masters[0]], [], []))
 
         with self.assertRaises(IndexError):
             pty._copy(masters[0])
diff --git a/Misc/ACKS b/Misc/ACKS
index 42ec059..be87556 100644
--- a/Misc/ACKS
+++ b/Misc/ACKS
@@ -2060,6 +2060,7 @@ Yuxiao Zeng
 Uwe Zessin
 Cheng Zhang
 George Zhang
+Youfu Zhang
 Charlie Zhao
 Kai Zhu
 Tarek Ziadé
diff --git a/Misc/NEWS.d/next/Library/2022-09-07-09-32-07.gh-issue-96522.t73oqp.rst b/Misc/NEWS.d/next/Library/2022-09-07-09-32-07.gh-issue-96522.t73oqp.rst
new file mode 100644
index 0000000..12ed52f
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2022-09-07-09-32-07.gh-issue-96522.t73oqp.rst
@@ -0,0 +1 @@
+Fix potential deadlock in pty.spawn()