mitogen/parent: first stage: Bail out if stdin/stdout is not accessible

mitogen/parent.py

@@ -1415,11 +1415,15 @@ def __repr__(self):
     #   W: write side of interpreter stdin.
     #   r: read side of core_src FD.
     #   w: write side of core_src FD.
-
-    # Final os.close(STDERR_FILENO) to avoid --py-debug build corrupting stream with
-    # "[1234 refs]" during exit.
     @staticmethod
     def _first_stage():
+        # Bail out in case STDIN or STDOUT is not accessible (e.g. closed).
+        # Otherwise, os.pipe() could reuse file descriptors 0 or 1, leading to
+        # unexpected behavior that is difficult to diagnose.
+        os.fstat(0)
+        os.fstat(1)
+        # If STDERR (2) is already closed, it might be possible that one of the
+        # returned pipe FDs is 2.
         R,W=os.pipe()
         r,w=os.pipe()
         if os.fork():
@@ -1455,6 +1459,10 @@ def _first_stage():
         f.write(C)
         f.close()
         os.write(1,'MITO001\n'.encode())
+        # Final os.close(STDERR_FILENO) to avoid `--py-debug` build corrupting
+        # stream with "[1234 refs]" during exit.
+        # If STDERR is already closed an OSError is raised, but no one cares
+        # as STDERR is closed and the exit status is not forwarded.
         os.close(2)
 
     def get_python_argv(self):

tests/first_stage_test.py

@@ -1,6 +1,7 @@
 import errno
 import operator
 import os
+import pty
 
 import mitogen.core
 import mitogen.parent
@@ -142,6 +143,28 @@ def create_child(*a, **kw):
         return proc
 
 
+class DummyConnectionClosedStdout(mitogen.parent.Connection):
+    """Dummy closed stdout connection"""
+
+    create_child = staticmethod(create_child_using_pipes)
+    name_prefix = "dummy_closed_stdout"
+
+    #: Dictionary of extra kwargs passed to :attr:`create_child`.
+    create_child_args = {"blocking": True,
+                         "preexec_fn": lambda: os.close(pty.STDOUT_FILENO)}
+
+
+class DummyConnectionClosedStdin(mitogen.parent.Connection):
+    """Dummy closed stdin connection"""
+
+    create_child = staticmethod(create_child_using_pipes)
+    name_prefix = "dummy_closed_stdin"
+
+    #: Dictionary of extra kwargs passed to :attr:`create_child`.
+    create_child_args = {"blocking": True,
+                         "preexec_fn": lambda: os.close(pty.STDIN_FILENO)}
+
+
 class DummyConnectionEndlessBlockingRead(mitogen.parent.Connection):
     """Dummy connection that triggers a non-returning read(STDIN) call in the
     first_stage.
@@ -200,6 +223,24 @@ def test_blocking_stdin(self):
             ctx = self.router._connect(DummyConnectionBlocking, connect_timeout=0.5)
             self.assertEqual(3, ctx.call(operator.add, 1, 2))
 
+    def test_closed_communication_channel(self):
+        """Test that first stage does not work with a closed STDIN or STDOUT
+
+        The first stage of the boot command should bail out as soon as it
+        detects that STDIN or STDOUT is closed/unavailable and this results in
+        disconnected streams and that is mapped by the broker to an
+        :class:`mitogen.parent.EofError`.
+
+        """
+        with testlib.LogCapturer() as _:
+            e = self.assertRaises(mitogen.parent.EofError,
+                                  self.router._connect, DummyConnectionClosedStdout, connect_timeout=0.5)
+            self.assertIn("Bad file descriptor", str(e))
+
+            e = self.assertRaises(mitogen.parent.EofError,
+                                  self.router._connect, DummyConnectionClosedStdin, connect_timeout=0.5)
+            self.assertIn("Bad file descriptor", str(e))
+
     def test_broker_connect_eof_error(self):
         """Test that broker takes care about EOF errors in the first stage
 
@@ -365,3 +406,87 @@ def test_premature_eof(self):
         finally:
             proc.stdout.close()
             proc.stderr.close()
+
+    def test_closed_stdin(self):
+        """This test closes STDIN of the child process.
+
+        1. The child process detects that STDIN is unavailable
+        2. The child process terminates early with an OSError exception, and
+           reports the issue via exception printed on STDERR.
+        3. The parent process correctly identifies this condition.
+
+        """
+
+        options = mitogen.parent.Options(max_message_size=123)
+        conn = mitogen.parent.Connection(options, self.router)
+        conn.context = mitogen.core.Context(None, 123)
+        proc = testlib.subprocess.Popen(
+            args=conn.get_boot_command(),
+            stdout=testlib.subprocess.PIPE,
+            stderr=testlib.subprocess.PIPE,
+            preexec_fn=lambda: os.close(pty.STDIN_FILENO),
+            close_fds=True,
+        )
+        try:
+            stdout, stderr = proc.communicate(timeout=1)
+        except testlib.subprocess.TimeoutExpired:
+            proc.kill()
+            proc.wait(timeout=3)
+            self.fail("Closed STDIN situation was not recognized")
+        self.assertEqual(1, proc.returncode)
+        self.assertEqual(stdout, b"")
+        self.assertIn(
+            b("Bad file descriptor"),
+            stderr,
+        )
+
+    def test_closed_stdout(self):
+        """Test that first stage bails out if STDOUT is closed
+
+        This test closes STDOUT of the child process.
+
+        1. The child process detects that STDOUT is unavailable
+        2. The child process terminates early with an OSError exception, and
+           reports the issue via exception printed on STDERR.
+        3. The parent process correctly identifies this condition.
+
+        """
+        options = mitogen.parent.Options(max_message_size=123)
+        conn = mitogen.parent.Connection(options, self.router)
+        conn.context = mitogen.core.Context(None, 123)
+
+        stdout_r, stdout_w = mitogen.core.pipe()
+        mitogen.core.set_cloexec(stdout_r.fileno())
+        stderr_r, stderr_w = mitogen.core.pipe()
+        mitogen.core.set_cloexec(stderr_r.fileno())
+        try:
+            proc = testlib.subprocess.Popen(
+                args=conn.get_boot_command(),
+                stdout=stdout_w,
+                stderr=stderr_w,
+                preexec_fn=lambda: os.close(pty.STDOUT_FILENO),
+            )
+        except Exception:
+            stdout_r.close()
+            stdout_w.close()
+            stderr_w.close()
+            stderr_r.close()
+            raise
+        stdout_w.close()
+        stderr_w.close()
+        try:
+            returncode = proc.wait(timeout=1)
+        except testlib.subprocess.TimeoutExpired:
+            proc.kill()
+            proc.wait(timeout=3)
+            self.fail("Closed STDOUT situation was not detected")
+        else:
+            stderr = stderr_r.read()
+        finally:
+            stderr_r.close()
+            stdout_r.close()
+        self.assertEqual(1, returncode)
+        self.assertIn(
+            b("Bad file descriptor"),
+            stderr,
+        )