diff --git a/ansible_mitogen/connection.py b/ansible_mitogen/connection.py index a81a4bd6..bd0d79ff 100644 --- a/ansible_mitogen/connection.py +++ b/ansible_mitogen/connection.py @@ -257,6 +257,14 @@ class Connection(ansible.plugins.connection.ConnectionBase): else: self.context = self.host + def get_context_name(self): + """ + Return the name of the target context we issue commands against, i.e. a + unique string useful as a key for related data, such as a list of + modules uploaded to the target. + """ + return self.context.name + def close(self): """ Arrange for the mitogen.master.Router running in the worker to diff --git a/ansible_mitogen/executor.py b/ansible_mitogen/executor.py deleted file mode 100644 index 00e91d05..00000000 --- a/ansible_mitogen/executor.py +++ /dev/null @@ -1,287 +0,0 @@ -# Copyright 2017, David Wilson -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, -# this list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its contributors -# may be used to endorse or promote products derived from this software without -# specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE -# POSSIBILITY OF SUCH DAMAGE. - -from __future__ import absolute_import -import json -import os -import tempfile - -import ansible_mitogen.helpers - -try: - from shlex import quote as shlex_quote -except ImportError: - from pipes import quote as shlex_quote - -# Prevent accidental import of an Ansible module from hanging on stdin read. -import ansible.module_utils.basic -ansible.module_utils.basic._ANSIBLE_ARGS = '{}' - - -class Exit(Exception): - """ - Raised when a module exits with success. - """ - def __init__(self, dct): - self.dct = dct - - -class ModuleError(Exception): - """ - Raised when a module voluntarily indicates failure via .fail_json(). - """ - def __init__(self, msg, dct): - Exception.__init__(self, msg) - self.dct = dct - - -class TemporaryEnvironment(object): - def __init__(self, env=None): - self.original = os.environ.copy() - self.env = env or {} - os.environ.update((k, str(v)) for k, v in self.env.iteritems()) - - def revert(self): - os.environ.clear() - os.environ.update(self.original) - - -class MethodOverrides(object): - @staticmethod - def exit_json(self, **kwargs): - """ - Replace AnsibleModule.exit_json() with something that doesn't try to - kill the process or JSON-encode the result dictionary. Instead, cause - Exit to be raised, with a `dct` attribute containing the successful - result dictionary. - """ - self.add_path_info(kwargs) - kwargs.setdefault('changed', False) - kwargs.setdefault('invocation', { - 'module_args': self.params - }) - kwargs = ansible.module_utils.basic.remove_values( - kwargs, - self.no_log_values, - ) - self.do_cleanup_files() - raise Exit(kwargs) - - @staticmethod - def fail_json(self, **kwargs): - """ - Replace AnsibleModule.fail_json() with something that raises - ModuleError, which includes a `dct` attribute. - """ - self.add_path_info(kwargs) - kwargs.setdefault('failed', True) - kwargs.setdefault('invocation', { - 'module_args': self.params - }) - kwargs = ansible.module_utils.basic.remove_values( - kwargs, - self.no_log_values, - ) - self.do_cleanup_files() - raise ModuleError(kwargs.get('msg'), kwargs) - - klass = ansible.module_utils.basic.AnsibleModule - - def __init__(self): - self._original_exit_json = self.klass.exit_json - self._original_fail_json = self.klass.fail_json - self.klass.exit_json = self.exit_json - self.klass.fail_json = self.fail_json - - def revert(self): - self.klass.exit_json = self._original_exit_json - self.klass.fail_json = self._original_fail_json - - -class ModuleArguments(object): - """ - Patch the ansible.module_utils.basic global arguments variable on - construction, and revert the changes on call to :meth:`revert`. - """ - def __init__(self, args): - self.original = ansible.module_utils.basic._ANSIBLE_ARGS - ansible.module_utils.basic._ANSIBLE_ARGS = json.dumps({ - 'ANSIBLE_MODULE_ARGS': args - }) - - def revert(self): - ansible.module_utils.basic._ANSIBLE_ARGS = self._original_args - - -class Runner(object): - def __init__(self, module, raw_params=None, args=None, env=None, - runner_params=None): - if args is None: - args = {} - if raw_params is not None: - args['_raw_params'] = raw_params - if runner_params is None: - runner_params = {} - - self.module = module - self.raw_params = raw_params - self.args = args - self.env = env - self.runner_params = runner_params - - def setup(self): - self._env = TemporaryEnvironment(self.env) - - def revert(self): - self._env.revert() - - def _run(self): - raise NotImplementedError() - - def run(self): - """ - Set up the process environment in preparation for running an Ansible - module. This monkey-patches the Ansible libraries in various places to - prevent it from trying to kill the process on completion, and to - prevent it from reading sys.stdin. - """ - self.setup() - try: - return self._run() - finally: - self.revert() - - -class PythonRunner(object): - """ - Execute a new-style Ansible module, where Module Replacer-related tricks - aren't required. - """ - def setup(self): - super(PythonRunner, self).setup() - self._overrides = MethodOverrides() - self._args = ModuleArguments(self.args) - - def revert(self): - super(PythonRunner, self).revert() - self._args.revert() - self._overrides.revert() - - def _run(self): - try: - mod = __import__(self.module, {}, {}, ['']) - # Ansible modules begin execution on import. Thus the above - # __import__ will cause either Exit or ModuleError to be raised. If - # we reach the line below, the module did not execute and must - # already have been imported for a previous invocation, so we need - # to invoke main explicitly. - mod.main() - except (Exit, ModuleError), e: - return json.dumps(e.dct) - - assert False, "Module returned no result." - - -class BinaryRunner(object): - def setup(self): - super(BinaryRunner, self).setup() - self._setup_binary() - self._setup_args() - - def _get_binary(self): - """ - Fetch the module binary from the master if necessary. - """ - return ansible_mitogen.helpers.get_file( - path=self.runner_params['path'], - ) - - def _get_args(self): - """ - Return the module arguments formatted as JSON. - """ - return json.dumps(self.args) - - def _setup_program(self): - """ - Create a temporary file containing the program code. The code is - fetched via :meth:`_get_binary`. - """ - self.bin_fp = tempfile.NamedTemporaryFile( - prefix='ansible_mitogen', - suffix='-binary', - ) - self.bin_fp.write(self._get_binary()) - self.bin_fp.flush() - os.chmod(self.fp.name, int('0700', 8)) - - def _setup_args(self): - """ - Create a temporary file containing the module's arguments. The - arguments are formatted via :meth:`_get_args`. - """ - self.args_fp = tempfile.NamedTemporaryFile( - prefix='ansible_mitogen', - suffix='-args', - ) - self.args_fp.write(self._get_args()) - self.args_fp.flush() - - def revert(self): - """ - Delete the temporary binary and argument files. - """ - self.args_fp.close() - self.bin_fp.close() - super(BinaryRunner, self).revert() - - def _run(self): - rc, stdout, stderr = ansible_mitogen.helpers.exec_args( - args=[self.bin_fp.name, self.args_fp.name], - ) - # ... - assert 0 - - -class WantJsonRunner(BinaryRunner): - def _get_binary(self): - s = super(WantJsonRunner, self)._get_binary() - # fix up shebang. - return s - - -class OldStyleRunner(BinaryRunner): - def _get_args(self): - """ - Mimic the argument formatting behaviour of - ActionBase._execute_module(). - """ - return ' '.join( - '%s=%s' % (key, shlex_quote(str(self.args[key]))) - for key in self.args - ) diff --git a/ansible_mitogen/helpers.py b/ansible_mitogen/helpers.py index 9e8693d9..270f3244 100644 --- a/ansible_mitogen/helpers.py +++ b/ansible_mitogen/helpers.py @@ -26,7 +26,9 @@ # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE # POSSIBILITY OF SUCH DAMAGE. +from __future__ import absolute_import import json +import logging import operator import os import pwd @@ -36,12 +38,18 @@ import stat import subprocess import tempfile import threading +import zlib import mitogen.core +import mitogen.service +import ansible_mitogen.runner +import ansible_mitogen.services -# Prevent accidental import of an Ansible module from hanging on stdin read. -import ansible.module_utils.basic -ansible.module_utils.basic._ANSIBLE_ARGS = '{}' + +LOG = logging.getLogger(__name__) + +#: Caching of fetched file data. +_file_cache = {} #: Mapping of job_id<->result dict _result_by_job_id = {} @@ -50,124 +58,52 @@ _result_by_job_id = {} _thread_by_job_id = {} -class Exit(Exception): - """ - Raised when a module exits with success. - """ - def __init__(self, dct): - self.dct = dct - - -class ModuleError(Exception): - """ - Raised when a module voluntarily indicates failure via .fail_json(). +def get_file(context, path): """ - def __init__(self, msg, dct): - Exception.__init__(self, msg) - self.dct = dct - - -def monkey_exit_json(self, **kwargs): - """ - Replace AnsibleModule.exit_json() with something that doesn't try to kill - the process or JSON-encode the result dictionary. Instead, cause Exit to be - raised, with a `dct` attribute containing the successful result dictionary. - """ - self.add_path_info(kwargs) - kwargs.setdefault('changed', False) - kwargs.setdefault('invocation', { - 'module_args': self.params - }) - kwargs = ansible.module_utils.basic.remove_values( - kwargs, - self.no_log_values - ) - self.do_cleanup_files() - raise Exit(kwargs) + Basic in-memory caching module fetcher. This generates an one roundtrip for + every previously unseen module, so it is only temporary. - -def monkey_fail_json(self, **kwargs): - """ - Replace AnsibleModule.fail_json() with something that raises ModuleError, - which includes a `dct` attribute. - """ - self.add_path_info(kwargs) - kwargs.setdefault('failed', True) - kwargs.setdefault('invocation', { - 'module_args': self.params - }) - kwargs = ansible.module_utils.basic.remove_values( - kwargs, - self.no_log_values - ) - self.do_cleanup_files() - raise ModuleError(kwargs.get('msg'), kwargs) - - -def module_fixups(mod): - """ - Apply fixups for known problems with mainline Ansible modules. + :param context: + Context we should direct FileService requests to. For now (and probably + forever) this is just the top-level Mitogen connection manager process. + :param path: + Path to fetch from FileService, must previously have been registered by + a privileged context using the `register` command. + :returns: + Bytestring file data. """ - if mod.__name__ == 'ansible.modules.packaging.os.yum_repository': - # https://github.com/dw/mitogen/issues/154 - mod.YumRepo.repofile = mod.configparser.RawConfigParser() + if path not in _file_cache: + _file_cache[path] = zlib.decompress( + mitogen.service.call( + context, + ansible_mitogen.services.FileService.handle, + ('fetch', path) + ) + ) + return _file_cache[path] -class TemporaryEnvironment(object): - def __init__(self, env=None): - self.original = os.environ.copy() - self.env = env or {} - os.environ.update((k, str(v)) for k, v in self.env.iteritems()) - def revert(self): - os.environ.clear() - os.environ.update(self.original) - - -def run_module(module, raw_params=None, args=None, env=None): +def run_module(kwargs): """ Set up the process environment in preparation for running an Ansible module. This monkey-patches the Ansible libraries in various places to prevent it from trying to kill the process on completion, and to prevent it from reading sys.stdin. """ - if args is None: - args = {} - if raw_params is not None: - args['_raw_params'] = raw_params - - ansible.module_utils.basic.AnsibleModule.exit_json = monkey_exit_json - ansible.module_utils.basic.AnsibleModule.fail_json = monkey_fail_json - ansible.module_utils.basic._ANSIBLE_ARGS = json.dumps({ - 'ANSIBLE_MODULE_ARGS': args - }) + runner_name = kwargs.pop('runner_name') + klass = getattr(ansible_mitogen.runner, runner_name) + impl = klass(**kwargs) + return impl.run() - temp_env = TemporaryEnvironment(env) - try: - try: - mod = __import__(module, {}, {}, ['']) - module_fixups(mod) - # Ansible modules begin execution on import. Thus the above __import__ - # will cause either Exit or ModuleError to be raised. If we reach the - # line below, the module did not execute and must already have been - # imported for a previous invocation, so we need to invoke main - # explicitly. - mod.main() - except (Exit, ModuleError), e: - result = json.dumps(e.dct) - finally: - temp_env.revert() - - return result - - -def _async_main(job_id, module, raw_params, args, env): + +def _async_main(job_id, runner_name, kwargs): """ Implementation for the thread that implements asynchronous module execution. """ try: - rc = run_module(module, raw_params, args, env) + rc = run_module(runner_name, kwargs) except Exception, e: rc = mitogen.core.CallError(e) @@ -189,7 +125,7 @@ def make_temp_directory(base_dir): prefix='ansible-mitogen-tmp-', ) -def run_module_async(module, raw_params=None, args=None): +def run_module_async(runner_name, kwargs): """ Arrange for an Ansible module to be executed in a thread of the current process, with results available via :py:func:`get_async_result`. @@ -200,9 +136,8 @@ def run_module_async(module, raw_params=None, args=None): target=_async_main, kwargs={ 'job_id': job_id, - 'module': module, - 'raw_params': raw_params, - 'args': args, + 'runner_name': runner_name, + 'kwargs': kwargs, } ) _thread_by_job_id[job_id].start() @@ -241,22 +176,23 @@ def get_user_shell(): return pw_shell or '/bin/sh' -def exec_command(cmd, in_data='', chdir=None, shell=None): +def exec_args(args, in_data='', chdir=None, shell=None): """ Run a command in a subprocess, emulating the argument handling behaviour of SSH. - :param bytes cmd: - String command line, passed to user's shell. + :param list[str]: + Argument vector. :param bytes in_data: Optional standard input for the command. :return: (return code, stdout bytes, stderr bytes) """ - assert isinstance(cmd, basestring) + LOG.debug('exec_args(%r, ..., chdir=%r)', args, chdir) + assert isinstance(args, list) proc = subprocess.Popen( - args=[get_user_shell(), '-c', cmd], + args=args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE, @@ -266,6 +202,27 @@ def exec_command(cmd, in_data='', chdir=None, shell=None): return proc.returncode, stdout, stderr +def exec_command(cmd, in_data='', chdir=None, shell=None): + """ + Run a command in a subprocess, emulating the argument handling behaviour of + SSH. + + :param bytes cmd: + String command line, passed to user's shell. + :param bytes in_data: + Optional standard input for the command. + :return: + (return code, stdout bytes, stderr bytes) + """ + assert isinstance(cmd, basestring) + return _exec_command( + args=[get_user_shell(), '-c', cmd], + in_data=in_Data, + chdir=chdir, + shell=shell, + ) + + def read_path(path): """ Fetch the contents of a filesystem `path` as bytes. diff --git a/ansible_mitogen/mixins.py b/ansible_mitogen/mixins.py index 9cef0c3e..56eaa281 100644 --- a/ansible_mitogen/mixins.py +++ b/ansible_mitogen/mixins.py @@ -52,6 +52,7 @@ import mitogen.master from mitogen.utils import cast import ansible_mitogen.connection +import ansible_mitogen.planner import ansible_mitogen.helpers from ansible.module_utils._text import to_text @@ -59,22 +60,6 @@ from ansible.module_utils._text import to_text LOG = logging.getLogger(__name__) -def get_command_module_name(module_name): - """ - Given the name of an Ansible command module, return its canonical module - path within the ansible. - - :param module_name: - "shell" - :return: - "ansible.modules.commands.shell" - """ - path = module_loader.find_plugin(module_name, '') - relpath = os.path.relpath(path, os.path.dirname(ansible.__file__)) - root, _ = os.path.splitext(relpath) - return 'ansible.' + root.replace('/', '.') - - class ActionModuleMixin(ansible.plugins.action.ActionBase): """ The Mitogen-patched PluginLoader dynamically mixes this into every action @@ -199,15 +184,18 @@ class ActionModuleMixin(ansible.plugins.action.ActionBase): """ LOG.debug('_make_tmp_path(remote_user=%r)', remote_user) + try: + remote_tmp = self._connection._shell.get_option('remote_tmp') + except AttributeError: + # Required for <2.4.x. + remote_tmp = '~/.ansible' + # _make_tmp_path() is basically a global stashed away as Shell.tmpdir. # The copy action plugin violates layering and grabs this attribute # directly. self._connection._shell.tmpdir = self.call( ansible_mitogen.helpers.make_temp_directory, - base_dir=self._remote_expand_user( - # ~/.ansible - self._connection._shell.get_option('remote_tmp') - ) + base_dir=self._remote_expand_user(remote_tmp), ) LOG.debug('Temporary directory: %r', self._connection._shell.tmpdir) self._cleanup_remote_tmp = True @@ -308,35 +296,46 @@ class ActionModuleMixin(ansible.plugins.action.ActionBase): helpers.run_module() or helpers.run_module_async() in the target context. """ - if task_vars is None: - task_vars = {} if module_name is None: module_name = self._task.action if module_args is None: module_args = self._task.args + if task_vars is None: + task_vars = {} self._update_module_args(module_name, module_args, task_vars) - if wrap_async: - helper = ansible_mitogen.helpers.run_module_async - else: - helper = ansible_mitogen.helpers.run_module - env = {} self._compute_environment_string(env) - js = self.call( - helper, - get_command_module_name(module_name), - args=cast(module_args), - env=cast(env), + return ansible_mitogen.planner.invoke( + ansible_mitogen.planner.Invocation( + action=self, + connection=self._connection, + module_name=mitogen.utils.cast(module_name), + module_args=mitogen.utils.cast(module_args), + task_vars=task_vars, + templar=self._templar, + env=mitogen.utils.cast(env), + wrap_async=wrap_async, + ) ) - data = self._parse_returned_data({ - 'rc': 0, - 'stdout': js, - 'stdout_lines': [js], - 'stderr': '' - }) + def _postprocess_response(self, result): + """ + Apply fixups mimicking ActionBase._execute_module(); this is copied + verbatim from action/__init__.py, the guts of _parse_returned_data are + garbage and should be removed or reimplemented once tests exist. + + :param dict result: + Dictionary with format:: + + { + "rc": int, + "stdout": "stdout data", + "stderr": "stderr data" + } + """ + data = self._parse_returned_data(result) # Cutpasted from the base implementation. if 'stdout' in data and 'stdout_lines' not in data: @@ -351,8 +350,8 @@ class ActionModuleMixin(ansible.plugins.action.ActionBase): encoding_errors='surrogate_then_replace', chdir=None): """ - Replace the mad rat's nest of logic in the base implementation by - simply calling helpers.exec_command() in the target context. + Override the base implementation by simply calling + helpers.exec_command() in the target context. """ LOG.debug('_low_level_execute_command(%r, in_data=%r, exe=%r, dir=%r)', cmd, type(in_data), executable, chdir) diff --git a/ansible_mitogen/planner.py b/ansible_mitogen/planner.py index 6c9f8441..49c0be06 100644 --- a/ansible_mitogen/planner.py +++ b/ansible_mitogen/planner.py @@ -27,19 +27,98 @@ # POSSIBILITY OF SUCH DAMAGE. """ -This exists to detect every case defined in [0] and prepare arguments necessary -for the executor implementation running within the target, including preloading -any requisite files/Python modules known to be missing. +Classes to detect each case from [0] and prepare arguments necessary for the +corresponding Runner class within the target, including preloading requisite +files/modules known missing. [0] "Ansible Module Architecture", developing_program_flow_modules.html """ from __future__ import absolute_import +import logging +import os + from ansible.executor import module_common +import ansible.errors + +try: + from ansible.plugins.loader import module_loader +except ImportError: # Ansible <2.4 + from ansible.plugins import module_loader import mitogen import mitogen.service import ansible_mitogen.helpers +import ansible_mitogen.services + + +LOG = logging.getLogger(__name__) + + +def parse_script_interpreter(source): + """ + Extract the script interpreter and its sole argument from the module + source code. + + :returns: + Tuple of `(interpreter, arg)`, where `intepreter` is the script + interpreter and `arg` is its solve argument if present, otherwise + :py:data:`None`. + """ + # Linux requires first 2 bytes with no whitespace, pretty sure it's the + # same everywhere. See binfmt_script.c. + if not source.startswith('#!'): + return None, None + + # Find terminating newline. Assume last byte of binprm_buf if absent. + nl = source.find('\n', 0, 128) + if nl == -1: + nl = min(128, len(source)) + + # Split once on the first run of whitespace. If no whitespace exists, + # bits just contains the interpreter filename. + bits = source[2:nl].strip().split(None, 1) + if len(bits) == 1: + return bits[0], None + return bits[0], bits[1] + + +class Invocation(object): + """ + Collect up a module's execution environment then use it to invoke + helpers.run_module() or helpers.run_module_async() in the target context. + """ + def __init__(self, action, connection, module_name, module_args, + task_vars, templar, env, wrap_async): + #: ActionBase instance invoking the module. Required to access some + #: output postprocessing methods that don't belong in ActionBase at + #: all. + self.action = action + #: Ansible connection to use to contact the target. Must be an + #: ansible_mitogen connection. + self.connection = connection + #: Name of the module ('command', 'shell', etc.) to execute. + self.module_name = module_name + #: Final module arguments. + self.module_args = module_args + #: Task variables, needed to extract ansible_*_interpreter. + self.task_vars = task_vars + #: Templar, needed to extract ansible_*_interpreter. + self.templar = templar + #: Final module environment. + self.env = env + #: Boolean, if :py:data:`True`, launch the module asynchronously. + self.wrap_async = wrap_async + + #: Initially ``None``, but set by :func:`invoke`. The path on the + #: master to the module's implementation file. + self.module_path = None + #: Initially ``None``, but set by :func:`invoke`. The raw source or + #: binary contents of the module. + self.module_source = None + + def __repr__(self): + return 'Invocation(module_name=%s)' % (self.module_name,) class Planner(object): @@ -48,57 +127,66 @@ class Planner(object): file, indicates whether or not it understands how to run the module, and exports a method to run the module. """ - def detect(self, name, source): - assert 0 + def detect(self, invocation): + raise NotImplementedError() - def run(self, connection, name, source, args, env): - assert 0 + def plan(self, invocation): + raise NotImplementedError() -class JsonArgsPlanner(Planner): +class BinaryPlanner(Planner): """ - Script that has its interpreter directive and the task arguments - substituted into its source as a JSON string. + Binary modules take their arguments and will return data to Ansible in the + same way as want JSON modules. """ - def detect(self, name, source): - return module_common.REPLACER_JSONARGS in source - - def run(self, name, source, args, env): - path = None # TODO - mitogen.service.call(501, ('register', path)) + runner_name = 'BinaryRunner' + + def detect(self, invocation): + return module_common._is_binary(invocation.module_source) + + def plan(self, invocation): + invocation.connection._connect() + mitogen.service.call( + invocation.connection.parent, + ansible_mitogen.services.FileService.handle, + ('register', invocation.module_path) + ) return { - 'func': 'run_json_args_module', - 'binary': source, - 'args': args, - 'env': env, + 'runner_name': self.runner_name, + 'module': invocation.module_name, + 'service_context': invocation.connection.parent, + 'path': invocation.module_path, + 'args': invocation.module_args, + 'env': invocation.env, } -class WantJsonPlanner(Planner): +class ScriptPlanner(BinaryPlanner): """ - If a module has the string WANT_JSON in it anywhere, Ansible treats it as a - non-native module that accepts a filename as its only command line - parameter. The filename is for a temporary file containing a JSON string - containing the module’s parameters. The module needs to open the file, read - and parse the parameters, operate on the data, and print its return data as - a JSON encoded dictionary to stdout before exiting. - - These types of modules are self-contained entities. As of Ansible 2.1, - Ansible only modifies them to change a shebang line if present. + Common functionality for script module planners -- handle interpreter + detection and rewrite. """ - def detect(self, name, source): - return 'WANT_JSON' in source - - def run(self, name, source, args, env): - return { - 'func': 'run_want_json_module', - 'binary': source, - 'args': args, - 'env': env, - } - - -class ReplacerPlanner(Planner): + def _rewrite_interpreter(self, interpreter, task_vars, templar): + key = u'ansible_%s_interpreter' % os.path.basename(interpreter).strip() + try: + return templar.template(task_vars[key].strip()) + except KeyError: + return interpreter + + def plan(self, invocation): + kwargs = super(ScriptPlanner, self).plan(invocation) + interpreter, arg = parse_script_interpreter(invocation.module_source) + return dict(kwargs, + interpreter_arg=arg, + interpreter=self._rewrite_interpreter( + interpreter=interpreter, + task_vars=invocation.task_vars, + templar=invocation.templar, + ) + ) + + +class ReplacerPlanner(BinaryPlanner): """ The Module Replacer framework is the original framework implementing new-style modules. It is essentially a preprocessor (like the C @@ -122,65 +210,114 @@ class ReplacerPlanner(Planner): "ansible/module_utils/powershell.ps1". It should only be used with new-style Powershell modules. """ - def detect(self, name, source): - return module_common.REPLACER in source + runner_name = 'ReplacerRunner' - def run(self, name, source, args, env): - return { - 'func': 'run_replacer_module', - 'binary': source, - 'args': args, - 'env': env, - } + def detect(self, invocation): + return module_common.REPLACER in invocation.module_source -class BinaryPlanner(Planner): +class JsonArgsPlanner(ScriptPlanner): """ - Binary modules take their arguments and will return data to Ansible in the - same way as want JSON modules. + Script that has its interpreter directive and the task arguments + substituted into its source as a JSON string. """ - helper = staticmethod(ansible_mitogen.helpers.run_binary) + runner_name = 'JsonArgsRunner' - def detect(self, name, source): - return module_common._is_binary(source) + def detect(self, invocation): + return module_common.REPLACER_JSONARGS in invocation.module_source - def run(self, name, source, args, env): - return { - 'func': 'run_binary_module', - 'binary': source, - 'args': args, - 'env': env, - } +class WantJsonPlanner(ScriptPlanner): + """ + If a module has the string WANT_JSON in it anywhere, Ansible treats it as a + non-native module that accepts a filename as its only command line + parameter. The filename is for a temporary file containing a JSON string + containing the module's parameters. The module needs to open the file, read + and parse the parameters, operate on the data, and print its return data as + a JSON encoded dictionary to stdout before exiting. + + These types of modules are self-contained entities. As of Ansible 2.1, + Ansible only modifies them to change a shebang line if present. + """ + runner_name = 'WantJsonRunner' + + def detect(self, invocation): + return 'WANT_JSON' in invocation.module_source -class PythonPlanner(Planner): + +class NewStylePlanner(ScriptPlanner): """ The Ansiballz framework differs from module replacer in that it uses real Python imports of things in ansible/module_utils instead of merely preprocessing the module. """ - helper = staticmethod(ansible_mitogen.helpers.run_module) + runner_name = 'NewStyleRunner' - def detect(self, name, source): - return True + def detect(self, invocation): + return 'from ansible.module_utils.' in invocation.module_source - def run(self, name, source, args, env): - return { - 'func': 'run_python_module', - 'module': name, - 'args': args, - 'env': env - } + +class ReplacerPlanner(NewStylePlanner): + runner_name = 'ReplacerRunner' + + def detect(self, invocation): + return module_common.REPLACER in invocation.module_source + + +class OldStylePlanner(ScriptPlanner): + runner_name = 'OldStyleRunner' + + def detect(self, invocation): + # Everything else. + return True _planners = [ - # JsonArgsPlanner, - # WantJsonPlanner, - # ReplacerPlanner, BinaryPlanner, - PythonPlanner, + # ReplacerPlanner, + NewStylePlanner, + JsonArgsPlanner, + WantJsonPlanner, + OldStylePlanner, ] -def plan(): - pass +NO_METHOD_MSG = 'Mitogen: no invocation method found for: ' +CRASHED_MSG = 'Mitogen: internal error: ' + + +def get_module_data(name): + path = module_loader.find_plugin(name, '') + with open(path, 'rb') as fp: + source = fp.read() + return path, source + + +def invoke(invocation): + """ + Find a suitable Planner that knows how to run `invocation`. + """ + (invocation.module_path, + invocation.module_source) = get_module_data(invocation.module_name) + + for klass in _planners: + planner = klass() + if planner.detect(invocation): + break + else: + raise ansible.errors.AnsibleError(NO_METHOD_MSG + repr(invocation)) + + kwargs = planner.plan(invocation) + if invocation.wrap_async: + helper = ansible_mitogen.helpers.run_module_async + else: + helper = ansible_mitogen.helpers.run_module + + try: + js = invocation.connection.call(helper, kwargs) + except mitogen.core.CallError as e: + LOG.exception('invocation crashed: %r', invocation) + summary = str(e).splitlines()[0] + raise ansible.errors.AnsibleInternalError(CRASHED_MSG + summary) + + return invocation.action._postprocess_response(js) diff --git a/ansible_mitogen/process.py b/ansible_mitogen/process.py index eb9bd2ff..8febea90 100644 --- a/ansible_mitogen/process.py +++ b/ansible_mitogen/process.py @@ -153,7 +153,8 @@ class MuxProcess(object): self.pool = mitogen.service.Pool( router=self.router, services=[ - ansible_mitogen.services.ContextService(self.router) + ansible_mitogen.services.ContextService(self.router), + ansible_mitogen.services.FileService(self.router), ], size=int(os.environ.get('MITOGEN_POOL_SIZE', '16')), ) diff --git a/ansible_mitogen/runner.py b/ansible_mitogen/runner.py new file mode 100644 index 00000000..f5cb1283 --- /dev/null +++ b/ansible_mitogen/runner.py @@ -0,0 +1,380 @@ +# Copyright 2017, David Wilson +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software without +# specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. + + +""" +These classes implement execution for each style of Ansible module. They are +instantiated in the target context by way of helpers.py::run_module(). + +Each class in here has a corresponding Planner class in planners.py that knows +how to build arguments for it, preseed related data, etc. +""" + +from __future__ import absolute_import +import cStringIO +import json +import logging +import os +import sys +import tempfile +import types + +import ansible_mitogen.helpers # TODO: circular import + +try: + from shlex import quote as shlex_quote +except ImportError: + from pipes import quote as shlex_quote + +# Prevent accidental import of an Ansible module from hanging on stdin read. +import ansible.module_utils.basic +ansible.module_utils.basic._ANSIBLE_ARGS = '{}' + + +LOG = logging.getLogger(__name__) + + +class Runner(object): + """ + Ansible module runner. After instantiation (with kwargs supplied by the + corresponding Planner), `.run()` is invoked, upon which `setup()`, + `_run()`, and `revert()` are invoked, with the return value of `_run()` + returned by `run()`. + + Subclasses may override `_run`()` and extend `setup()` and `revert()`. + """ + def __init__(self, module, raw_params=None, args=None, env=None): + if args is None: + args = {} + if raw_params is not None: + args['_raw_params'] = raw_params + + self.module = module + self.raw_params = raw_params + self.args = args + self.env = env + + def setup(self): + """ + Prepare the current process for running a module. The base + implementation simply prepares the environment. + """ + self._env = TemporaryEnvironment(self.env) + + def revert(self): + """ + Revert any changes made to the process after running a module. The base + implementation simply restores the original environment. + """ + self._env.revert() + + def _run(self): + """ + The _run() method is expected to return a dictionary in the form of + ActionBase._low_level_execute_command() output, i.e. having:: + + { + "rc": int, + "stdout": "stdout data", + "stderr": "stderr data" + } + """ + raise NotImplementedError() + + def run(self): + """ + Set up the process environment in preparation for running an Ansible + module. This monkey-patches the Ansible libraries in various places to + prevent it from trying to kill the process on completion, and to + prevent it from reading sys.stdin. + + :returns: + Module result dictionary. + """ + self.setup() + try: + return self._run() + finally: + self.revert() + + +class TemporaryEnvironment(object): + def __init__(self, env=None): + self.original = os.environ.copy() + self.env = env or {} + os.environ.update((k, str(v)) for k, v in self.env.iteritems()) + + def revert(self): + os.environ.clear() + os.environ.update(self.original) + + +class TemporaryArgv(object): + def __init__(self, argv): + self.original = sys.argv[:] + sys.argv[:] = argv + + def revert(self): + sys.argv[:] = self.original + + +class NewStyleStdio(object): + """ + Patch ansible.module_utils.basic argument globals. + """ + def __init__(self, args): + self.original_stdout = sys.stdout + self.original_stderr = sys.stderr + self.original_stdin = sys.stdin + sys.stdout = cStringIO.StringIO() + sys.stderr = cStringIO.StringIO() + ansible.module_utils.basic._ANSIBLE_ARGS = json.dumps({ + 'ANSIBLE_MODULE_ARGS': args + }) + sys.stdin = cStringIO.StringIO( + ansible.module_utils.basic._ANSIBLE_ARGS + ) + + def revert(self): + sys.stdout = self.original_stdout + sys.stderr = self.original_stderr + sys.stdin = self.original_stdin + ansible.module_utils.basic._ANSIBLE_ARGS = '{}' + + +class ProgramRunner(Runner): + def __init__(self, path, service_context, **kwargs): + super(ProgramRunner, self).__init__(**kwargs) + self.path = path + self.service_context = service_context + + def setup(self): + super(ProgramRunner, self).setup() + self._setup_program() + + def _setup_program(self): + """ + Create a temporary file containing the program code. The code is + fetched via :meth:`_get_program`. + """ + self.program_fp = tempfile.NamedTemporaryFile( + prefix='ansible_mitogen', + suffix='-binary', + ) + self.program_fp.write(self._get_program()) + self.program_fp.flush() + os.chmod(self.program_fp.name, int('0700', 8)) + + def _get_program(self): + """ + Fetch the module binary from the master if necessary. + """ + return ansible_mitogen.helpers.get_file( + context=self.service_context, + path=self.path, + ) + + def _get_program_args(self): + return [self.program_fp.name] + + def revert(self): + """ + Delete the temporary program file. + """ + super(ProgramRunner, self).revert() + self.program_fp.close() + + def _run(self): + try: + rc, stdout, stderr = ansible_mitogen.helpers.exec_args( + args=self._get_program_args(), + ) + except Exception, e: + return { + 'rc': 1, + 'stdout': '', + 'stderr': '%s: %s' % (type(e), e), + } + + return { + 'rc': rc, + 'stdout': stdout, + 'stderr': stderr + } + + +class ArgsFileRunner(Runner): + def setup(self): + super(ArgsFileRunner, self).setup() + self._setup_args() + + def _setup_args(self): + """ + Create a temporary file containing the module's arguments. The + arguments are formatted via :meth:`_get_args`. + """ + self.args_fp = tempfile.NamedTemporaryFile( + prefix='ansible_mitogen', + suffix='-args', + ) + self.args_fp.write(self._get_args_contents()) + self.args_fp.flush() + + def _get_args_contents(self): + """ + Return the module arguments formatted as JSON. + """ + return json.dumps(self.args) + + def _get_program_args(self): + return [self.program_fp.name, self.args_fp.name] + + def revert(self): + """ + Delete the temporary argument file. + """ + super(ArgsFileRunner, self).revert() + self.args_fp.close() + + +class BinaryRunner(ArgsFileRunner, ProgramRunner): + pass + + +class ScriptRunner(ProgramRunner): + def __init__(self, interpreter, interpreter_arg, **kwargs): + super(ScriptRunner, self).__init__(**kwargs) + self.interpreter = interpreter + self.interpreter_arg = interpreter_arg + + b_ENCODING_STRING = b'# -*- coding: utf-8 -*-' + + def _get_program(self): + return self._rewrite_source( + super(ScriptRunner, self)._get_program() + ) + + def _rewrite_source(self, s): + """ + Mutate the source according to the per-task parameters. + """ + # Couldn't find shebang, so let shell run it, because shell assumes + # executables like this are just shell scripts. + if not self.interpreter: + return s + + shebang = '#!' + self.interpreter + if self.interpreter_arg: + shebang += ' ' + self.interpreter_arg + + new = [shebang] + if os.path.basename(self.interpreter).startswith('python'): + new.append(self.b_ENCODING_STRING) + + _, _, rest = s.partition('\n') + new.append(rest) + return '\n'.join(new) + + +class NewStyleRunner(ScriptRunner): + """ + Execute a new-style Ansible module, where Module Replacer-related tricks + aren't required. + """ + #: path => new-style module bytecode. + _code_by_path = {} + + def setup(self): + super(NewStyleRunner, self).setup() + self._stdio = NewStyleStdio(self.args) + self._argv = TemporaryArgv([self.path]) + + def revert(self): + self._argv.revert() + self._stdio.revert() + super(NewStyleRunner, self).revert() + + def _get_code(self): + try: + return self._code_by_path[self.path] + except KeyError: + return self._code_by_path.setdefault(self.path, compile( + source=ansible_mitogen.helpers.get_file( + context=self.service_context, + path=self.path, + ), + filename=self.path, + mode='exec', + dont_inherit=True, + )) + + def _run(self): + code = self._get_code() + mod = types.ModuleType('__main__') + d = vars(mod) + e = None + + try: + exec code in d, d + except SystemExit, e: + pass + + return { + 'rc': e[0] if e else 2, + 'stdout': sys.stdout.getvalue(), + 'stderr': sys.stderr.getvalue(), + } + + +class JsonArgsRunner(ScriptRunner): + JSON_ARGS = '<>' + + def _get_args_contents(self): + return json.dumps(self.args) + + def _rewrite_source(self, s): + return ( + super(JsonArgsRunner, self)._rewrite_source(s) + .replace(self.JSON_ARGS, self._get_args_contents()) + ) + + +class WantJsonRunner(ArgsFileRunner, ScriptRunner): + pass + + +class OldStyleRunner(ArgsFileRunner, ScriptRunner): + def _get_args_contents(self): + """ + Mimic the argument formatting behaviour of + ActionBase._execute_module(). + """ + return ' '.join( + '%s=%s' % (key, shlex_quote(str(self.args[key]))) + for key in self.args + ) + ' ' # Bug-for-bug :( diff --git a/ansible_mitogen/services.py b/ansible_mitogen/services.py index 0266d533..7bc59135 100644 --- a/ansible_mitogen/services.py +++ b/ansible_mitogen/services.py @@ -112,6 +112,9 @@ class FileService(mitogen.service.Service): """ handle = 501 max_message_size = 1000 + policies = ( + mitogen.service.AllowAny(), + ) unprivileged_msg = 'Cannot register from unprivileged context.' unregistered_msg = 'Path is not registered with FileService.' @@ -125,17 +128,21 @@ class FileService(mitogen.service.Service): isinstance(args, tuple) and len(args) == 2 and args[0] in ('register', 'fetch') and - isinstance(args[1], str) + isinstance(args[1], basestring) ) def dispatch(self, args, msg): - cmd, path = msg + cmd, path = args return getattr(self, cmd)(path, msg) def register(self, path, msg): - if msg.auth_id not in mitogen.parent_ids: + if not mitogen.core.has_parent_authority(msg): raise mitogen.core.CallError(self.unprivileged_msg) + if path in self._paths: + return + + LOG.info('%r: registering %r', self, path) with open(path, 'rb') as fp: self._paths[path] = zlib.compress(fp.read()) diff --git a/examples/playbook/Makefile b/examples/playbook/Makefile new file mode 100644 index 00000000..8c9e4480 --- /dev/null +++ b/examples/playbook/Makefile @@ -0,0 +1,4 @@ + +all: \ + modules/binary_producing_junk \ + modules/binary_producing_json diff --git a/examples/playbook/modules/bin_bash_module.sh b/examples/playbook/modules/bin_bash_module.sh new file mode 100755 index 00000000..01abff0c --- /dev/null +++ b/examples/playbook/modules/bin_bash_module.sh @@ -0,0 +1,5 @@ +#!/bin/bash +exec >/tmp/derp +echo "$1" +cat "$1" + diff --git a/examples/playbook/modules/binary_producing_json.c b/examples/playbook/modules/binary_producing_json.c new file mode 100644 index 00000000..989e5e3e --- /dev/null +++ b/examples/playbook/modules/binary_producing_json.c @@ -0,0 +1,13 @@ +#include + + +int main(void) +{ + fprintf(stderr, "binary_producing_json: oh noes\n"); + printf("{" + "\"changed\": true, " + "\"failed\": false, " + "\"msg\": \"Hello, world.\"" + "}\n"); + return 0; +} diff --git a/examples/playbook/modules/binary_producing_junk.c b/examples/playbook/modules/binary_producing_junk.c new file mode 100644 index 00000000..f6b68462 --- /dev/null +++ b/examples/playbook/modules/binary_producing_junk.c @@ -0,0 +1,9 @@ +#include + + +int main(void) +{ + fprintf(stderr, "binary_producing_junk: oh noes\n"); + printf("Hello, world.\n"); + return 0; +} diff --git a/examples/playbook/modules/json_args_python.py b/examples/playbook/modules/json_args_python.py new file mode 100644 index 00000000..45689584 --- /dev/null +++ b/examples/playbook/modules/json_args_python.py @@ -0,0 +1,13 @@ +#!/usr/bin/python +# I am an Ansible Python JSONARGS module. I should receive an encoding string. + +import json +import sys + +json_arguments = """<>""" + +print "{" +print " \"changed\": false," +print " \"msg\": \"Here is my input\"," +print " \"input\": [%s]" % (json_arguments,) +print "}" diff --git a/examples/playbook/modules/old_style_module.sh b/examples/playbook/modules/old_style_module.sh new file mode 100755 index 00000000..47e6afbd --- /dev/null +++ b/examples/playbook/modules/old_style_module.sh @@ -0,0 +1,16 @@ +#!/bin/bash +# I am an Ansible old-style module. + +INPUT=$1 + +[ ! -r "$INPUT" ] && { + echo "Usage: $0 " >&2 + exit 1 +} + +echo "{" +echo " \"changed\": false," +echo " \"msg\": \"Here is my input\"," +echo " \"filname\": \"$INPUT\"," +echo " \"input\": [\"$(cat $INPUT | tr \" \' )\"]" +echo "}" diff --git a/examples/playbook/modules/python_new_style_module.py b/examples/playbook/modules/python_new_style_module.py new file mode 100755 index 00000000..1ae50d50 --- /dev/null +++ b/examples/playbook/modules/python_new_style_module.py @@ -0,0 +1,27 @@ +#!/usr/bin/python +# I am an Ansible new-style Python module. I should receive an encoding string. + +import json +import sys + +# This is the magic marker Ansible looks for: +# from ansible.module_utils. + + +def usage(): + sys.stderr.write('Usage: %s \n' % (sys.argv[0],)) + sys.exit(1) + +# Also must slurp in our own source code, to verify the encoding string was +# added. +with open(sys.argv[0]) as fp: + me = fp.read() + +input_json = sys.stdin.read() + +print "{" +print " \"changed\": false," +print " \"msg\": \"Here is my input\"," +print " \"source\": [%s]," % (json.dumps(me),) +print " \"input\": [%s]" % (input_json,) +print "}" diff --git a/examples/playbook/modules/python_want_json_module.py b/examples/playbook/modules/python_want_json_module.py new file mode 100755 index 00000000..bd12704e --- /dev/null +++ b/examples/playbook/modules/python_want_json_module.py @@ -0,0 +1,33 @@ +#!/usr/bin/python +# I am an Ansible Python WANT_JSON module. I should receive an encoding string. + +import json +import sys + +WANT_JSON = 1 + + +def usage(): + sys.stderr.write('Usage: %s \n' % (sys.argv[0],)) + sys.exit(1) + +if len(sys.argv) < 2: + usage() + +# Also must slurp in our own source code, to verify the encoding string was +# added. +with open(sys.argv[0]) as fp: + me = fp.read() + +try: + with open(sys.argv[1]) as fp: + input_json = fp.read() +except IOError: + usage() + +print "{" +print " \"changed\": false," +print " \"msg\": \"Here is my input\"," +print " \"source\": [%s]," % (json.dumps(me),) +print " \"input\": [%s]" % (input_json,) +print "}" diff --git a/examples/playbook/modules/single_null_binary b/examples/playbook/modules/single_null_binary new file mode 100755 index 00000000..1f2a4f5e Binary files /dev/null and b/examples/playbook/modules/single_null_binary differ diff --git a/examples/playbook/modules/want_json_module.sh b/examples/playbook/modules/want_json_module.sh new file mode 100755 index 00000000..6053eacd --- /dev/null +++ b/examples/playbook/modules/want_json_module.sh @@ -0,0 +1,16 @@ +#!/bin/bash +# I am an Ansible WANT_JSON module. + +WANT_JSON=1 +INPUT=$1 + +[ ! -r "$INPUT" ] && { + echo "Usage: $0 " >&2 + exit 1 +} + +echo "{" +echo " \"changed\": false," +echo " \"msg\": \"Here is my input\"," +echo " \"input\": [$(< $INPUT)]" +echo "}" diff --git a/examples/playbook/regtest.py b/examples/playbook/regtest.py new file mode 100644 index 00000000..15fc9e3c --- /dev/null +++ b/examples/playbook/regtest.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python + +import difflib +import logging +import re +import subprocess +import tempfile + + +LOG = logging.getLogger(__name__) + +suffixes = [ + '-m bin_bash_module', + '-m binary_producing_json', + '-m binary_producing_junk', + '-m old_style_module', + '-m python_new_style_module', + '-m python_want_json_module', + '-m single_null_binary', + '-m want_json_module', + '-m json_args_python', + '-m setup', +] + +fixups = [ + ('Shared connection to localhost closed\\.(\r\n)?', ''), # TODO +] + + +def fixup(s): + for regex, to in fixups: + s = re.sub(regex, to, s, re.DOTALL|re.M) + return s + + +def run(s): + LOG.debug('running: %r', s) + with tempfile.NamedTemporaryFile() as fp: + # https://www.systutorials.com/docs/linux/man/1-ansible-playbook/#lbAG + returncode = subprocess.call(s, stdout=fp, stderr=fp, shell=True) + fp.write('\nReturn code: %s\n' % (returncode,)) + fp.seek(0) + return fp.read() + + +logging.basicConfig(level=logging.DEBUG) + +for suffix in suffixes: + ansible = run('ansible localhost %s' % (suffix,)) + mitogen = run('ANSIBLE_STRATEGY=mitogen ansible localhost %s' % (suffix,)) + + diff = list(difflib.unified_diff( + a=fixup(ansible).splitlines(), + b=fixup(mitogen).splitlines(), + fromfile='ansible-output.txt', + tofile='mitogen-output.txt', + )) + if diff: + print '++ differ! suffix: %r' % (suffix,) + for line in diff: + print line + print + print diff --git a/mitogen/core.py b/mitogen/core.py index beeedd63..d60fbafc 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -157,8 +157,9 @@ def _unpickle_dead(): _DEAD = Dead() -def has_parent_authority(msg, _stream): - return msg.auth_id in mitogen.parent_ids +def has_parent_authority(msg, _stream=None): + return (msg.auth_id == mitogen.context_id or + msg.auth_id in mitogen.parent_ids) def listen(obj, name, func): @@ -491,6 +492,7 @@ class Importer(object): 'fork', 'master', 'parent', + 'service', 'ssh', 'sudo', 'utils', @@ -955,8 +957,10 @@ class Context(object): def _unpickle_context(router, context_id, name): if not (isinstance(router, Router) and - isinstance(context_id, (int, long)) and context_id >= 0 and - isinstance(name, basestring) and len(name) < 100): + isinstance(context_id, (int, long)) and context_id >= 0 and ( + (name is None) or + (isinstance(name, basestring) and len(name) < 100)) + ): raise TypeError('cannot unpickle Context: bad input') return router.context_class(router, context_id, name) diff --git a/mitogen/service.py b/mitogen/service.py index b9fad1fe..018fa17d 100644 --- a/mitogen/service.py +++ b/mitogen/service.py @@ -82,6 +82,12 @@ class Service(object): self.handle = self.recv.handle self.running = True + def __repr__(self): + return '%s.%s()' % ( + self.__class__.__module__, + self.__class__.__name__, + ) + def validate_args(self, args): return ( isinstance(args, dict) and @@ -108,6 +114,7 @@ class Service(object): isinstance(args, mitogen.core.CallError) or not self.validate_args(args)): LOG.warning('Received junk message: %r', args) + msg.reply(mitogen.core.CallError('Received junk message')) return try: diff --git a/tests/io_op_test.py b/tests/io_op_test.py new file mode 100644 index 00000000..8ec204b6 --- /dev/null +++ b/tests/io_op_test.py @@ -0,0 +1,119 @@ + +import errno +import select + +import mock +import unittest2 + +import testlib +import mitogen.core + + +class RestartTest(object): + func = staticmethod(mitogen.core.io_op) + exception_class = None + + def test_eintr_restarts(self): + m = mock.Mock() + m.side_effect = [ + self.exception_class(errno.EINTR), + self.exception_class(errno.EINTR), + self.exception_class(errno.EINTR), + 'yay', + ] + rc, disconnected = self.func(m, 'input') + self.assertEquals(rc, 'yay') + self.assertFalse(disconnected) + self.assertEquals(4, m.call_count) + self.assertEquals(m.mock_calls, [ + mock.call('input'), + mock.call('input'), + mock.call('input'), + mock.call('input'), + ]) + + +class SelectRestartTest(RestartTest, testlib.TestCase): + exception_class = select.error + + +class OsErrorRestartTest(RestartTest, testlib.TestCase): + exception_class = OSError + + +class DisconnectTest(object): + func = staticmethod(mitogen.core.io_op) + errno = None + exception_class = None + + def test_disconnection(self): + m = mock.Mock() + m.side_effect = self.exception_class(self.errno) + rc, disconnected = self.func(m, 'input') + self.assertEquals(rc, None) + self.assertTrue(disconnected) + self.assertEquals(1, m.call_count) + self.assertEquals(m.mock_calls, [ + mock.call('input'), + ]) + + +class SelectDisconnectEioTest(DisconnectTest, testlib.TestCase): + errno = errno.EIO + exception_class = select.error + + +class SelectDisconnectEconnresetTest(DisconnectTest, testlib.TestCase): + errno = errno.ECONNRESET + exception_class = select.error + + +class SelectDisconnectEpipeTest(DisconnectTest, testlib.TestCase): + errno = errno.EPIPE + exception_class = select.error + + +class OsErrorDisconnectEioTest(DisconnectTest, testlib.TestCase): + errno = errno.EIO + exception_class = OSError + + +class OsErrorDisconnectEconnresetTest(DisconnectTest, testlib.TestCase): + errno = errno.ECONNRESET + exception_class = OSError + + +class OsErrorDisconnectEpipeTest(DisconnectTest, testlib.TestCase): + errno = errno.EPIPE + exception_class = OSError + + +class ExceptionTest(object): + func = staticmethod(mitogen.core.io_op) + errno = None + exception_class = None + + def test_exception(self): + m = mock.Mock() + m.side_effect = self.exception_class(self.errno) + e = self.assertRaises(self.exception_class, + lambda: self.func(m, 'input')) + self.assertEquals(e, m.side_effect) + self.assertEquals(1, m.call_count) + self.assertEquals(m.mock_calls, [ + mock.call('input'), + ]) + + +class SelectExceptionTest(ExceptionTest, testlib.TestCase): + errno = errno.EBADF + exception_class = select.error + + +class OsErrorExceptionTest(ExceptionTest, testlib.TestCase): + errno = errno.EBADF + exception_class = OSError + + +if __name__ == '__main__': + unittest2.main()