From d8e0c9e12c16ea26fa9bad61ef1dfb0a32bfffe2 Mon Sep 17 00:00:00 2001 From: David Wilson Date: Tue, 10 Jul 2018 18:29:49 +0100 Subject: [PATCH] issue #297: local commands must execute with WorkerProcess environment. --- ansible_mitogen/connection.py | 9 ++++- ansible_mitogen/planner.py | 1 + .../plugins/connection/mitogen_local.py | 36 +++++++++++++++++++ ansible_mitogen/process.py | 7 ++++ ansible_mitogen/runner.py | 28 +++++++++++---- tests/ansible/ansible.cfg | 1 + tests/ansible/integration/local/all.yml | 2 +- .../integration/local/env_preserved.yml | 8 +++++ .../lib/vars/custom_modifies_os_environ.py | 16 +++++++++ 9 files changed, 99 insertions(+), 9 deletions(-) create mode 100644 tests/ansible/integration/local/env_preserved.yml create mode 100644 tests/ansible/lib/vars/custom_modifies_os_environ.py diff --git a/ansible_mitogen/connection.py b/ansible_mitogen/connection.py index 5b9407fe..c8b1870c 100644 --- a/ansible_mitogen/connection.py +++ b/ansible_mitogen/connection.py @@ -550,7 +550,14 @@ class Connection(ansible.plugins.connection.ConnectionBase): def get_default_cwd(self): """ Overridden by connections/mitogen_local.py to emulate behaviour of CWD - inherited from WorkerProcess. + being fixed to that of ActionBase._loader.get_basedir(). + """ + return None + + def get_default_env(self): + """ + Overridden by connections/mitogen_local.py to emulate behaviour of + WorkProcess environment inherited from WorkerProcess. """ return None diff --git a/ansible_mitogen/planner.py b/ansible_mitogen/planner.py index edf1101e..3542756b 100644 --- a/ansible_mitogen/planner.py +++ b/ansible_mitogen/planner.py @@ -176,6 +176,7 @@ class Planner(object): new = dict((mitogen.core.UnicodeType(k), kwargs[k]) for k in kwargs) new.setdefault('cwd', self._inv.connection.get_default_cwd()) + new.setdefault('extra_env', self._inv.connection.get_default_env()) new.setdefault('emulate_tty', True) new.setdefault('service_context', self._inv.connection.parent) return new diff --git a/ansible_mitogen/plugins/connection/mitogen_local.py b/ansible_mitogen/plugins/connection/mitogen_local.py index 16a69364..c33cd64d 100644 --- a/ansible_mitogen/plugins/connection/mitogen_local.py +++ b/ansible_mitogen/plugins/connection/mitogen_local.py @@ -37,6 +37,32 @@ except ImportError: del base_dir import ansible_mitogen.connection +import ansible_mitogen.process + + +if sys.version_info > (2, 7): + viewkeys = dict.viewkeys +else: + viewkeys = lambda dct: set(dct) + + +def dict_diff(old, new): + """ + Return a dict representing the differences between the dicts `old` and + `new`. Deleted keys appear as a key with the value :data:`None`, added and + changed keys appear as a key with the new value. + """ + old_keys = viewkeys(old) + new_keys = viewkeys(dict(new)) + out = {} + for key in new_keys - old_keys: + out[key] = new[key] + for key in old_keys - new_keys: + out[key] = None + for key in old_keys & new_keys: + if old[key] != new[key]: + out[key] = new[key] + return out class Connection(ansible_mitogen.connection.Connection): @@ -45,3 +71,13 @@ class Connection(ansible_mitogen.connection.Connection): def get_default_cwd(self): # https://github.com/ansible/ansible/issues/14489 return self.loader_basedir + + def get_default_env(self): + """ + Vanilla Ansible local commands execute with an environment inherited + from WorkerProcess, we must emulate that. + """ + return dict_diff( + old=ansible_mitogen.process.MuxProcess.original_env, + new=os.environ, + ) diff --git a/ansible_mitogen/process.py b/ansible_mitogen/process.py index 69e70571..f19079ee 100644 --- a/ansible_mitogen/process.py +++ b/ansible_mitogen/process.py @@ -84,6 +84,12 @@ class MuxProcess(object): #: that was spawned. worker_pid = None + #: A copy of :data:`os.environ` at the time the multiplexer process was + #: started. It's used by mitogen_local.py to find changes made to the + #: top-level environment (e.g. vars plugins -- issue #297) that must be + #: applied to locally executed commands and modules. + original_env = None + #: In both processes, this is the temporary UNIX socket used for #: forked WorkerProcesses to contact the MuxProcess unix_listener_path = None @@ -108,6 +114,7 @@ class MuxProcess(object): mitogen.core.set_cloexec(cls.worker_sock.fileno()) mitogen.core.set_cloexec(cls.child_sock.fileno()) + cls.original_env = dict(os.environ) cls.child_pid = os.fork() ansible_mitogen.logging.setup() if cls.child_pid: diff --git a/ansible_mitogen/runner.py b/ansible_mitogen/runner.py index baace940..00ebccf3 100644 --- a/ansible_mitogen/runner.py +++ b/ansible_mitogen/runner.py @@ -137,13 +137,14 @@ class Runner(object): When :data:`True`, indicate the runner should detach the context from its parent after setup has completed successfully. """ - def __init__(self, module, service_context, json_args, cwd=None, env=None, - econtext=None, detach=False): + def __init__(self, module, service_context, json_args, extra_env=None, + cwd=None, env=None, econtext=None, detach=False): self.module = module self.service_context = service_context self.econtext = econtext self.detach = detach self.args = json.loads(json_args) + self.extra_env = extra_env self.env = env self.cwd = cwd @@ -154,7 +155,10 @@ class Runner(object): execution. The base implementation simply prepares the environment. """ self._cwd = TemporaryCwd(self.cwd) - self._env = TemporaryEnvironment(self.env) + env = dict(self.extra_env or {}) + if self.env: + env.update(self.env) + self._env = TemporaryEnvironment(env) def revert(self): """ @@ -275,14 +279,24 @@ class TemporaryCwd(object): class TemporaryEnvironment(object): + """ + Apply environment changes from `env` until :meth:`revert` is called. Values + in the dict may be :data:`None` to indicate the relevant key should be + deleted. + """ def __init__(self, env=None): - self.original = os.environ.copy() + self.original = dict(os.environ) self.env = env or {} - os.environ.update((k, str(v)) for k, v in iteritems(self.env)) + for key, value in iteritems(self.env): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = str(value) def revert(self): - os.environ.clear() - os.environ.update(self.original) + if self.env: + os.environ.clear() + os.environ.update(self.original) class TemporaryArgv(object): diff --git a/tests/ansible/ansible.cfg b/tests/ansible/ansible.cfg index eeb82109..e00cc974 100644 --- a/tests/ansible/ansible.cfg +++ b/tests/ansible/ansible.cfg @@ -5,6 +5,7 @@ strategy_plugins = ../../ansible_mitogen/plugins/strategy action_plugins = lib/action callback_plugins = lib/callback stdout_callback = nice_stdout +vars_plugins = lib/vars library = lib/modules # module_utils = lib/module_utils retry_files_enabled = False diff --git a/tests/ansible/integration/local/all.yml b/tests/ansible/integration/local/all.yml index 2ca1660d..383a9108 100644 --- a/tests/ansible/integration/local/all.yml +++ b/tests/ansible/integration/local/all.yml @@ -1,4 +1,4 @@ - import_playbook: cwd_preserved.yml -#- import_playbook: env_preserved.yml +- import_playbook: env_preserved.yml diff --git a/tests/ansible/integration/local/env_preserved.yml b/tests/ansible/integration/local/env_preserved.yml new file mode 100644 index 00000000..1375391e --- /dev/null +++ b/tests/ansible/integration/local/env_preserved.yml @@ -0,0 +1,8 @@ + +# Execution environment should be that of WorkerProcess -- +# https://github.com/dw/mitogen/issues/297 + +- hosts: localhost + connection: local + tasks: + - shell: "env | grep EVIL_VARS_PLUGIN" diff --git a/tests/ansible/lib/vars/custom_modifies_os_environ.py b/tests/ansible/lib/vars/custom_modifies_os_environ.py new file mode 100644 index 00000000..4039dfa7 --- /dev/null +++ b/tests/ansible/lib/vars/custom_modifies_os_environ.py @@ -0,0 +1,16 @@ +# https://github.com/dw/mitogen/issues/297 + +from __future__ import (absolute_import, division, print_function) +__metaclass__ = type + +from ansible.plugins.vars import BaseVarsPlugin +import os + +class VarsModule(BaseVarsPlugin): + def __init__(self, *args): + super(VarsModule, self).__init__(*args) + os.environ['EVIL_VARS_PLUGIN'] = 'YIPEEE' + + def get_vars(self, loader, path, entities, cache=True): + super(VarsModule, self).get_vars(loader, path, entities) + return {}