From 28cf5beac000930ac5e5a718e06a883588915852 Mon Sep 17 00:00:00 2001 From: Cornelius Diekmann Date: Fri, 23 Dec 2016 20:22:12 +0100 Subject: [PATCH] Handling closed STDIN in Lib/pty.py --- Lib/pty.py | 10 ++++++-- Lib/test/test_pty.py | 64 +++++++++++++++++++++++++++++++++++++++++++++------- 2 files changed, 64 insertions(+), 10 deletions(-) diff --git a/Lib/pty.py b/Lib/pty.py index e841f12..012c053 100644 --- a/Lib/pty.py +++ b/Lib/pty.py @@ -133,17 +133,23 @@ def _copy(master_fd, master_read=_read, stdin_read=_read): standard input -> pty master (stdin_read)""" fds = [master_fd, STDIN_FILENO] while True: + if not fds: + # STDIN or master reached EOF + # Note the *or* + return rfds, wfds, xfds = select(fds, [], []) if master_fd in rfds: data = master_read(master_fd) if not data: # Reached EOF. - fds.remove(master_fd) + fds = [] + # Checking for STDIN in rfds may still occur one last time. + # Triggers BrokenPipeError if master is closed for writing. else: os.write(STDOUT_FILENO, data) if STDIN_FILENO in rfds: data = stdin_read(STDIN_FILENO) if not data: - fds.remove(STDIN_FILENO) + fds = [] else: _writen(master_fd, data) diff --git a/Lib/test/test_pty.py b/Lib/test/test_pty.py index 15f88e4..54cfade 100644 --- a/Lib/test/test_pty.py +++ b/Lib/test/test_pty.py @@ -268,8 +268,8 @@ class SmallPtyTests(unittest.TestCase): self.assertEqual(os.read(read_from_stdout_fd, 20), b'from master') self.assertEqual(os.read(masters[1], 20), b'from stdin') - def test__copy_eof_on_all(self): - """Test the empty read EOF case on both master_fd and stdin.""" + def __copy_eof_helper(self, close_stdin, close_master): + """Helper to test the empty read EOF case on master_fd and/or stdin.""" read_from_stdout_fd, mock_stdout_fd = self._pipe() pty.STDOUT_FILENO = mock_stdout_fd mock_stdin_fd, write_to_stdin_fd = self._pipe() @@ -277,21 +277,69 @@ class SmallPtyTests(unittest.TestCase): socketpair = self._socketpair() masters = [s.fileno() for s in socketpair] - socketpair[1].close() - os.close(write_to_stdin_fd) + # If both not close_master and not close_stdin, we expect to block + # forever in the select call of pty._copy. The mock_select does not + # model this. Reject these test parameters. + if not close_master and not close_stdin: + raise ValueError("At least one fd needs to be closed") - # Expect two select calls, the last one will cause IndexError + + # close fds or fill with dummy data in order to prevent blocking on + # one read call + if close_master: + socketpair[1].close() + else: + os.write(masters[1], b'from master') + if close_stdin: + os.close(write_to_stdin_fd) + else: + os.write(write_to_stdin_fd, b'from stdin') + + + # Expect two select calls, then a normal return on EOF pty.select = self._mock_select self.select_rfds_lengths.append(2) self.select_rfds_results.append([mock_stdin_fd, masters[0]]) - # We expect that both fds were removed from the fds list as they - # both encountered an EOF before the second select call. - self.select_rfds_lengths.append(0) + # We expect that both fds were removed from the fds list. + + if close_master and not close_stdin: + with self.assertRaises(BrokenPipeError): + # We expect to find a closed master but some data in STDIN. + # The _copy loop will try to write the data from STDIN to + # master, which is a broken pipe. + # + # Only closed master and not closed STDIN is a problem: + # closed STDIN but not closed master is not a broken pipe. + # In this case, we just read from master and write to STDOUT + pty._copy(masters[0]) + else: + # We expect that this returns immediately + pty._copy(masters[0]) + # We expect that everything is consumed + self.assertEquals(self.select_rfds_results, []) + self.assertEquals(self.select_rfds_lengths, []) + + # We expect that calling again raises an error with self.assertRaises(IndexError): pty._copy(masters[0]) + def test__copy_eof_on_all(self): + """Test the empty read EOF case on both master_fd and stdin.""" + self.__copy_eof_helper(close_stdin=True, close_master=True) + + def test__copy_eof_on_stdin(self): + """Test the empty read EOF case on only stdin.""" + self.__copy_eof_helper(close_stdin=True, close_master=False) + + def test__copy_eof_on_master(self): + """Test the empty read EOF case on only master_fd.""" + self.__copy_eof_helper(close_stdin=False, close_master=True) + def test__copy_eof_on_helper(self): + """Test that the helper function rejects wrong usage""" + with self.assertRaises(ValueError): + self.__copy_eof_helper(close_stdin=False, close_master=False) def tearDownModule(): reap_children() -- 2.1.4