tests: Add first_stage_test that checks that the stage handles an too early EOF

Add first_stage_test that checks that the first stage handles an too early EOF.

Without the next commit, the process enters an infinite loop, causing the Python
interpreter to consume up to 100% of a CPU core. The test kills this process
after running into the test timeout of 10s.

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

Author: mhartmay

tests/first_stage_test.py

@@ -91,3 +91,53 @@ def test_stage(self):
                 b(""),
                 stderr,
             )
+
+    def test_eof_too_early(self):
+        options = mitogen.parent.Options(max_message_size=123)
+        conn = mitogen.parent.Connection(options, self.router)
+        conn.context = mitogen.core.Context(None, 123)
+        args = conn.get_boot_command()
+        preamble = conn.get_preamble()
+
+        # The boot command should write an ECO marker to stdout, read the
+        # preamble from stdin, then execute it.
+
+        # 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
+
+        proc = mitogen.parent.popen(args=args,
+                                    stdout=subprocess.PIPE,
+                                    stderr=subprocess.PIPE,
+                                    stdin=subprocess.PIPE,
+                                    )
+
+        # Do not send all of the data from the preamble
+        proc.stdin.write(preamble[:-128])
+        proc.stdin.flush()
+        proc.stdin.close()
+        try:
+            try:
+                returncode = proc.wait(timeout=10)
+            except TypeError:
+                returncode = proc.wait()
+        except:
+            proc.kill()
+            self.fail("First stage did not handle EOF on STDIN")
+        else:
+            try:
+                self.assertEqual(0, returncode)
+                self.assertEqual(proc.stdout.read(),
+                                 mitogen.parent.BootstrapProtocol.EC0_MARKER+b('\n'))
+                self.assertIn(
+                    b("Error -5 while decompressing data"),
+                    proc.stderr.read(),
+                )
+            finally:
+                proc.stdout.close()
+                proc.stderr.close()