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