mitogen: Fix non-blocking IO errors in first stage of bootstrap

.ci/ci_lib.py

@@ -41,6 +41,8 @@
     'MITOGEN_TEST_IMAGE_TEMPLATE',
     'ghcr.io/mitogen-hq/%(distro)s-test:2021',
 )
+SUDOERS_DEFAULTS_SRC = './tests/image_prep/files/sudoers_defaults'
+SUDOERS_DEFAULTS_DEST = '/etc/sudoers.d/mitogen_test_defaults'
 TESTS_SSH_PRIVATE_KEY_FILE = os.path.join(GIT_ROOT, 'tests/data/docker/mitogen__has_sudo_pubkey.key')
 
 
@@ -58,6 +60,7 @@ def _have_cmd(args):
     try:
         subprocess.run(
             args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+            check=True,
         )
     except OSError as exc:
         if exc.errno == errno.ENOENT:

.ci/mitogen_tests.py

@@ -2,6 +2,7 @@
 # Run the Mitogen tests.
 
 import os
+import subprocess
 
 import ci_lib
 
@@ -13,6 +14,14 @@
 if not ci_lib.have_docker():
     os.environ['SKIP_DOCKER_TESTS'] = '1'
 
+subprocess.check_call(
+    "umask 0022; sudo cp '%s' '%s'"
+    % (ci_lib.SUDOERS_DEFAULTS_SRC, ci_lib.SUDOERS_DEFAULTS_DEST),
+    shell=True,
+)
+subprocess.check_call(['sudo', 'visudo', '-cf', ci_lib.SUDOERS_DEFAULTS_DEST])
+subprocess.check_call(['sudo', '-l'])
+
 interesting = ci_lib.get_interesting_procs()
 ci_lib.run('./run_tests -v')
 ci_lib.check_stray_processes(interesting)

docs/changelog.rst

@@ -21,6 +21,14 @@ To avail of fixes in an unreleased version, please download a ZIP file
 In progress (unreleased)
 ------------------------
 
+* :gh:issue:`1306` :mod:`ansible_mitogen`: Fix non-blocking IO errors in
+  first stage of bootstrap
+* :gh:issue:`1306` CI: Report sudo version on Ansible targets
+* :gh:issue:`1306` CI: Move sudo test users defaults into ``/etc/sudoers.d``
+* :gh:issue:`1306` preamble_size: Fix variability of measured command size
+* :gh:issue:`1306` tests: Count bytes written in ``stdio_test.StdIOTest``
+* :gh:issue:`1306` tests: Check stdio is blocking in sudo contexts
+
 
 v0.3.27 (2025-08-20)
 --------------------

mitogen/parent.py

@@ -1440,9 +1440,9 @@ def _first_stage():
             os.environ['ARGV0']=sys.executable
             os.execl(sys.executable,sys.executable+'(mitogen:CONTEXT_NAME)')
         os.write(1,'MITO000\n'.encode())
-        fp=os.fdopen(0,'rb')
-        C=zlib.decompress(fp.read(PREAMBLE_COMPRESSED_LEN))
-        fp.close()
+        C=''.encode()
+        while PREAMBLE_COMPRESSED_LEN-len(C)and select.select([0],[],[]):C+=os.read(0,PREAMBLE_COMPRESSED_LEN-len(C))
+        C=zlib.decompress(C)
         fp=os.fdopen(W,'wb',0)
         fp.write(C)
         fp.close()
@@ -1478,11 +1478,12 @@ def get_boot_command(self):
 
         # Just enough to decode, decompress, and exec the first stage.
         # Priorities: wider compatibility, faster startup, shorter length.
-        # `import os` here, instead of stage 1, to save a few bytes.
         # `sys.path=...` for https://github.com/python/cpython/issues/115911.
+        # `import os,select` here (not stage 1) to save a few bytes overall.
         return self.get_python_argv() + [
             '-c',
-            'import sys;sys.path=[p for p in sys.path if p];import binascii,os,zlib;'
+            'import sys;sys.path=[p for p in sys.path if p];'
+            'import binascii,os,select,zlib;'
             'exec(zlib.decompress(binascii.a2b_base64("%s")))' % (encoded.decode(),),
         ]
 

preamble_size.py

@@ -8,6 +8,7 @@
 import sys
 import zlib
 
+import mitogen.core
 import mitogen.fakessh
 import mitogen.fork
 import mitogen.master
@@ -18,14 +19,28 @@
 import mitogen.ssh
 import mitogen.sudo
 
