diff --git a/Lib/pty.py b/Lib/pty.py index 6571050886b..1d97994abef 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 STDIN_FILENO in rfds: + if master_fd in wfds: + n = os.write(master_fd, i_buf) + i_buf = i_buf[n:] + + 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 c723bb362c5..c9c2b42861c 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 42ec059a7c4..be8755637ff 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 00000000000..12ed52f9712 --- /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()