From 6eed3aa1fa20b072cf148a168fbe6a825b57d6eb Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sun, 1 Apr 2018 15:54:17 +0100
Subject: [PATCH] issue #177: fetch and cache HOME value during connection
 setup.

This ensures only 1 roundtrip is required for every invocation of
_remote_expand_user().
---
 ansible_mitogen/connection.py | 20 +++++++++++++++-----
 ansible_mitogen/mixins.py     | 15 ++++++++++++---
 ansible_mitogen/services.py   | 13 +++++++++++--
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/ansible_mitogen/connection.py b/ansible_mitogen/connection.py
index 1c404410..a81a4bd6 100644
--- a/ansible_mitogen/connection.py
+++ b/ansible_mitogen/connection.py
@@ -82,6 +82,9 @@ class Connection(ansible.plugins.connection.ConnectionBase):
     #: Set to 'mitogen_ssh_discriminator' by on_action_run()
     mitogen_ssh_discriminator = None
 
+    #: Set after connection to the target context's home directory.
+    _homedir = None
+
     def __init__(self, play_context, new_stdin, original_transport, **kwargs):
         assert ansible_mitogen.process.MuxProcess.unix_listener_path, (
             'The "mitogen" connection plug-in may only be instantiated '
@@ -125,6 +128,11 @@ class Connection(ansible.plugins.connection.ConnectionBase):
             'sudo'
         )
 
+    @property
+    def homedir(self):
+        self._connect()
+        return self._homedir
+
     @property
     def connected(self):
         return self.broker is not None
@@ -232,18 +240,20 @@ class Connection(ansible.plugins.connection.ConnectionBase):
 
         if self.original_transport == 'local':
             if self._play_context.become:
-                self.context = self._connect_sudo(python_path=sys.executable)
+                self.context, self._homedir = self._connect_sudo(
+                    python_path=sys.executable
+                )
             else:
-                self.context = self._connect_local()
+                self.context, self._homedir = self._connect_local()
             return
 
         if self.original_transport == 'docker':
-            self.host = self._connect_docker()
+            self.host, self._homedir = self._connect_docker()
         elif self.original_transport == 'ssh':
-            self.host = self._connect_ssh()
+            self.host, self._homedir = self._connect_ssh()
 
         if self._play_context.become:
-            self.context = self._connect_sudo(via=self.host)
+            self.context, self._homedir = self._connect_sudo(via=self.host)
         else:
             self.context = self.host
 
diff --git a/ansible_mitogen/mixins.py b/ansible_mitogen/mixins.py
index ad020f67..70f63ec6 100644
--- a/ansible_mitogen/mixins.py
+++ b/ansible_mitogen/mixins.py
@@ -272,10 +272,19 @@ class ActionModuleMixin(ansible.plugins.action.ActionBase):
         Replace the base implementation's attempt to emulate
         os.path.expanduser() with an actual call to os.path.expanduser().
         """
-        LOG.debug('_remove_expand_user(%r, sudoable=%r)', path, sudoable)
+        LOG.debug('_remote_expand_user(%r, sudoable=%r)', path, sudoable)
+        if not path.startswith('~'):
+            # /home/foo -> /home/foo
+            return path
+        if path == '~':
+            # ~ -> /home/dmw
+            return self._connection.homedir
+        if path.startswith('~/'):
+            # ~/.ansible -> /home/dmw/.ansible
+            return os.path.join(self._connection.homedir, path[2:])
         if path.startswith('~'):
-            path = self.call(os.path.expanduser, path)
-        return path
+            # ~root/.ansible -> /root/.ansible
+            return self.call(os.path.expanduser, path)
 
     def _execute_module(self, module_name=None, module_args=None, tmp=None,
                         task_vars=None, persist_files=False,
diff --git a/ansible_mitogen/services.py b/ansible_mitogen/services.py
index 98a231d4..0266d533 100644
--- a/ansible_mitogen/services.py
+++ b/ansible_mitogen/services.py
@@ -28,6 +28,7 @@
 
 from __future__ import absolute_import
 import logging
+import os.path
 import zlib
 
 import mitogen
@@ -62,7 +63,13 @@ class ContextService(mitogen.service.DeduplicatingService):
           existing connection, but popped from the list of arguments passed to
           the connection method.
 
-    :returns mitogen.master.Context:
+    :returns tuple:
+        Tuple of `(context, home_dir)`, where:
+            * `context` is the mitogen.master.Context referring to the target
+              context.
+            * `home_dir` is a cached copy of the remote directory.
+
+    mitogen.master.Context:
         Corresponding Context instance.
     """
     handle = 500
@@ -74,7 +81,9 @@ class ContextService(mitogen.service.DeduplicatingService):
     def get_response(self, args):
         args.pop('discriminator', None)
         method = getattr(self.router, args.pop('method'))
-        return method(**args)
+        context = method(**args)
+        home_dir = context.call(os.path.expanduser, '~')
+        return context, home_dir
 
 
 class FileService(mitogen.service.Service):