+
+class Table(object):
+    HEADERS = (' ', 'Original', 'Minimized', 'Compressed')
+    HEAD_FMT = '{:20} {:^15}  {:^19}  {:^19}'
+    ROW_FMT =  '%-20s %6i %5.1fKiB  %5i %4.1fKiB %4.1f%%  %5i %4.1fKiB %4.1f%%'
+
+    def header(self):
+        return self.HEAD_FMT.format(*self.HEADERS)
+
+
 router = mitogen.master.Router()
 context = mitogen.parent.Context(router, 0)
-options = mitogen.ssh.Options(max_message_size=0, hostname='foo')
+options = mitogen.ssh.Options(
+    hostname='foo',
+    max_message_size=0,
+    remote_name='alice@host:1234',
+)
 conn = mitogen.ssh.Connection(options, router)
 conn.context = context
 
 print('SSH command size: %s' % (len(' '.join(conn.get_boot_command())),))
-print('Bootstrap (mitogen.core) size: %s (%.2fKiB)' % (
+print('Preamble (mitogen.core + econtext) size: %s (%.2fKiB)' % (
     len(conn.get_preamble()),
     len(conn.get_preamble()) / 1024.0,
 ))
@@ -36,17 +51,10 @@
     exit()
 
 
-print(
-    '                           '
-    ' '
-    '  Original   '
-    '  '
-    '     Minimized     '
-    '  '
-    '    Compressed     '
-)
-
+table = Table()
+print(table.header())
 for mod in (
+        mitogen.core,
         mitogen.parent,
         mitogen.fork,
         mitogen.ssh,
@@ -63,13 +71,7 @@
     compressed = zlib.compress(minimized.encode(), 9)
     compressed_size = len(compressed)
     print(
-        '%-25s'
-        ' '
-        '%5i %4.1fKiB'
-        '  '
-        '%5i %4.1fKiB %.1f%%'
-        '  '
-        '%5i %4.1fKiB %.1f%%'
+        table.ROW_FMT
     % (
         mod.__name__,
         original_size,

tests/ansible/setup/report_targets.yml

@@ -13,3 +13,23 @@
     - debug: {var: ansible_facts.osversion}
     - debug: {var: ansible_facts.python}
     - debug: {var: ansible_facts.system}
+
+- name: Check target versions
+  hosts: localhost:test-targets
+  check_mode: false
+  tasks:
+    - name: Get command versions
+      command:
+        cmd: "{{ item.cmd }}"
+      changed_when: false
+      check_mode: false
+      loop:
+        - cmd: sudo -V
+      register: command_versions
+
+    - name: Show command versions
+      debug:
+        msg: |
+          cmd: {{ item.item.cmd }}
+          {{ item.stdout }}
+      loop: "{{ command_versions.results }}"

tests/data/stdio_checks.py

@@ -3,9 +3,24 @@
 import sys
 
 
+def _shout_stdout_py3(size):
+    nwritten = sys.stdout.write('A' * size)
+    return nwritten
+
+
+def _shout_stdout_py2(size):
+    shout = 'A' * size
+    nwritten = 0
+    while nwritten < size:
+        nwritten += os.write(sys.stdout.fileno(), shout[-nwritten:])
+    return nwritten
+
+
 def shout_stdout(size):
-    sys.stdout.write('A' * size)
-    return 'success'
+    if sys.version_info > (3, 0):
+        return _shout_stdout_py3(size)
+    else:
+        return _shout_stdout_py2(size)
 
 
 def file_is_blocking(fobj):

tests/first_stage_test.py

@@ -1,5 +1,6 @@
 import subprocess
 
+import mitogen.core
 import mitogen.parent
 from mitogen.core import b
 
@@ -21,14 +22,18 @@ def test_valid_syntax(self):
         conn.context = mitogen.core.Context(None, 123)
         args = conn.get_boot_command()
 
-        # Executing the boot command will print "EC0" and expect to read from
-        # stdin, which will fail because it's pointing at /dev/null, causing
-        # the forked child to crash with an EOFError and disconnect its write
-        # pipe. The forked and freshly execed parent will get a 0-byte read
-        # from the pipe, which is a valid script, and therefore exit indicating
-        # success.
+        # The boot command should write an ECO marker to stdout, read the
+        # preamble from stdin, then execute it.
 
-        fp = open("/dev/null", "r")
+        # This test attaches /dev/zero to stdin to create a specific failure
+        # 1. Fork child reads PREAMBLE_COMPRESSED_LEN bytes of junk (all `\0`)
+        # 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
+
+        fp = open("/dev/zero", "r")
         try:
             proc = subprocess.Popen(args,
                 stdin=fp,
@@ -39,6 +44,9 @@ def test_valid_syntax(self):
             self.assertEqual(0, proc.returncode)
             self.assertEqual(stdout,
                 mitogen.parent.BootstrapProtocol.EC0_MARKER+b('\n'))
-            self.assertIn(b("Error -5 while decompressing data"), stderr)
+            self.assertIn(
+                b("Error -3 while decompressing data"),  # Unknown compression method
+                stderr,
+            )
         finally:
             fp.close()

tests/image_prep/_user_accounts.yml

@@ -157,15 +157,14 @@
             owner: mitogen__has_sudo_pubkey
             group: mitogen__group
 
-    - name: Configure sudoers defaults
-      blockinfile:
-        path: /etc/sudoers
-        marker: "# {mark} Mitogen test defaults"
-        block: |
-          Defaults>mitogen__pw_required targetpw
-          Defaults>mitogen__require_tty requiretty
-          Defaults>mitogen__require_tty_pw_required requiretty,targetpw
+    - name: Configure sudoers
+      copy:
+        src: "{{ item.src }}"
+        dest: "{{ item.dest }}"
+        mode: ug=r,o=
         validate: '/usr/sbin/visudo -cf %s'
+      with_items:
+        - {src: sudoers_defaults, dest: /etc/sudoers.d/mitogen_test_defaults}
 
     - name: Configure sudoers users
       blockinfile:

tests/image_prep/files/sudoers_defaults

@@ -0,0 +1,7 @@
+# Testing non-blocking stdio during bootstrap
+# https://github.com/mitogen-hq/mitogen/issues/1306
+Defaults log_output
+
+Defaults>mitogen__pw_required targetpw
+Defaults>mitogen__require_tty requiretty
+Defaults>mitogen__require_tty_pw_required requiretty,targetpw

tests/stdio_test.py

@@ -1,28 +1,46 @@
+import unittest
+
 import testlib
 
 import stdio_checks
 
 
-class StdIOTest(testlib.RouterMixin, testlib.TestCase):
+class StdIOMixin(testlib.RouterMixin):
     """
     Test that stdin, stdout, and stderr conform to common expectations,
     such as blocking IO.
     """
-    def test_can_write_stdout_1_mib(self):
+    def check_can_write_stdout_1_mib(self, context):
         """
         Writing to stdout should not raise EAGAIN. Regression test for
         https://github.com/mitogen-hq/mitogen/issues/712.
         """
         size = 1 * 2**20
-        context = self.router.local()
-        result = context.call(stdio_checks.shout_stdout, size)
-        self.assertEqual('success', result)
+        nwritten = context.call(stdio_checks.shout_stdout, size)
+        self.assertEqual(nwritten, size)
 
-    def test_stdio_is_blocking(self):
-        context = self.router.local()
+    def check_stdio_is_blocking(self, context):
         stdin_blocking, stdout_blocking, stderr_blocking = context.call(
             stdio_checks.stdio_is_blocking,
         )
         self.assertTrue(stdin_blocking)
         self.assertTrue(stdout_blocking)
         self.assertTrue(stderr_blocking)
+
+
+class LocalTest(StdIOMixin, testlib.TestCase):
+    def test_can_write_stdout_1_mib(self):
+        self.check_can_write_stdout_1_mib(self.router.local())
+
+    def test_stdio_is_blocking(self):
+        self.check_stdio_is_blocking(self.router.local())
+
+
+class SudoTest(StdIOMixin, testlib.TestCase):
+    @unittest.skipIf(not testlib.have_sudo_nopassword(), 'Needs passwordless sudo')
+    def test_can_write_stdout_1_mib(self):
+        self.check_can_write_stdout_1_mib(self.router.sudo())
+
+    @unittest.skipIf(not testlib.have_sudo_nopassword(), 'Needs passwordless sudo')
+    def test_stdio_is_blocking(self):
+        self.check_stdio_is_blocking(self.router.sudo())

tests/testlib.py

@@ -160,6 +160,7 @@ def _have_cmd(args):
     try:
         subprocess.run(
             args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+            check=True,
         )
     except OSError as exc:
         if exc.errno == errno.ENOENT:
@@ -178,6 +179,15 @@ def have_python3():
     return _have_cmd(['python3'])
 
 
+def have_sudo_nopassword():
+    """
+    Return True if we can run `sudo` with no password, otherwise False.
+
+    Any cached credentials are ignored.
+    """
+    return _have_cmd(['sudo', '-kn', 'true'])
+
+
 def retry(fn, on, max_attempts, delay):
     for i in range(max_attempts):
         try: