Merge pull request #298 from dw/dmw

Dmw

ansible_mitogen/connection.py

@@ -347,6 +347,9 @@ class Connection(ansible.plugins.connection.ConnectionBase):
     #: Set to 'hostvars' by on_action_run()
     host_vars = None
 
+    #: Set to '_loader.get_basedir()' by on_action_run().
+    loader_basedir = None
+
     #: Set after connection to the target context's home directory.
     home_dir = None
 
@@ -366,7 +369,7 @@ class Connection(ansible.plugins.connection.ConnectionBase):
         # https://github.com/dw/mitogen/issues/140
         self.close()
 
-    def on_action_run(self, task_vars):
+    def on_action_run(self, task_vars, loader_basedir):
         """
         Invoked by ActionModuleMixin to indicate a new task is about to start
         executing. We use the opportunity to grab relevant bits from the
@@ -384,6 +387,7 @@ class Connection(ansible.plugins.connection.ConnectionBase):
         self.mitogen_ssh_debug_level = task_vars.get('mitogen_ssh_debug_level')
         self.inventory_hostname = task_vars['inventory_hostname']
         self.host_vars = task_vars['hostvars']
+        self.loader_basedir = loader_basedir
         self.close(new_task=True)
 
     @property
@@ -543,6 +547,20 @@ class Connection(ansible.plugins.connection.ConnectionBase):
         """
         return self.call(ansible_mitogen.target.create_fork_child)
 
+    def get_default_cwd(self):
+        """
+        Overridden by connections/mitogen_local.py to emulate behaviour of CWD
+        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
+
     def exec_command(self, cmd, in_data='', sudoable=True, mitogen_chdir=None):
         """
         Implement exec_command() by calling the corresponding
@@ -560,7 +578,7 @@ class Connection(ansible.plugins.connection.ConnectionBase):
             ansible_mitogen.target.exec_command,
             cmd=mitogen.utils.cast(cmd),
             in_data=mitogen.utils.cast(in_data),
-            chdir=mitogen_chdir,
+            chdir=mitogen_chdir or self.get_default_cwd(),
             emulate_tty=emulate_tty,
         )
 

ansible_mitogen/mixins.py

@@ -108,7 +108,10 @@ class ActionModuleMixin(ansible.plugins.action.ActionBase):
         Override run() to notify Connection of task-specific data, so it has a
         chance to know e.g. the Python interpreter in use.
         """
-        self._connection.on_action_run(task_vars)
+        self._connection.on_action_run(
+            task_vars=task_vars,
+            loader_basedir=self._loader.get_basedir(),
+        )
         return super(ActionModuleMixin, self).run(tmp, task_vars)
 
     def call(self, func, *args, **kwargs):

ansible_mitogen/planner.py

@@ -175,6 +175,8 @@ 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

ansible_mitogen/plugins/connection/mitogen_local.py

@@ -37,7 +37,49 @@ except ImportError:
     del base_dir
 
 import ansible_mitogen.connection
+import ansible_mitogen.process
+
+
+if sys.version_info > (3,):
+    viewkeys = dict.keys
+elif 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):
     transport = 'local'
+
+    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,
+        )

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:

ansible_mitogen/runner.py

@@ -126,8 +126,10 @@ class Runner(object):
         implicit bytes/str conversion behaviour of a 2.x controller running
         against a 3.x target.
     :param dict env:
-        Additional environment variables to set during the run.
-
+        Additional environment variables to set during the run. Keys with
+        :data:`None` are unset if present.
+    :param str cwd:
+        If not :data:`None`, change to this directory before executing.
     :param mitogen.core.ExternalContext econtext:
         When `detach` is :data:`True`, a reference to the ExternalContext the
         runner is executing in.
@@ -135,14 +137,16 @@ 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, 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
 
     def setup(self):
         """
@@ -150,7 +154,16 @@ class Runner(object):
         from the parent, as :meth:`run` may detach prior to beginning
         execution. The base implementation simply prepares the environment.
         """
-        self._env = TemporaryEnvironment(self.env)
+        if self.cwd:
+            # For situations like sudo to another non-privileged account, the
+            # CWD could be $HOME of the old account, which could have mode go=,
+            # which means it is impossible to restore the old directory, so
+            # don't even bother.
+            os.chdir(self.cwd)
+        env = dict(self.extra_env or {})
+        if self.env:
+            env.update(self.env)
+        self._env = TemporaryEnvironment(env)
 
     def revert(self):
         """
@@ -260,14 +273,24 @@ class ModuleUtilsImporter(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):

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

tests/ansible/integration/all.yml

@@ -8,6 +8,7 @@
 - import_playbook: become/all.yml
 - import_playbook: connection_loader/all.yml
 - import_playbook: context_service/all.yml
+- import_playbook: local/all.yml
 #- import_playbook: module_utils/all.yml
 - import_playbook: playbook_semantics/all.yml
 - import_playbook: remote_tmp/all.yml

tests/ansible/integration/local/all.yml

@@ -0,0 +1,4 @@
+
+- import_playbook: cwd_preserved.yml
+- import_playbook: env_preserved.yml
+

tests/ansible/integration/local/cwd_preserved.yml

@@ -0,0 +1,22 @@
+
+# Current working directory should be that of WorkerProcess -- which is the
+# same directory as the currently executing playbook, i.e. integration/local/
+#
+# https://github.com/ansible/ansible/issues/14489
+
+- name: integration/local/cwd_preserved.yml
+  any_errors_fatal: true
+  hosts: test-targets
+  tasks:
+  - connection: local
+    command: pwd
+    register: pwd
+
+  - connection: local
+    stat:
+      path: "{{pwd.stdout}}/cwd_preserved.yml"
+    register: stat
+
+  - assert:
+      that: stat.stat.exists
+

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"

tests/ansible/integration/ssh/timeouts.yml

@@ -4,7 +4,14 @@
   hosts: test-targets
   tasks:
     - connection: local
-      command: ansible -vvv -i "{{inventory_file}}" test-targets -m custom_python_detect_environment -e ansible_user=mitogen__slow_user -e ansible_password=slow_user_password
+      command: |
+        ansible -vvv
+        -i "{{inventory_file}}"
+        test-targets
+        -m custom_python_detect_environment
+        -e ansible_user=mitogen__slow_user -e ansible_password=slow_user_password
+      args:
+        chdir: ../..
       register: out
       ignore_errors: true
       when: is_mitogen

tests/ansible/integration/ssh/variables.yml

@@ -21,6 +21,8 @@
         ansible -m shell -a whoami -i "{{inventory_file}}" test-targets
         -e ansible_ssh_user=mitogen__has_sudo
         -e ansible_ssh_pass=has_sudo_password
+      args:
+        chdir: ../..
       register: out
       when: is_mitogen
 
@@ -30,6 +32,8 @@
         ansible -m shell -a whoami -i "{{inventory_file}}" test-targets
         -e ansible_ssh_user=mitogen__has_sudo
         -e ansible_ssh_pass=wrong_password
+      args:
+        chdir: ../..
       register: out
       ignore_errors: true
       when: is_mitogen
@@ -46,6 +50,8 @@
         ansible -m shell -a whoami -i "{{inventory_file}}" test-targets
         -e ansible_user=mitogen__has_sudo
         -e ansible_ssh_pass=has_sudo_password
+      args:
+        chdir: ../..
       register: out
       when: is_mitogen
 
@@ -55,6 +61,8 @@
         ansible -m shell -a whoami -i "{{inventory_file}}" test-targets
         -e ansible_user=mitogen__has_sudo
         -e ansible_ssh_pass=wrong_password
+      args:
+        chdir: ../..
       register: out
       ignore_errors: true
       when: is_mitogen
@@ -71,6 +79,8 @@
         ansible -m shell -a whoami -i "{{inventory_file}}" test-targets
         -e ansible_user=mitogen__has_sudo
         -e ansible_password=has_sudo_password
+      args:
+        chdir: ../..
       register: out
       when: is_mitogen
 
@@ -80,6 +90,8 @@
         ansible -m shell -a whoami -i "{{inventory_file}}" test-targets
         -e ansible_user=mitogen__has_sudo
         -e ansible_password=wrong_password
+      args:
+        chdir: ../..
       register: out
       ignore_errors: true
       when: is_mitogen
@@ -96,6 +108,8 @@
         ansible -m shell -a whoami -i "{{inventory_file}}" test-targets
         -e ansible_user=mitogen__has_sudo_pubkey
         -e ansible_ssh_private_key_file=../data/docker/mitogen__has_sudo_pubkey.key
+      args:
+        chdir: ../..
       register: out
       when: is_mitogen
 
@@ -105,6 +119,8 @@
         ansible -m shell -a whoami -i "{{inventory_file}}" test-targets
         -e ansible_user=mitogen__has_sudo
         -e ansible_ssh_private_key_file=/dev/null
+      args:
+        chdir: ../..
       register: out
       ignore_errors: true
       when: is_mitogen

tests/ansible/integration/strategy/_mixed_mitogen_vanilla.yml

@@ -0,0 +1,45 @@
+
+# issue #294: ensure running mixed vanilla/Mitogen succeeds.
+
+- name: integration/strategy/_mixed_mitogen_vanilla.yml (mitogen_linear)
+  hosts: test-targets
+  any_errors_fatal: true
+  strategy: mitogen_linear
+  tasks:
+    - custom_python_detect_environment:
+      register: out
+    - assert:
+        that: out.mitogen_loaded
+
+    - determine_strategy:
+    - assert:
+        that: strategy == 'ansible.plugins.strategy.mitogen_linear.StrategyModule'
+
+- name: integration/strategy/_mixed_mitogen_vanilla.yml (linear)
+  hosts: test-targets
+  any_errors_fatal: true
+  strategy: linear
+  tasks:
+    - custom_python_detect_environment:
+      register: out
+    - assert:
+        that: not out.mitogen_loaded
+
+    - determine_strategy:
+    - assert:
+        that: strategy == 'ansible.plugins.strategy.linear.StrategyModule'
+
+
+- name: integration/strategy/_mixed_mitogen_vanilla.yml (mitogen_linear)
+  hosts: test-targets
+  any_errors_fatal: true
+  strategy: mitogen_linear
+  tasks:
+    - custom_python_detect_environment:
+      register: out
+    - assert:
+        that: out.mitogen_loaded
+
+    - determine_strategy:
+    - assert:
+        that: strategy == 'ansible.plugins.strategy.mitogen_linear.StrategyModule'

tests/ansible/integration/strategy/_mixed_vanilla_mitogen.yml

@@ -1,7 +1,7 @@
 
 # issue #294: ensure running mixed vanilla/Mitogen succeeds.
 
-- name: integration/strategy/mixed_vanilla_ansible.yml (linear)
+- name: integration/strategy/_mixed_vanilla_mitogen.yml (linear)
   hosts: test-targets
   any_errors_fatal: true
   strategy: linear
@@ -15,7 +15,7 @@
     - assert:
         that: strategy == 'ansible.plugins.strategy.linear.StrategyModule'
 
-- name: integration/strategy/mixed_vanilla_ansible.yml (mitogen_linear)
+- name: integration/strategy/_mixed_vanilla_mitogen.yml (mitogen_linear)
   hosts: test-targets
   any_errors_fatal: true
   strategy: mitogen_linear
@@ -30,7 +30,7 @@
         that: strategy == 'ansible.plugins.strategy.mitogen_linear.StrategyModule'
 
 
-- name: integration/strategy/mixed_vanilla_ansible.yml (linear)
+- name: integration/strategy/_mixed_vanilla_mitogen.yml (linear)
   hosts: test-targets
   any_errors_fatal: true
   strategy: linear

tests/ansible/integration/strategy/all.yml

@@ -1 +1 @@
-- import_playbook: mixed_vanilla_ansible.yml
+- import_playbook: mixed_vanilla_mitogen.yml

tests/ansible/integration/strategy/mixed_vanilla_mitogen.yml

@@ -0,0 +1,22 @@
+
+- name: integration/strategy/mixed_vanilla_mitogen.yml (linear->mitogen->linear)
+  hosts: test-targets
+  any_errors_fatal: true
+  tasks:
+    - connection: local
+      command: |
+        ansible-playbook
+        -i "{{inventory_file}}"
+        integration/strategy/_mixed_mitogen_vanilla.yml
+      args:
+        chdir: ../..
+      register: out
+
+    - connection: local
+      command: |
+        ansible-playbook
+        -i "{{inventory_file}}"
+        integration/strategy/_mixed_vanilla_mitogen.yml
+      args:
+        chdir: ../..
+      register: out

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 {}