first_stage_test: Add test_eof_too_early test

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

Author: mhartmay

tests/first_stage_test.py

@@ -292,3 +292,50 @@ def test_stdin_blocking(self):
             finally:
                 proc.stdout.close()
                 proc.stderr.close()
+
+    def test_eof_too_early(self):
+        """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 = subprocess.Popen(
+            args=self.args,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            stdin=subprocess.PIPE,
+        )
+
+        # Do not send all of the data from the preamble
+        proc.stdin.write(self.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()