first_stage_test: Add test_closed_stdin and test_closed_stdout

Signed-off-by: Marc Hartmayer <[email protected]>

Author: mhartmay

tests/first_stage_test.py

@@ -1,6 +1,7 @@
 import fcntl
 import functools
 import operator
+import os
 
 import mitogen.core
 import mitogen.parent
@@ -269,3 +270,93 @@ def test_timeout_error(self):
             b(""),
             stderr,
         )
+
+    def test_closed_stdin(self):
+        """
+        The boot command should write an ECO marker to stdout, read the
+        preamble from stdin, then execute it.
+
+        TODO UNDER DISCUSSION
+
+        This test attaches closes stdin to create a specific failure
+        1. Fork child tries to read from STDIN, but fails as it is closed
+        2. Fork child raises TimeoutError
+        3. Fork child's file descriptors (write pipes) are closed by the OS
+        4. Fork parent does `dup(<read pipe>, <stdin>)` and `exec(<python>)`
+        5. Python reads `b''` (i.e. EOF) from stdin (a closed pipe)
+        6. Python runs `''` (a valid script) and exits with success
+        """
+        # We do not want to wait the default of 10s, change it to 0.1s
+        self.conn._first_stage_timeout = 0.1
+        args = self.conn.get_boot_command()
+
+        proc = testlib.subprocess.Popen(
+            args=args,
+            stdout=testlib.subprocess.PIPE,
+            stderr=testlib.subprocess.PIPE,
+            preexec_fn=lambda: os.close(0),
+            close_fds=True,
+        )
+        try:
+            stdout, stderr = proc.communicate(timeout=12)
+        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):
+        """The boot command should write an ECO marker to stdout, read the
+        preamble from stdin, then execute it.
+
+        TODO UNDER DISCUSSION
+
+        This test writes some data to STDIN and closes it then to create an
+        EOF situation.
+        1. Fork child tries to read from STDIN, but stops as EOF is received.
+        2. Fork child crashes (trying to decompress the junk data)
+        3. Fork child's file descriptors (write pipes) are closed by the OS
+        4. Fork parent does `dup(<read pipe>, <stdin>)` and `exec(<python>)`
+        5. Python reads `b''` (i.e. EOF) from stdin (a closed pipe)
+        6. Python runs `''` (a valid script) and exits with success"""
+
+        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=self.args,
+                stdout=stdout_w,
+                stderr=stderr_w,
+                preexec_fn=lambda: os.close(0),
+            )
+        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,
+        )